jaclang 0.8.7__py3-none-any.whl → 0.8.9__py3-none-any.whl

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -64,7 +64,6 @@ TOKEN_AST_MAP: dict[Tok, type[ast3.AST]] = {
     Tok.MINUS: ast3.Sub,
     Tok.SUB_EQ: ast3.Sub,
     Tok.BW_NOT: ast3.Invert,
-    Tok.BW_NOT_EQ: ast3.Invert,
     Tok.NOT: ast3.Not,
     Tok.EQ: ast3.NotEq,
     Tok.EE: ast3.Eq,
@@ -92,11 +91,34 @@ UNARY_OP_MAP: dict[Tok, type[ast3.unaryop]] = {
 class PyastGenPass(UniPass):
     """Jac blue transpilation to python pass."""
 
+    # Builtins that should be imported from jaclang.runtimelib.builtin
+    KNOWN_BUILTINS = {
+        "abstractmethod",
+        "ClassVar",
+        "override",
+        "printgraph",
+        "jid",
+        "jobj",
+        "grant",
+        "revoke",
+        "allroots",
+        "save",
+        "commit",
+        "NoPerm",
+        "ReadPerm",
+        "ConnectPerm",
+        "WritePerm",
+    }
+
     def before_pass(self) -> None:
+        self.child_passes: list[PyastGenPass] = []
         for i in self.ir_in.impl_mod + self.ir_in.test_mod:
-            PyastGenPass(ir_in=i, prog=self.prog)
+            child_pass = PyastGenPass(ir_in=i, prog=self.prog)
+            self.child_passes.append(child_pass)
         self.debuginfo: dict[str, list[str]] = {"jac_mods": []}
         self.already_added: list[str] = []
+        self.jaclib_imports: set[str] = set()  # Track individual jaclib imports
+        self.builtin_imports: set[str] = set()  # Track individual builtin imports
         self.preamble: list[ast3.AST] = [
             self.sync(
                 ast3.ImportFrom(
@@ -106,40 +128,6 @@ class PyastGenPass(UniPass):
                 ),
                 jac_node=self.ir_out,
             ),
-            (
-                self.sync(
-                    ast3.ImportFrom(
-                        module="jaclang.runtimelib.builtin",
-                        names=[
-                            self.sync(
-                                ast3.alias(
-                                    name="*",
-                                    asname=None,
-                                )
-                            )
-                        ],
-                        level=0,
-                    ),
-                    jac_node=self.ir_out,
-                )
-            ),
-            (
-                self.sync(
-                    ast3.ImportFrom(
-                        module="jaclang",
-                        names=[
-                            self.sync(
-                                ast3.alias(
-                                    name="JacMachineInterface",
-                                    asname=settings.pyout_jaclib_alias,
-                                )
-                            ),
-                        ],
-                        level=0,
-                    ),
-                    jac_node=self.ir_out,
-                )
-            ),
         ]
 
     def enter_node(self, node: uni.UniNode) -> None:
@@ -161,15 +149,31 @@ class PyastGenPass(UniPass):
 
     def jaclib_obj(self, obj_name: str) -> ast3.Name | ast3.Attribute:
         """Return the object from jaclib as ast node based on the import config."""
-        return self.sync(
-            ast3.Attribute(
-                value=self.sync(
-                    ast3.Name(id=settings.pyout_jaclib_alias, ctx=ast3.Load())
-                ),
-                attr=obj_name,
-                ctx=ast3.Load(),
+        if settings.library_mode:
+            self.jaclib_imports.add(obj_name)
+            return self.sync(ast3.Name(id=obj_name, ctx=ast3.Load()))
+        else:
+            self.needs_jaclib()
+            return self.sync(
+                ast3.Attribute(
+                    value=self.sync(
+                        ast3.Name(id=settings.pyout_jaclib_alias, ctx=ast3.Load())
+                    ),
+                    attr=obj_name,
+                    ctx=ast3.Load(),
+                )
             )
-        )
+
+    def builtin_name(self, name: str) -> ast3.Name:
+        """Return a builtin name and track it for importing.
+
+        Note: Some names like 'Enum' are provided by other imports (e.g., needs_enum)
+        and should not be added to builtin_imports.
+        """
+        # Enum is imported via needs_enum, not from builtins
+        if name not in ["Enum"]:
+            self.builtin_imports.add(name)
+        return self.sync(ast3.Name(id=name, ctx=ast3.Load()))
 
     def _add_preamble_once(self, key: str, node: ast3.AST) -> None:
         """Append an import statement to the preamble once."""
@@ -187,15 +191,6 @@ class PyastGenPass(UniPass):
             ),
         )
 
-    def needs_mtllm(self) -> None:
-        """Ensure byLLM is imported only once."""
-        self._add_preamble_once(
-            self.needs_mtllm.__name__,
-            ast3.Import(
-                names=[self.sync(ast3.alias(name="byllm"), jac_node=self.ir_out)]
-            ),
-        )
-
     def needs_enum(self) -> None:
         """Ensure Enum utilities are imported only once."""
         self._add_preamble_once(
@@ -221,6 +216,24 @@ class PyastGenPass(UniPass):
             ),
         )
 
