jaclang 0.8.8__py3-none-any.whl → 0.8.10__py3-none-any.whl

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -26,8 +26,8 @@ from typing import List, Optional, Sequence, TypeVar, Union, cast
 
 import jaclang.compiler.unitree as uni
 from jaclang.compiler.constant import Constants as Con, EdgeDir, Tokens as Tok
-from jaclang.compiler.passes import UniPass
-from jaclang.settings import settings
+from jaclang.compiler.passes.ast_gen import BaseAstGenPass
+from jaclang.compiler.passes.ast_gen.jsx_processor import PyJsxProcessor
 
 T = TypeVar("T", bound=ast3.AST)
 
@@ -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,
@@ -89,14 +88,17 @@ UNARY_OP_MAP: dict[Tok, type[ast3.unaryop]] = {
 }
 
 
-class PyastGenPass(UniPass):
+class PyastGenPass(BaseAstGenPass[ast3.AST]):
     """Jac blue transpilation to python pass."""
 
     def before_pass(self) -> None:
-        for i in self.ir_in.impl_mod + self.ir_in.test_mod:
-            PyastGenPass(ir_in=i, prog=self.prog)
+        self.child_passes: list[PyastGenPass] = self._init_child_passes(PyastGenPass)
         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._temp_name_counter: int = 0
+        self._hoisted_funcs: list[ast3.FunctionDef | ast3.AsyncFunctionDef] = []
         self.preamble: list[ast3.AST] = [
             self.sync(
                 ast3.ImportFrom(
@@ -106,44 +108,17 @@ 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,
-                )
-            ),
         ]
+        self.jsx_processor = PyJsxProcessor(self)
 
     def enter_node(self, node: uni.UniNode) -> None:
         """Enter node."""
+        if self._should_skip_client(node):
+            node.gen.py_ast = []
+            if hasattr(node.gen, "py"):
+                node.gen.py = ""  # type: ignore[attr-defined]
+            self.prune()
+            return
         if node.gen.py_ast:
             self.prune()
             return
@@ -151,25 +126,28 @@ class PyastGenPass(UniPass):
 
     def exit_node(self, node: uni.UniNode) -> None:
         """Exit node."""
+        if self._should_skip_client(node):
+            node.gen.py_ast = []
+            if hasattr(node.gen, "py"):
+                node.gen.py = ""  # type: ignore[attr-defined]
+            return
         super().exit_node(node)
-        # for i in node.gen.py_ast:  # Internal validation
-        #     self.node_compilable_test(i)
-
-        # TODO: USE THIS TO SYNC
-        #     if isinstance(i, ast3.AST):
-        #         i.jac_link = node
 
-    def jaclib_obj(self, obj_name: str) -> ast3.Name | ast3.Attribute:
+    def jaclib_obj(self, obj_name: str) -> ast3.Name:
         """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(),
-            )
-        )
+        self.jaclib_imports.add(obj_name)
+        return self.sync(ast3.Name(id=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 +165,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 +190,33 @@ class PyastGenPass(UniPass):
             ),
         )
 
+    def _next_temp_name(self, prefix: str) -> str:
+        """Generate a deterministic temporary name for synthesized definitions."""
+        self._temp_name_counter += 1
+        return f"__jac_{prefix}_{self._temp_name_counter}"
+
+    def _function_expr_from_def(
+        self,
+        func_def: ast3.FunctionDef | ast3.AsyncFunctionDef,
+        jac_node: uni.UniNode,
+        filename_hint: str,
+    ) -> ast3.Name:
+        """Convert a synthesized function definition into an executable expression.
+
+        Instead of using make_block_lambda at runtime, we hoist the function
+        definition to be emitted before the current statement, then return
+        a reference to the function name.
+        """
+        # Ensure locations are fixed
+        ast3.fix_missing_locations(func_def)
+        # Add the function to the hoisted list - it will be emitted before the current statement
+        self._hoisted_funcs.append(func_def)
+        # Return a reference to the function name
+        return self.sync(
+            ast3.Name(id=func_def.name, ctx=ast3.Load()),
+            jac_node=jac_node,
+        )
+
     def _get_sem_decorator(self, node: uni.UniNode) -> ast3.Call | None:
         """Create a semstring decorator for the given semantic strings.
 
@@ -283,7 +279,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(
@@ -308,15 +307,9 @@ class PyastGenPass(UniPass):
             )
         )
 
