jaclang 0.5.7__py3-none-any.whl → 0.5.8__py3-none-any.whl

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -210,9 +210,18 @@ class PyastGenPass(Pass):
                 i.col_offset = jac_node.loc.col_start
                 i.end_lineno = jac_node.loc.last_line
                 i.end_col_offset = jac_node.loc.col_end
-                i.jac_link = jac_node  # type: ignore
+                i.jac_link: list[ast3.AST] = [jac_node]  # type: ignore
         return py_node
 
+    def link_jac_py_nodes(
+        self, jac_node: ast.AstNode, py_nodes: list[ast3.AST]
+    ) -> None:
+        """Link jac name ast to py ast nodes."""
+        jac_node.gen.py_ast = py_nodes
+        for i in py_nodes:
+            if isinstance(i.jac_link, list):  # type: ignore
+                i.jac_link.append(jac_node)  # type: ignore
+
     def pyinline_sync(
         self,
         py_nodes: list[ast3.AST],
@@ -225,7 +234,7 @@ class PyastGenPass(Pass):
                         i.lineno += self.cur_node.loc.first_line
                     if hasattr(i, "end_lineno") and i.end_lineno is not None:
                         i.end_lineno += self.cur_node.loc.first_line
-                    i.jac_link = self.cur_node  # type: ignore
+                    i.jac_link: ast3.AST = [self.cur_node]  # type: ignore
         return py_nodes
 
     def resolve_stmt_block(
@@ -247,7 +256,7 @@ class PyastGenPass(Pass):
                     [
                         x.gen.py_ast
                         for x in node.items
-                        # if not isinstance(x, ast.AstImplOnlyNode)
+                        if not isinstance(x, ast.AstImplOnlyNode)
                     ]
                 )
                 if node and isinstance(node.gen.py_ast, list)
@@ -305,7 +314,8 @@ class PyastGenPass(Pass):
         body: Sequence[ElementStmt],
         is_imported: bool,
         """
-        pre_body = [*node.impl_mod.body, *node.body] if node.impl_mod else node.body
+        clean_body = [i for i in node.body if not isinstance(i, ast.AstImplOnlyNode)]
+        pre_body = [*node.impl_mod.body, *clean_body] if node.impl_mod else clean_body
         pre_body = [*pre_body, *node.test_mod.body] if node.test_mod else pre_body
         body = (
             [
@@ -567,6 +577,8 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        if node.alias:
+            self.link_jac_py_nodes(jac_node=node.alias, py_nodes=node.gen.py_ast)
 
     def exit_module_item(self, node: ast.ModuleItem) -> None:
         """Sub objects.
@@ -691,6 +703,9 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        self.link_jac_py_nodes(jac_node=node.name, py_nodes=node.gen.py_ast)
+        if isinstance(node.body, ast.ArchDef):
+            self.link_jac_py_nodes(jac_node=node.body, py_nodes=node.gen.py_ast)
 
     def collect_events(
         self, node: ast.Architype
@@ -743,6 +758,12 @@ class PyastGenPass(Pass):
         doc: Optional[String],
         decorators: Optional[SubNodeList[ExprType]],
         """
+        for i in node.target.archs:
+            if i.sym_link:
+                self.link_jac_py_nodes(jac_node=i, py_nodes=i.sym_link.decl.gen.py_ast)
+                self.link_jac_py_nodes(
+                    jac_node=i.name_ref, py_nodes=i.sym_link.decl.gen.py_ast
+                )
 
     def exit_enum(self, node: ast.Enum) -> None:
         """Sub objects.
@@ -783,6 +804,8 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        if isinstance(node.body, ast.EnumDef):
+            self.link_jac_py_nodes(jac_node=node.body, py_nodes=node.gen.py_ast)
 
     def exit_enum_def(self, node: ast.EnumDef) -> None:
         """Sub objects.
@@ -792,6 +815,12 @@ class PyastGenPass(Pass):
         doc: Optional[String],
         decorators: Optional[SubNodeList[ExprType]],
         """
+        for i in node.target.archs:
+            if i.sym_link:
+                self.link_jac_py_nodes(jac_node=i, py_nodes=i.sym_link.decl.gen.py_ast)
+                self.link_jac_py_nodes(
+                    jac_node=i.name_ref, py_nodes=i.sym_link.decl.gen.py_ast
+                )
 
     def exit_ability(self, node: ast.Ability) -> None:
         """Sub objects.
@@ -815,12 +844,12 @@ class PyastGenPass(Pass):
                 [
                     self.sync(
                         ast3.Expr(value=node.doc.gen.py_ast[0]), jac_node=node.doc
-                    ),  # type: ignore
-                    self.sync(ast3.Pass(), node.body),  # type: ignore
+                    ),
+                    self.sync(ast3.Pass(), node.body),
                 ]
                 if node.doc and node.is_abstract
                 else (
-                    [self.sync(ast3.Pass(), node.body)]  # type: ignore
+                    [self.sync(ast3.Pass(), node.body)]
                     if node.is_abstract
                     else self.resolve_stmt_block(
                         (
@@ -875,6 +904,7 @@ class PyastGenPass(Pass):
             )
         if not body and not isinstance(node.body, ast.FuncCall):
             self.error("Ability has no body. Perhaps an impl must be imported.", node)
+            body = [self.sync(ast3.Pass(), node)]
 
         node.gen.py_ast = [
             self.sync(
@@ -892,6 +922,9 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        self.link_jac_py_nodes(jac_node=node.name_ref, py_nodes=node.gen.py_ast)
+        if isinstance(node.body, ast.AbilityDef):
+            self.link_jac_py_nodes(jac_node=node.body, py_nodes=node.gen.py_ast)
 
     def gen_llm_body(self, node: ast.Ability) -> list[ast3.AST]:
         """Generate llm body."""
@@ -1077,6 +1110,15 @@ class PyastGenPass(Pass):
         doc: Optional[String],
         decorators: Optional[SubNodeList[ExprType]],
         """
+        for i in node.target.archs:
+            if i.sym_link:
+                self.link_jac_py_nodes(jac_node=i, py_nodes=i.sym_link.decl.gen.py_ast)
+                self.link_jac_py_nodes(
+                    jac_node=i.name_ref, py_nodes=i.sym_link.decl.gen.py_ast
+                )
+        if isinstance(node.parent, ast.Ability) and node.parent.signature:
+            # TODO: Here we need to do a link for each subnode to the original parent signature
+            pass
 
     def exit_func_signature(self, node: ast.FuncSignature) -> None:
         """Sub objects.
@@ -1233,6 +1275,7 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        self.link_jac_py_nodes(jac_node=node.name, py_nodes=node.gen.py_ast)
 
     def exit_arch_has(self, node: ast.ArchHas) -> None:
         """Sub objects.
@@ -1488,6 +1531,8 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        if node.name:
+            self.link_jac_py_nodes(jac_node=node.name, py_nodes=node.gen.py_ast)
 
     def exit_finally_stmt(self, node: ast.FinallyStmt) -> None:
         """Sub objects.
@@ -1595,6 +1640,9 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        self.link_jac_py_nodes(jac_node=node.expr, py_nodes=node.gen.py_ast)
+        if node.alias:
+            self.link_jac_py_nodes(jac_node=node.alias, py_nodes=node.gen.py_ast)
 
     def exit_raise_stmt(self, node: ast.RaiseStmt) -> None:
         """Sub objects.
@@ -1852,6 +1900,7 @@ class PyastGenPass(Pass):
                     jac_node=x,
                 )
             )