+    def needs_jaclib(self) -> None:
+        """Ensure JacMachineInterface is imported only once."""
+        self._add_preamble_once(
+            self.needs_jaclib.__name__,
+            ast3.ImportFrom(
+                module="jaclang",
+                names=[
+                    self.sync(
+                        ast3.alias(
+                            name="JacMachineInterface",
+                            asname=settings.pyout_jaclib_alias,
+                        )
+                    )
+                ],
+                level=0,
+            ),
+        )
+
     def _get_sem_decorator(self, node: uni.UniNode) -> ast3.Call | None:
         """Create a semstring decorator for the given semantic strings.
 
@@ -283,7 +296,10 @@ class PyastGenPass(UniPass):
                 else {}
             )
 
-        if not semstr and not inner_semstr:
+        # Only add sem decorator if there's actual semantic content
+        if not semstr and (
+            not inner_semstr or all(not v for v in inner_semstr.values())
+        ):
             return None
 
         return self.sync(
@@ -417,6 +433,48 @@ class PyastGenPass(UniPass):
         node.gen.py_ast = node.tag.gen.py_ast
 
     def exit_module(self, node: uni.Module) -> None:
+        # Check if any child passes (impl_mod or test_mod) needed jaclib
+        for child_pass in self.child_passes:
+            if "needs_jaclib" in child_pass.already_added:
+                self.needs_jaclib()
+                break
+            # Merge jaclib and builtin imports from child passes
+            if settings.library_mode:
+                self.jaclib_imports.update(child_pass.jaclib_imports)
+            self.builtin_imports.update(child_pass.builtin_imports)
+
+        # Add builtin imports if any were used
+        if self.builtin_imports:
+            self.preamble.append(
+                self.sync(
+                    ast3.ImportFrom(
+                        module="jaclang.runtimelib.builtin",
+                        names=[
+                            self.sync(ast3.alias(name=name, asname=None))
+                            for name in sorted(self.builtin_imports)
+                        ],
+                        level=0,
+                    ),
+                    jac_node=self.ir_out,
+                )
+            )
+
+        # Add library mode imports at the end of preamble
+        if settings.library_mode and self.jaclib_imports:
+            self.preamble.append(
+                self.sync(
+                    ast3.ImportFrom(
+                        module="jaclang.lib",
+                        names=[
+                            self.sync(ast3.alias(name=name, asname=None))
+                            for name in sorted(self.jaclib_imports)
+                        ],
+                        level=0,
+                    ),
+                    jac_node=self.ir_out,
+                )
+            )
+
         clean_body = [i for i in node.body if not isinstance(i, uni.ImplDef)]
         pre_body: list[uni.UniNode] = []
         for pbody in node.impl_mod:
@@ -722,7 +780,7 @@ class PyastGenPass(UniPass):
             decorators.append(sem_decorator)
 
         base_classes = [cast(ast3.expr, i.gen.py_ast[0]) for i in node.base_classes]
-        base_classes.append(self.sync(ast3.Name(id="Enum", ctx=ast3.Load())))
+        base_classes.append(self.builtin_name("Enum"))
         node.gen.py_ast = [
             self.sync(
                 ast3.ClassDef(
@@ -744,23 +802,9 @@ class PyastGenPass(UniPass):
         self, model: ast3.expr, caller: ast3.expr, args: ast3.Dict
     ) -> ast3.Call:
         """Reusable method to codegen call_llm(model, caller, args)."""
-        self.needs_mtllm()
-        mtir_cls_ast = self.sync(
-            ast3.Attribute(
-                value=self.sync(ast3.Name(id="byllm", ctx=ast3.Load())),
-                attr="MTIR",
-                ctx=ast3.Load(),
-            )
-        )
         mtir_ast = self.sync(
             ast3.Call(
-                func=self.sync(
-                    ast3.Attribute(
-                        value=mtir_cls_ast,
-                        attr="factory",
-                        ctx=ast3.Load(),
-                    )
-                ),
+                func=self.jaclib_obj("get_mtir"),
                 args=[],
                 keywords=[
                     self.sync(
@@ -928,7 +972,9 @@ class PyastGenPass(UniPass):
         if isinstance(node.signature, uni.EventSignature):
             decorator_list.append(
                 self.jaclib_obj(
-                    "entry" if node.signature.event.name == Tok.KW_ENTRY else "exit"
+                    "on_entry"
+                    if node.signature.event.name == Tok.KW_ENTRY
+                    else "on_exit"
                 )
             )
 
@@ -943,11 +989,9 @@ class PyastGenPass(UniPass):
                 )
             )
         if node.is_abstract:
-            decorator_list.append(
-                self.sync(ast3.Name(id="abstractmethod", ctx=ast3.Load()))
-            )
+            decorator_list.append(self.builtin_name("abstractmethod"))
         if node.is_override:
-            decorator_list.append(self.sync(ast3.Name(id="override", ctx=ast3.Load())))
+            decorator_list.append(self.builtin_name("override"))
         if node.is_static:
             decorator_list.insert(
                 0, self.sync(ast3.Name(id="staticmethod", ctx=ast3.Load()))
@@ -1000,6 +1044,9 @@ class PyastGenPass(UniPass):
         pass
 
     def exit_func_signature(self, node: uni.FuncSignature) -> None:
+        posonlyargs = [i.gen.py_ast[0] for i in node.posonly_params]
+        vararg = node.varargs.gen.py_ast[0] if node.varargs else None
+        kwarg = node.kwargs.gen.py_ast[0] if node.kwargs else None
         params = (
             [self.sync(ast3.arg(arg="self", annotation=None))]
             if (abl := node.parent)
@@ -1009,29 +1056,33 @@ class PyastGenPass(UniPass):
             and not node.is_in_py_class
             else []
         )
-        vararg = None
-        kwarg = None
-        for i in node.params:
-            if i.unpack and i.unpack.value == "*":
-                vararg = i.gen.py_ast[0]
-            elif i.unpack and i.unpack.value == "**":
-                kwarg = i.gen.py_ast[0]
+        if posonlyargs:
+            posonlyargs = params + posonlyargs
+            params = [i.gen.py_ast[0] for i in node.params]
+        else:
+            params = params + [i.gen.py_ast[0] for i in node.params]
+        defaults = []
+        for i in [*node.posonly_params, *node.params]:
+            if i.value:
+                defaults.append(cast(ast3.expr, i.value.gen.py_ast[0]))
+        kwonly_args = [i.gen.py_ast[0] for i in node.kwonlyargs]
+        # kw_defaults must be the same length as kwonlyargs
+        # it will have None for args that don't have defaults
+        kw_defaults: list[ast3.expr | None] = []
+        for i in node.kwonlyargs:
+            if i.value:
+                kw_defaults.append(cast(ast3.expr, i.value.gen.py_ast[0]))
             else:
-                (
-                    params.append(i.gen.py_ast[0])
-                    if isinstance(i.gen.py_ast[0], ast3.arg)
-                    else self.ice("This list should only be Args")
-                )
-        defaults = [x.value.gen.py_ast[0] for x in node.params if x.value]
+                kw_defaults.append(None)
         node.gen.py_ast = [
             self.sync(
                 ast3.arguments(
-                    posonlyargs=[],
+                    posonlyargs=[cast(ast3.arg, param) for param in posonlyargs],
                     args=[cast(ast3.arg, param) for param in params],
-                    kwonlyargs=[],
+                    kwonlyargs=kwonly_args,
                     vararg=cast(ast3.arg, vararg) if vararg else None,
                     kwarg=cast(ast3.arg, kwarg) if kwarg else None,
-                    kw_defaults=[],
+                    kw_defaults=kw_defaults,
                     defaults=[cast(ast3.expr, default) for default in defaults],
                 )
             )
@@ -1138,7 +1189,7 @@ class PyastGenPass(UniPass):
         elif is_static_var:
             annotation = self.sync(
                 ast3.Subscript(
-                    value=self.sync(ast3.Name(id="ClassVar", ctx=ast3.Load())),
+                    value=self.builtin_name("ClassVar"),
                     slice=cast(ast3.expr, annotation),
                     ctx=ast3.Load(),
                 )
@@ -1215,8 +1266,29 @@ class PyastGenPass(UniPass):
         ]
 
     def exit_typed_ctx_block(self, node: uni.TypedCtxBlock) -> None:
-        # TODO: Come back
-        pass
+        loc = self.sync(
+            ast3.Name(id=Con.HERE.value, ctx=ast3.Load())
+            if node.from_walker
+            else ast3.Name(id=Con.VISITOR.value, ctx=ast3.Load())
+        )
+        node.gen.py_ast = [
+            self.sync(
+                ast3.If(
+                    test=self.sync(
+                        ast3.Call(
+                            func=self.sync(ast3.Name(id="isinstance", ctx=ast3.Load())),
+                            args=[
+                                loc,
+                                cast(ast3.expr, node.type_ctx.gen.py_ast[0]),
+                            ],
+                            keywords=[],
+                        )
+                    ),
+                    body=cast(list[ast3.stmt], self.resolve_stmt_block(node.body)),
+                    orelse=[],
+                )
+            )
+        ]
 
     def exit_if_stmt(self, node: uni.IfStmt) -> None:
         node.gen.py_ast = [
@@ -1467,20 +1539,23 @@ class PyastGenPass(UniPass):
         ]
 
     def exit_assert_stmt(self, node: uni.AssertStmt) -> None:
-        node.gen.py_ast = [
-            self.sync(
-                ast3.Assert(
-                    test=cast(ast3.expr, node.condition.gen.py_ast[0]),
-                    msg=(
-                        cast(ast3.expr, node.error_msg.gen.py_ast[0])
-                        if node.error_msg
-                        else None
-                    ),
+        if isinstance(node.parent, uni.Test):
+            self.assert_helper(node)
+        else:
+            node.gen.py_ast = [
+                self.sync(
+                    ast3.Assert(
+                        test=cast(ast3.expr, node.condition.gen.py_ast[0]),
+                        msg=(
+                            cast(ast3.expr, node.error_msg.gen.py_ast[0])
+                            if node.error_msg
+                            else None
+                        ),
+                    )
                 )
-            )
-        ]
+            ]
 
-    def exit_check_stmt(self, node: uni.CheckStmt) -> None:
+    def assert_helper(self, node: uni.AssertStmt) -> None:
         """Sub objects.
 
         target: ExprType,