-    def flatten(self, body: list[T | list[T] | None]) -> list[T]:
-        """Flatten a list of items or lists into a single list."""
-        new_body: list[T] = []
-        for item in body:
-            if isinstance(item, list):
-                new_body.extend(item)
-            elif item is not None:
-                new_body.append(item)
-        return new_body
+    def _should_skip_client(self, node: uni.UniNode) -> bool:
+        """Check if node is a client-facing declaration that should skip Python codegen."""
+        return isinstance(node, uni.ClientFacingNode) and node.is_client_decl
 
     def sync(
         self, py_node: T, jac_node: Optional[uni.UniNode] = None, deep: bool = False
@@ -373,7 +366,7 @@ class PyastGenPass(UniPass):
             [self.sync(ast3.Pass())]
             if isinstance(node, Sequence) and not valid_stmts
             else (
-                self.flatten(
+                self._flatten_ast_list(
                     [
                         x.gen.py_ast
                         for x in valid_stmts
@@ -417,31 +410,55 @@ class PyastGenPass(UniPass):
         node.gen.py_ast = node.tag.gen.py_ast
 
     def exit_module(self, node: uni.Module) -> None:
-        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:
-            pre_body = [*pre_body, *pbody.body]
-        pre_body = [*pre_body, *clean_body]
-        for pbody in node.test_mod:
-            pre_body = [*pre_body, *pbody.body]
-        body = (
-            [
+        # Merge jaclib and builtin imports from child passes
+        for child_pass in self.child_passes:
+            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.Expr(value=cast(ast3.expr, node.doc.gen.py_ast[0])),
-                    jac_node=node.doc,
-                ),
-                *self.preamble,
-                *[x.gen.py_ast for x in pre_body],
-            ]
-            if node.doc
-            else [*self.preamble, *[x.gen.py_ast for x in pre_body]]
-        )
-        new_body = []
-        for i in body:
-            if isinstance(i, list):
-                new_body += i
-            else:
-                new_body.append(i) if i else None
+                    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 imports at the end of preamble
+        if 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,
+                )
+            )
+
+        merged_body = self._merge_module_bodies(node)
+        body_items: list[ast3.AST | list[ast3.AST] | None] = [*self.preamble]
+        body_items.extend(item.gen.py_ast for item in merged_body)
+
+        if node.doc:
+            doc_stmt = self.sync(
+                ast3.Expr(value=cast(ast3.expr, node.doc.gen.py_ast[0])),
+                jac_node=node.doc,
+            )
+            body_items.insert(0, doc_stmt)
+
+        new_body = self._flatten_ast_list(body_items)
         node.gen.py_ast = [
             self.sync(
                 ast3.Module(
@@ -458,7 +475,7 @@ class PyastGenPass(UniPass):
                 ast3.Expr(value=cast(ast3.expr, node.doc.gen.py_ast[0])),
                 jac_node=node.doc,
             )
-            assigns_ast: list[ast3.AST] = self.flatten(
+            assigns_ast: list[ast3.AST] = self._flatten_ast_list(
                 [a.gen.py_ast for a in node.assignments]
             )
             if isinstance(doc, ast3.AST):
@@ -466,7 +483,9 @@ class PyastGenPass(UniPass):
             else:
                 raise self.ice()
         else:
-            node.gen.py_ast = self.flatten([a.gen.py_ast for a in node.assignments])
+            node.gen.py_ast = self._flatten_ast_list(
+                [a.gen.py_ast for a in node.assignments]
+            )
 
     def exit_test(self, node: uni.Test) -> None:
         test_name = node.name.sym_name
@@ -687,18 +706,17 @@ class PyastGenPass(UniPass):
         if node.arch_type.name != Tok.KW_CLASS:
             base_classes.append(self.jaclib_obj(node.arch_type.value.capitalize()))
 
-        node.gen.py_ast = [
-            self.sync(
-                ast3.ClassDef(
-                    name=node.name.sym_name,
-                    bases=[cast(ast3.expr, i) for i in base_classes],
-                    keywords=[],
-                    body=[cast(ast3.stmt, i) for i in body],
-                    decorator_list=[cast(ast3.expr, i) for i in decorators],
-                    type_params=[],
-                )
+        class_def = self.sync(
+            ast3.ClassDef(
+                name=node.name.sym_name,
+                bases=[cast(ast3.expr, i) for i in base_classes],
+                keywords=[],
+                body=[cast(ast3.stmt, i) for i in body],
+                decorator_list=[cast(ast3.expr, i) for i in decorators],
+                type_params=[],
             )
-        ]
+        )
+        node.gen.py_ast = [class_def]
 
     def enter_enum(self, node: uni.Enum) -> None:
         if isinstance(node.body, uni.ImplDef):
@@ -722,19 +740,18 @@ 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())))
-        node.gen.py_ast = [
-            self.sync(
-                ast3.ClassDef(
-                    name=node.name.sym_name,
-                    bases=[cast(ast3.expr, i) for i in base_classes],
-                    keywords=[],
-                    body=[cast(ast3.stmt, i) for i in body],
-                    decorator_list=[cast(ast3.expr, i) for i in decorators],
-                    type_params=[],
-                )
+        base_classes.append(self.builtin_name("Enum"))
+        class_def = self.sync(
+            ast3.ClassDef(
+                name=node.name.sym_name,
+                bases=[cast(ast3.expr, i) for i in base_classes],
+                keywords=[],
+                body=[cast(ast3.stmt, i) for i in body],
+                decorator_list=[cast(ast3.expr, i) for i in decorators],
+                type_params=[],
             )
-        ]
+        )
+        node.gen.py_ast = [class_def]
 
     def enter_ability(self, node: uni.Ability) -> None:
         if isinstance(node.body, uni.ImplDef):
@@ -744,23 +761,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 +931,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 +948,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()))
@@ -962,36 +965,35 @@ class PyastGenPass(UniPass):
         if isinstance(node.signature, uni.FuncSignature) and node.signature.return_type:
             ast_returns = cast(ast3.expr, node.signature.return_type.gen.py_ast[0])
 
-        node.gen.py_ast = [
-            self.sync(
-                func_type(
-                    name=node.name_ref.sym_name,
-                    args=(
-                        cast(ast3.arguments, node.signature.gen.py_ast[0])
-                        if node.signature
-                        else self.sync(
-                            ast3.arguments(
-                                posonlyargs=[],
-                                args=(
-                                    [self.sync(ast3.arg(arg="self", annotation=None))]
-                                    if node.is_method
-                                    else []
-                                ),
-                                vararg=None,
-                                kwonlyargs=[],
-                                kw_defaults=[],
-                                kwarg=None,
-                                defaults=[],
-                            )
+        func_def = self.sync(
+            func_type(
+                name=node.name_ref.sym_name,
+                args=(
+                    cast(ast3.arguments, node.signature.gen.py_ast[0])
+                    if node.signature
+                    else self.sync(
+                        ast3.arguments(
+                            posonlyargs=[],
+                            args=(
+                                [self.sync(ast3.arg(arg="self", annotation=None))]
+                                if node.is_method
+                                else []
+                            ),
+                            vararg=None,
+                            kwonlyargs=[],
+                            kw_defaults=[],
+                            kwarg=None,
+                            defaults=[],
                         )
-                    ),
-                    body=[cast(ast3.stmt, i) for i in body],
-                    decorator_list=[cast(ast3.expr, i) for i in decorator_list],
-                    returns=ast_returns,
-                    type_params=[],
-                )
+                    )
+                ),
+                body=[cast(ast3.stmt, i) for i in body],
+                decorator_list=[cast(ast3.expr, i) for i in decorator_list],
+                returns=ast_returns,
+                type_params=[],
             )
-        ]
+        )
+        node.gen.py_ast = [func_def]
 
     def exit_impl_def(self, node: uni.ImplDef) -> None:
         pass
@@ -1111,7 +1113,9 @@ class PyastGenPass(UniPass):
             ]
 
     def exit_arch_has(self, node: uni.ArchHas) -> None:
-        vars_py: list[ast3.AST] = self.flatten([v.gen.py_ast for v in node.vars])
+        vars_py: list[ast3.AST] = self._flatten_ast_list(
+            [v.gen.py_ast for v in node.vars]
+        )
         if node.doc:
             doc = self.sync(
                 ast3.Expr(value=cast(ast3.expr, node.doc.gen.py_ast[0])),
@@ -1145,7 +1149,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(),
                 )
@@ -1222,8 +1226,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 = [
@@ -1259,19 +1284,24 @@ class PyastGenPass(UniPass):
         node.gen.py_ast = self.resolve_stmt_block(node.body)
 
     def exit_expr_stmt(self, node: uni.ExprStmt) -> None:
-        node.gen.py_ast = [
-            (
-                self.sync(ast3.Expr(value=cast(ast3.expr, node.expr.gen.py_ast[0])))
-                if not node.in_fstring
-                else self.sync(
-                    ast3.FormattedValue(
-                        value=cast(ast3.expr, node.expr.gen.py_ast[0]),
-                        conversion=-1,
-                        format_spec=None,
-                    )
+        # Collect any hoisted functions that were generated during expression evaluation
+        hoisted = self._hoisted_funcs[:]
+        self._hoisted_funcs.clear()
+
+        expr_stmt = (
+            self.sync(ast3.Expr(value=cast(ast3.expr, node.expr.gen.py_ast[0])))
+            if not node.in_fstring
+            else self.sync(
+                ast3.FormattedValue(
+                    value=cast(ast3.expr, node.expr.gen.py_ast[0]),
+                    conversion=-1,
+                    format_spec=None,
                 )
             )
-        ]
+        )
+
+        # Emit hoisted functions before the expression statement
+        node.gen.py_ast = [*hoisted, expr_stmt]
 
     def exit_concurrent_expr(self, node: uni.ConcurrentExpr) -> None:
         func = ""
@@ -1708,7 +1738,7 @@ class PyastGenPass(UniPass):
                     value=self.sync(
                         self.sync(
                             ast3.Call(
-                                func=self.jaclib_obj("report"),
+                                func=self.jaclib_obj("log_report"),
                                 args=cast(list[ast3.expr], node.expr.gen.py_ast),
                                 keywords=[],
                             )
@@ -1865,48 +1895,49 @@ class PyastGenPass(UniPass):
         )
         targets_ast = [cast(ast3.expr, t.gen.py_ast[0]) for t in node.target]
 
+        # Collect any hoisted functions that were generated during expression evaluation
+        hoisted = self._hoisted_funcs[:]
+        self._hoisted_funcs.clear()
+
         if node.type_tag:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.AnnAssign(
-                        target=cast(ast3.Name, targets_ast[0]),
-                        annotation=cast(ast3.expr, node.type_tag.gen.py_ast[0]),
-                        value=(
-                            cast(ast3.expr, node.value.gen.py_ast[0])
-                            if node.value
-                            else None
-                        ),
-                        simple=int(isinstance(targets_ast[0], ast3.Name)),
-                    )
+            assignment_stmt: ast3.AnnAssign | ast3.Assign | ast3.AugAssign = self.sync(
+                ast3.AnnAssign(
+                    target=cast(ast3.Name, targets_ast[0]),
+                    annotation=cast(ast3.expr, node.type_tag.gen.py_ast[0]),
+                    value=(
+                        cast(ast3.expr, node.value.gen.py_ast[0])
+                        if node.value
+                        else None
+                    ),
+                    simple=int(isinstance(targets_ast[0], ast3.Name)),
                 )
-            ]
+            )
         elif node.aug_op:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.AugAssign(
-                        target=cast(ast3.Name, targets_ast[0]),
-                        op=cast(ast3.operator, node.aug_op.gen.py_ast[0]),
-                        value=(
-                            cast(ast3.expr, value)
-                            if isinstance(value, ast3.expr)
-                            else ast3.Constant(value=None)
-                        ),
-                    )
+            assignment_stmt = self.sync(
+                ast3.AugAssign(
+                    target=cast(ast3.Name, targets_ast[0]),
+                    op=cast(ast3.operator, node.aug_op.gen.py_ast[0]),
+                    value=(
+                        cast(ast3.expr, value)
+                        if isinstance(value, ast3.expr)
+                        else ast3.Constant(value=None)
+                    ),
                 )
-            ]
+            )
         else:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.Assign(
-                        targets=cast(list[ast3.expr], targets_ast),
-                        value=(
-                            cast(ast3.expr, value)
-                            if isinstance(value, ast3.expr)
-                            else ast3.Constant(value=None)
-                        ),
-                    )
+            assignment_stmt = self.sync(
+                ast3.Assign(
+                    targets=cast(list[ast3.expr], targets_ast),
+                    value=(
+                        cast(ast3.expr, value)
+                        if isinstance(value, ast3.expr)
+                        else ast3.Constant(value=None)
+                    ),
                 )
-            ]
+            )
+
+        # Emit hoisted functions before the assignment
+        node.gen.py_ast = [*hoisted, assignment_stmt]
 
     def exit_binary_expr(self, node: uni.BinaryExpr) -> None:
         if isinstance(node.op, uni.ConnectOp):
@@ -1999,7 +2030,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],
@@ -2134,28 +2165,80 @@ class PyastGenPass(UniPass):
         ]
 
     def exit_lambda_expr(self, node: uni.LambdaExpr) -> None:
-        # Python lambda expressions don't support type annotations
+        # Multi-statement lambda emits a synthesized function definition that is
+        # hoisted before the current statement.
+        if isinstance(node.body, list):
+            if node.signature and node.signature.gen.py_ast:
+                arguments = cast(ast3.arguments, node.signature.gen.py_ast[0])
+            else:
+                arguments = self.sync(
+                    ast3.arguments(
+                        posonlyargs=[],
+                        args=[],
+                        kwonlyargs=[],
+                        kw_defaults=[],
+                        defaults=[],
+                    ),
+                    jac_node=node,
+                )
+            body_stmts = [
+                cast(ast3.stmt, stmt)
+                for stmt in self.resolve_stmt_block(node.body, doc=None)
+            ]
+            if not body_stmts:
+                body_stmts = [self.sync(ast3.Pass(), jac_node=node)]
+            func_name = self._next_temp_name("lambda")
+            returns = (
+                cast(ast3.expr, node.signature.return_type.gen.py_ast[0])
+                if node.signature
+                and node.signature.return_type
+                and node.signature.return_type.gen.py_ast
+                else None
+            )
+            func_def = self.sync(
+                ast3.FunctionDef(
+                    name=func_name,
+                    args=arguments,
+                    body=body_stmts,
+                    decorator_list=[],
+                    returns=returns,
+                    type_params=[],
+                ),
+                jac_node=node,
+            )
+            node.gen.py_ast = [
+                self._function_expr_from_def(
+                    func_def=func_def,
+                    jac_node=node,
+                    filename_hint=f"<jac_lambda:{func_name}>",
+                )
+            ]
+            return
+
+        # Single-expression lambda maps directly to Python lambda; strip annotations.
         if node.signature:
             self._remove_lambda_param_annotations(node.signature)
-
+        if node.signature and node.signature.gen.py_ast:
+            arguments = cast(ast3.arguments, node.signature.gen.py_ast[0])
+        else:
+            arguments = self.sync(
+                ast3.arguments(
+                    posonlyargs=[],
+                    args=[],
+                    kwonlyargs=[],
+                    kw_defaults=[],
+                    defaults=[],
+                ),
+                jac_node=node,
+            )
+        body_expr = cast(ast3.expr, node.body.gen.py_ast[0])
         node.gen.py_ast = [
             self.sync(
                 ast3.Lambda(
-                    args=(
-                        cast(ast3.arguments, node.signature.gen.py_ast[0])
-                        if node.signature
-                        else self.sync(
-                            ast3.arguments(
-                                posonlyargs=[],
-                                args=[],
-                                kwonlyargs=[],
-                                kw_defaults=[],
-                                defaults=[],
-                            )
-                        )
-                    ),
-                    body=cast(ast3.expr, node.body.gen.py_ast[0]),
-                )
+                    args=arguments,
+                    body=body_expr,
+                ),
+                jac_node=node,
             )
         ]
 
@@ -2206,7 +2289,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(
@@ -2239,9 +2322,8 @@ class PyastGenPass(UniPass):
             for i in str_seq:
                 if isinstance(i, uni.String):
                     pieces.append(i.lit_value)
-                elif isinstance(i, uni.FString):
-                    pieces.extend(get_pieces(i.parts)) if i.parts else None
-                elif isinstance(i, uni.ExprStmt):
+                elif isinstance(i, (uni.FString, uni.ExprStmt)):
+                    # pieces.extend(get_pieces(i.parts)) if i.parts else None
                     pieces.append(i.gen.py_ast[0])
                 elif isinstance(i, uni.Token) and i.name in [Tok.LBRACE, Tok.RBRACE]:
                     continue
@@ -2281,11 +2363,32 @@ class PyastGenPass(UniPass):
             node.gen.py_ast = [combined_multi[0]]
 
     def exit_f_string(self, node: uni.FString) -> None:
-        py_parts: list[list[ast3.AST]] = [
-            cast(list[ast3.AST], p.gen.py_ast) for p in node.parts
+        node.gen.py_ast = [
+            self.sync(
+                ast3.JoinedStr(
+                    values=[
+                        cast(ast3.expr, part.gen.py_ast[0])
+                        for part in node.parts
+                        if part.gen.py_ast
+                    ],
+                )
+            )
+        ]
+
+    def exit_formatted_value(self, node: uni.FormattedValue) -> None:
+        node.gen.py_ast = [
+            self.sync(
+                ast3.FormattedValue(
+                    value=cast(ast3.expr, node.format_part.gen.py_ast[0]),
+                    conversion=node.conversion,
+                    format_spec=(
+                        cast(ast3.expr, node.format_spec.gen.py_ast[0])
+                        if node.format_spec
+                        else None
+                    ),
+                )
+            )
         ]
-        parts = self.flatten(cast(list[list[ast3.AST] | ast3.AST | None], py_parts))
-        node.gen.py_ast = parts if parts else [self.sync(ast3.Constant(value=""))]
 
     def exit_list_val(self, node: uni.ListVal) -> None:
         elts = [cast(ast3.expr, v.gen.py_ast[0]) for v in node.values]
@@ -2431,7 +2534,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(
@@ -2454,7 +2557,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]],
@@ -2501,6 +2604,24 @@ class PyastGenPass(UniPass):
             ]
 
     def exit_atom_unit(self, node: uni.AtomUnit) -> None:
+        if (
+            isinstance(node.value, uni.Ability)
+            and node.value.gen.py_ast
+            and isinstance(
+                node.value.gen.py_ast[0], (ast3.FunctionDef, ast3.AsyncFunctionDef)
+            )
+        ):
+            func_ast = cast(
+                ast3.FunctionDef | ast3.AsyncFunctionDef, node.value.gen.py_ast[0]
+            )
+            node.gen.py_ast = [
+                self._function_expr_from_def(
+                    func_def=func_ast,
+                    jac_node=node,
+                    filename_hint=f"<jac_iife:{func_ast.name}>",
+                )
+            ]
+            return
         node.gen.py_ast = node.value.gen.py_ast
 
     def gen_call_args(
@@ -2649,7 +2770,7 @@ class PyastGenPass(UniPass):
 
         pynode = self.sync(
             ast3.Call(
-                func=self.jaclib_obj("Path"),
+                func=self.jaclib_obj("OPath"),
                 args=[cast(ast3.expr, origin or cur.gen.py_ast[0])],
                 keywords=[],
             )
@@ -2685,7 +2806,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(),
                         )
                     ),
@@ -3026,6 +3147,11 @@ class PyastGenPass(UniPass):
 
     def exit_name(self, node: uni.Name) -> None:
         name = node.sym_name
+        # Track if this name is a known builtin
+        import jaclang.runtimelib.builtin
+
+        if name in set(jaclang.runtimelib.builtin.__all__):
+            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:
@@ -3051,6 +3177,38 @@ class PyastGenPass(UniPass):
     def exit_ellipsis(self, node: uni.Ellipsis) -> None:
         node.gen.py_ast = [self.sync(ast3.Constant(value=...))]
 
+    def exit_jsx_element(self, node: uni.JsxElement) -> None:
+        """Generate Python AST for JSX elements."""
+        self.jsx_processor.element(node)
+
+    def exit_jsx_element_name(self, node: uni.JsxElementName) -> None:
+        """Generate Python AST for JSX element names."""
+        self.jsx_processor.element_name(node)
+
+    def exit_jsx_attribute(self, node: uni.JsxAttribute) -> None:
+        """Bases class for JSX attributes - handled by subclasses."""
+        pass
+
+    def exit_jsx_spread_attribute(self, node: uni.JsxSpreadAttribute) -> None:
+        """Generate Python AST for JSX spread attributes."""
+        self.jsx_processor.spread_attribute(node)
+
+    def exit_jsx_normal_attribute(self, node: uni.JsxNormalAttribute) -> None:
+        """Generate Python AST for JSX normal attributes."""
+        self.jsx_processor.normal_attribute(node)
+
+    def exit_jsx_child(self, node: uni.JsxChild) -> None:
+        """Bases class for JSX children - handled by subclasses."""
+        pass
+
+    def exit_jsx_text(self, node: uni.JsxText) -> None:
+        """Generate Python AST for JSX text nodes."""
+        self.jsx_processor.text(node)
+
+    def exit_jsx_expression(self, node: uni.JsxExpression) -> None:
+        """Generate Python AST for JSX expression children."""
+        self.jsx_processor.expression(node)
+
     def exit_semi(self, node: uni.Semi) -> None:
         pass