+            self.link_jac_py_nodes(jac_node=x, py_nodes=[py_nodes[-1]])
         node.gen.py_ast = [*py_nodes]
 
     def exit_non_local_stmt(self, node: ast.NonLocalStmt) -> None:
@@ -1867,6 +1916,7 @@ class PyastGenPass(Pass):
                     jac_node=x,
                 )
             )
+            self.link_jac_py_nodes(jac_node=x, py_nodes=[py_nodes[-1]])
         node.gen.py_ast = [*py_nodes]
 
     def exit_assignment(self, node: ast.Assignment) -> None:
@@ -1877,6 +1927,21 @@ class PyastGenPass(Pass):
         type_tag: Optional[SubTag[ExprType]],
         mutable: bool =True,
         """
+        value = (
+            node.value.gen.py_ast[0]
+            if node.value
+            else (
+                self.sync(
+                    ast3.Call(
+                        func=self.sync(ast3.Name(id="__jac_auto__", ctx=ast3.Load())),
+                        args=[],
+                        keywords=[],
+                    )
+                )
+                if node.is_enum_stmt
+                else self.ice()
+            )
+        )
         if node.type_tag:
             node.gen.py_ast = [
                 self.sync(
@@ -1888,25 +1953,19 @@ class PyastGenPass(Pass):
                     )
                 )
             ]
-        elif not node.value:
-            self.ice()
         elif node.aug_op:
             node.gen.py_ast = [
                 self.sync(
                     ast3.AugAssign(
                         target=node.target.items[0].gen.py_ast[0],
                         op=node.aug_op.gen.py_ast[0],
-                        value=node.value.gen.py_ast[0],
+                        value=value,
                     )
                 )
             ]
         else:
             node.gen.py_ast = [
-                self.sync(
-                    ast3.Assign(
-                        targets=node.target.gen.py_ast, value=node.value.gen.py_ast[0]
-                    )
-                )
+                self.sync(ast3.Assign(targets=node.target.gen.py_ast, value=value))
             ]
 
     def exit_binary_expr(self, node: ast.BinaryExpr) -> None:
@@ -2037,7 +2096,9 @@ class PyastGenPass(Pass):
                 params=(
                     node.left.values
                     if isinstance(node.left, ast.TupleVal)
-                    else ast.SubNodeList(items=[node.left], kid=[node.left])
+                    else ast.SubNodeList(
+                        items=[node.left], delim=Tok.COMMA, kid=[node.left]
+                    )
                 ),
                 kid=node.kid,
             )
@@ -2071,7 +2132,9 @@ class PyastGenPass(Pass):
                 params=(
                     node.right.values
                     if isinstance(node.right, ast.TupleVal)
-                    else ast.SubNodeList(items=[node.right], kid=[node.right])
+                    else ast.SubNodeList(
+                        items=[node.right], delim=Tok.COMMA, kid=[node.right]
+                    )
                 ),
                 kid=node.kid,
             )
@@ -2293,13 +2356,6 @@ class PyastGenPass(Pass):
             else [self.sync(ast3.Constant(value=""))]
         )
 
-    def exit_expr_list(self, node: ast.ExprList) -> None:
-        """Sub objects.
-
-        values: Optional[SubNodeList[ExprType]],
-        """
-        node.gen.py_ast = node.values.gen.py_ast if node.values else []
-
     def exit_list_val(self, node: ast.ListVal) -> None:
         """Sub objects.
 