@@ -1536,16 +1611,16 @@ class PyastGenPass(UniPass):
 
         # By default the check expression will become assertTrue(<expr>), unless any pattern detected.
         assert_func_name = "assertTrue"
-        assert_args_list = node.target.gen.py_ast
+        assert_args_list = node.condition.gen.py_ast
 
         # Compare operations. Note that We're only considering the compare
         # operation with a single operation ie. a < b < c is  ignored here.
         if (
-            isinstance(node.target, uni.CompareExpr)
-            and isinstance(node.target.gen.py_ast[0], ast3.Compare)
-            and len(node.target.ops) == 1
+            isinstance(node.condition, uni.CompareExpr)
+            and isinstance(node.condition.gen.py_ast[0], ast3.Compare)
+            and len(node.condition.ops) == 1
         ):
-            expr: uni.CompareExpr = node.target
+            expr: uni.CompareExpr = node.condition
             opty: uni.Token = expr.ops[0]
 
             optype2fn = {
@@ -1579,10 +1654,10 @@ class PyastGenPass(UniPass):
                     assert_args_list.pop()
 
         # Check if 'isinstance' is called.
-        elif isinstance(node.target, uni.FuncCall) and isinstance(
-            node.target.gen.py_ast[0], ast3.Call
+        elif isinstance(node.condition, uni.FuncCall) and isinstance(
+            node.condition.gen.py_ast[0], ast3.Call
         ):
-            res = check_node_isinstance_call(node.target)
+            res = check_node_isinstance_call(node.condition)
             if res.isit:
                 # These assertions will make mypy happy.
                 assert isinstance(res.inst, ast3.AST)
@@ -1592,12 +1667,12 @@ class PyastGenPass(UniPass):
 
         # Check if 'not isinstance(<expr>, <expr>)' is called.
         elif (
-            isinstance(node.target, uni.UnaryExpr)
-            and isinstance(node.target, uni.UnaryExpr)
-            and isinstance(node.target.operand, uni.FuncCall)
-            and isinstance(node.target.operand, uni.UnaryExpr)
+            isinstance(node.condition, uni.UnaryExpr)
+            and isinstance(node.condition, uni.UnaryExpr)
+            and isinstance(node.condition.operand, uni.FuncCall)
+            and isinstance(node.condition.operand, uni.UnaryExpr)
         ):
-            res = check_node_isinstance_call(node.target.operand)
+            res = check_node_isinstance_call(node.condition.operand)
             if res.isit:
                 # These assertions will make mypy happy.
                 assert isinstance(res.inst, ast3.AST)
@@ -1610,14 +1685,14 @@ class PyastGenPass(UniPass):
         # the almost equal functionality (snice there is no almost equal operator in jac and never needed ig.).
 
         # Check if 'almostEqual' is called.
-        if isinstance(node.target, uni.FuncCall) and isinstance(
-            node.target.gen.py_ast[0], ast3.Call
+        if isinstance(node.condition, uni.FuncCall) and isinstance(
+            node.condition.gen.py_ast[0], ast3.Call
         ):
-            func = node.target.target
+            func = node.condition.target
             if isinstance(func, uni.Name) and func.value == "almostEqual":
                 assert_func_name = "assertAlmostEqual"
                 assert_args_list = []
-                for param in node.target.params:
+                for param in node.condition.params:
                     assert_args_list.append(param.gen.py_ast[0])
 
         # assert_func_expr = "Con.JAC_CHECK.value.assertXXX"
@@ -1989,7 +2064,7 @@ class PyastGenPass(UniPass):
                 keywords.append(
                     self.sync(
                         ast3.keyword(
-                            arg="filter",
+                            arg="filter_on",
                             value=cast(
                                 ast3.expr,
                                 node.op.edge_spec.filter_cond.gen.py_ast[0],
@@ -2196,7 +2271,7 @@ class PyastGenPass(UniPass):
             node.gen.py_ast = [
                 self.sync(
                     ast3.Call(
-                        func=self.sync(ast3.Name(id="jobj", ctx=ast3.Load())),
+                        func=self.builtin_name("jobj"),
                         args=[],
                         keywords=[
                             self.sync(
@@ -2421,7 +2496,7 @@ class PyastGenPass(UniPass):
             node.gen.py_ast = [
                 self.sync(
                     ast3.Call(
-                        func=self.jaclib_obj("filter"),
+                        func=self.jaclib_obj("filter_on"),
                         args=[],
                         keywords=[
                             self.sync(
@@ -2444,7 +2519,7 @@ class PyastGenPass(UniPass):
             node.gen.py_ast = [
                 self.sync(
                     ast3.Call(
-                        func=self.jaclib_obj("assign"),
+                        func=self.jaclib_obj("assign_all"),
                         args=cast(
                             list[ast3.expr],
                             [node.target.gen.py_ast[0], node.right.gen.py_ast[0]],
@@ -2675,7 +2750,7 @@ class PyastGenPass(UniPass):
                     func=self.sync(
                         ast3.Attribute(
                             value=pynode,
-                            attr=f"_{cur.edge_dir.name.lower()}",
+                            attr=f"edge_{cur.edge_dir.name.lower()}",
                             ctx=ast3.Load(),
                         )
                     ),
@@ -2718,14 +2793,22 @@ class PyastGenPass(UniPass):
                     keywords=[],
                 )
             )
-
-        pynode = self.sync(
-            ast3.Call(
-                func=self.jaclib_obj("refs"),
-                args=[pynode],
-                keywords=[],
+        if node.is_async:
+            pynode = self.sync(
+                ast3.Call(
+                    func=self.jaclib_obj("arefs"),
+                    args=[pynode],
+                    keywords=[],
+                )
+            )
+        else:
+            pynode = self.sync(
+                ast3.Call(
+                    func=self.jaclib_obj("refs"),
+                    args=[pynode],
+                    keywords=[],
+                )
             )
-        )
 
         node.gen.py_ast = [pynode]
 
@@ -3008,6 +3091,9 @@ class PyastGenPass(UniPass):
 
     def exit_name(self, node: uni.Name) -> None:
         name = node.sym_name
+        # Track if this name is a known builtin
+        if name in self.KNOWN_BUILTINS:
+            self.builtin_imports.add(name)
         node.gen.py_ast = [self.sync(ast3.Name(id=name, ctx=node.py_ctx_func()))]
 
     def exit_float(self, node: uni.Float) -> None: