Make is_dataclass robust (lcompilers#1385)

seifashraf1 · Dec 27, 2022 · 51bbc8e · 51bbc8e
1 parent e6cff59
commit 51bbc8e
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 70 deletions.
diff --git a/integration_tests/structs_14.py b/integration_tests/structs_14.py
@@ -1,4 +1,4 @@
-from ltypes import i8, dataclass, i32, f32, c32, f64, i16, i64, c64, ccallable
+from ltypes import i8, dataclass, i32, f32, c32, f64, i16, i64, c64, ccallable, packed
 from numpy import empty, int8, int16, int32, int64, float32, complex64, complex128, float64
 from copy import deepcopy
 
@@ -14,6 +14,7 @@ class buffer_struct:
     buffer7: c64[32]
 
 @ccallable
+@packed
 @dataclass
 class buffer_struct_clink:
     buffer: i8[32]

diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp
@@ -2133,80 +2133,43 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
 
     bool is_dataclass(AST::expr_t** decorators, size_t n,
         ASR::expr_t*& aligned_expr, bool& is_packed) {
-        if( n == 1 ) {
-            AST::expr_t* decorator = decorators[0];
-            if( !AST::is_a<AST::Name_t>(*decorator) ) {
-                return false;
-            }
-
-            AST::Name_t* dec_name = AST::down_cast<AST::Name_t>(decorator);
-            return std::string(dec_name->m_id) == "dataclass";
-        } else if( n == 2 ) {
-
-            if( AST::is_a<AST::Name_t>(*decorators[0]) &&
-                AST::is_a<AST::Name_t>(*decorators[1]) ) {
-                AST::Name_t* dc1_name = AST::down_cast<AST::Name_t>(decorators[0]);
-                AST::Name_t* dc2_name = AST::down_cast<AST::Name_t>(decorators[1]);
-                std::string dc1_str = dc1_name->m_id;
-                std::string dc2_str = dc2_name->m_id;
-                if( dc2_str == "packed" || dc1_str == "packed" ) {
+        bool is_dataclass_ = false;
+        for( size_t i = 0; i < n; i++ ) {
+            if( AST::is_a<AST::Name_t>(*decorators[i]) ) {
+                AST::Name_t* dc_name = AST::down_cast<AST::Name_t>(decorators[i]);
+                is_dataclass_ = std::string(dc_name->m_id) == "dataclass";
+            } else if( AST::is_a<AST::Call_t>(*decorators[i]) ) {
+                AST::Call_t* dc_call = AST::down_cast<AST::Call_t>(decorators[i]);
+                AST::Name_t* dc_name = AST::down_cast<AST::Name_t>(dc_call->m_func);
+                if( std::string(dc_name->m_id) == "packed" ) {
                     is_packed = true;
+                    if( dc_call->n_keywords == 1 && dc_call->n_args == 0 ) {
+                        AST::keyword_t kwarg = dc_call->m_keywords[0];
+                        std::string kwarg_name = kwarg.m_arg;
+                        if( kwarg_name == "aligned" ) {
+                            this->visit_expr(*kwarg.m_value);
+                            aligned_expr = ASRUtils::EXPR(tmp);
+                            ASR::expr_t* aligned_expr_value = ASRUtils::expr_value(aligned_expr);
+                            std::string msg = "Alignment should always evaluate to a constant expressions.";
+                            if( !aligned_expr_value ) {
+                                throw SemanticError(msg, aligned_expr->base.loc);
+                            }
+                            int64_t alignment_int;
+                            if( !ASRUtils::extract_value(aligned_expr_value, alignment_int) ) {
+                                throw SemanticError(msg, aligned_expr->base.loc);
+                            }
+                            if( alignment_int == 0 || ((alignment_int & (alignment_int - 1)) != 0) ) {
+                                throw SemanticError("Alignment " + std::to_string(alignment_int) +
+                                                    " is not a positive power of 2.",
+                                                    aligned_expr->base.loc);
+                            }
+                        }
+                    }
                 }
-                aligned_expr = nullptr;
-                return ((dc1_str == "dataclass" && (dc2_str == "packed" || dc2_str == "ccallable")) ||
-                        ((dc1_str == "packed" || dc1_str == "ccallable") && dc2_str == "dataclass"));
             }
-
-            int32_t dc_idx = -1, pck_idx = -1;
-            if( AST::is_a<AST::Name_t>(*decorators[0]) &&
-                AST::is_a<AST::Call_t>(*decorators[1]) ) {
-                dc_idx = 0, pck_idx = 1;
-            } else if( AST::is_a<AST::Call_t>(*decorators[0]) &&
-                    AST::is_a<AST::Name_t>(*decorators[1]) ) {
-                dc_idx = 1, pck_idx = 0;
-            }
-            if( dc_idx == -1 || pck_idx == -1 ) {
-                return false;
-            }
-            AST::Name_t* dc_name = AST::down_cast<AST::Name_t>(decorators[dc_idx]);
-            if( std::string(dc_name->m_id) != "dataclass" ) {
-                 return false;
-            }
-            AST::Call_t* pck_call = AST::down_cast<AST::Call_t>(decorators[pck_idx]);
-            AST::Name_t* pck_name = AST::down_cast<AST::Name_t>(pck_call->m_func);
-            if( std::string(pck_name->m_id) != "packed" ) {
-                return false;
-            }
-            if( !(pck_call->n_keywords == 1 && pck_call->n_args == 0) ) {
-                return false;
-            }
-
-            AST::keyword_t kwarg = pck_call->m_keywords[0];
-            std::string kwarg_name = kwarg.m_arg;
-            if( kwarg_name != "aligned" ) {
-                return false;
-            }
-            this->visit_expr(*kwarg.m_value);
-            aligned_expr = ASRUtils::EXPR(tmp);
-            ASR::expr_t* aligned_expr_value = ASRUtils::expr_value(aligned_expr);
-            std::string msg = "Alignment should always evaluate to a constant expressions.";
-            if( !aligned_expr_value ) {
-                throw SemanticError(msg, aligned_expr->base.loc);
-            }
-            int64_t alignment_int;
-            if( !ASRUtils::extract_value(aligned_expr_value, alignment_int) ) {
-                throw SemanticError(msg, aligned_expr->base.loc);
-            }
-            if( alignment_int == 0 || ((alignment_int & (alignment_int - 1)) != 0) ) {
-                throw SemanticError("Alignment " + std::to_string(alignment_int) +
-                                    " is not a positive power of 2.",
-                                    aligned_expr->base.loc);
-            }
-            is_packed = true;
-            return true;
         }
 
-        return false;
+        return is_dataclass_;
     }
 
     bool is_enum(AST::expr_t** bases, size_t n) {