@@ -2380,6 +2436,8 @@ class PyastGenPass(Pass):
                 )
             )
         ]
+        if node.key:
+            self.link_jac_py_nodes(jac_node=node.key, py_nodes=node.gen.py_ast)
 
     def exit_inner_compr(self, node: ast.InnerCompr) -> None:
         """Sub objects.
@@ -2394,7 +2452,11 @@ class PyastGenPass(Pass):
                 ast3.comprehension(
                     target=node.target.gen.py_ast[0],
                     iter=node.collection.gen.py_ast[0],
-                    ifs=node.conditional.gen.py_ast if node.conditional else [],
+                    ifs=(
+                        [x.gen.py_ast[0] for x in node.conditional]
+                        if node.conditional
+                        else []
+                    ),
                     is_async=0,
                 )
             )
@@ -2404,7 +2466,7 @@ class PyastGenPass(Pass):
         """Sub objects.
 
         out_expr: ExprType,
-        compr: InnerCompr,
+        compr: list[InnerCompr]
         """
         node.gen.py_ast = [
             self.sync(
@@ -2419,7 +2481,7 @@ class PyastGenPass(Pass):
         """Sub objects.
 
         out_expr: ExprType,
-        compr: InnerCompr,
+        compr: list[InnerCompr]
         """
         node.gen.py_ast = [
             self.sync(
@@ -2434,7 +2496,7 @@ class PyastGenPass(Pass):
         """Sub objects.
 
         out_expr: ExprType,
-        compr: InnerCompr,
+        compr: list[InnerCompr]
         """
         node.gen.py_ast = [
             self.sync(
@@ -2468,27 +2530,25 @@ class PyastGenPass(Pass):
 
         target: Expr,
         right: AtomExpr | Expr,
-        is_attr: Optional[Token],
+        is_attr: bool,
         is_null_ok: bool,
         """
         if node.is_attr:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.Attribute(
-                        value=node.target.gen.py_ast[0],
-                        attr=(
-                            node.right.sym_name
-                            if isinstance(node.right, ast.AstSymbolNode)
-                            else ""
-                        ),
-                        ctx=(
-                            node.right.py_ctx_func()
-                            if isinstance(node.right, ast.AstSymbolNode)
-                            else ast3.Load()
-                        ),
+            if isinstance(node.right, ast.AstSymbolNode):
+                node.gen.py_ast = [
+                    self.sync(
+                        ast3.Attribute(
+                            value=node.target.gen.py_ast[0],
+                            attr=(node.right.sym_name),
+                            ctx=(node.right.py_ctx_func()),
+                        )
                     )
+                ]
+                self.link_jac_py_nodes(
+                    jac_node=node.right.sym_name_node, py_nodes=node.gen.py_ast
                 )
-            ]
+            else:
+                self.error("Invalid attribute access")
         elif isinstance(node.right, ast.FilterCompr):
             node.gen.py_ast = [
                 self.sync(
@@ -2568,19 +2628,11 @@ class PyastGenPass(Pass):
         target: Expr,
         params: Optional[SubNodeList[Expr | KWPair]],
         """
-        node.gen.py_ast = [self.sync(self.gen_func_call(node.target, node.params))]
-
-    def gen_func_call(
-        self,
-        target: ast.Expr,
-        params: Optional[ast.SubNodeList[ast.Expr | ast.KWPair]],
-    ) -> ast3.Call:
-        """Generate a function call."""
-        func = target.gen.py_ast[0]
+        func = node.target.gen.py_ast[0]
         args = []
         keywords = []
-        if params and len(params.items) > 0:
-            for x in params.items:
+        if node.params and len(node.params.items) > 0:
+            for x in node.params.items:
                 if isinstance(x, ast.UnaryExpr) and x.op.name == Tok.STAR_POW:
                     keywords.append(
                         self.sync(ast3.keyword(value=x.operand.gen.py_ast[0]), x)
@@ -2593,7 +2645,9 @@ class PyastGenPass(Pass):
                     keywords.append(x.gen.py_ast[0])
                 else:
                     self.ice("Invalid Parameter")
-        return ast3.Call(func=func, args=args, keywords=keywords)
+        node.gen.py_ast = [
+            self.sync(ast3.Call(func=func, args=args, keywords=keywords))
+        ]
 
     def exit_index_slice(self, node: ast.IndexSlice) -> None:
         """Sub objects.
@@ -2976,14 +3030,14 @@ class PyastGenPass(Pass):
 
         pattern: MatchPattern,
         guard: Optional[ExprType],
-        body: SubNodeList[CodeBlockStmt],
+        body: list[CodeBlockStmt],
         """
         node.gen.py_ast = [
             self.sync(
                 ast3.match_case(
                     pattern=node.pattern.gen.py_ast[0],
                     guard=node.guard.gen.py_ast[0] if node.guard else None,
-                    body=self.resolve_stmt_block(node.body),
+                    body=[x.gen.py_ast[0] for x in node.body],
                 )
             )
         ]
@@ -3204,24 +3258,6 @@ class PyastGenPass(Pass):
         node.gen.py_ast = [
             self.sync(ast3.Name(id=node.sym_name, ctx=node.py_ctx_func()))
         ]
-        if node.is_enum_singleton and isinstance(node.gen.py_ast[0], ast3.Name):
-            node.gen.py_ast[0].ctx = ast3.Store()
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.Assign(
-                        targets=node.gen.py_ast,
-                        value=self.sync(
-                            ast3.Call(
-                                func=self.sync(
-                                    ast3.Name(id="__jac_auto__", ctx=ast3.Load())
-                                ),
-                                args=[],
-                                keywords=[],
-                            )
-                        ),
-                    )
-                )
-            ]
 
     def exit_float(self, node: ast.Float) -> None:
         """Sub objects.