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

jaclang/compiler/passes/main/pyast_load_pass.py

@@ -16,6 +16,8 @@ from __future__ import annotations
 
 import ast as py_ast
 import os
+import re
+from threading import Event
 from typing import Optional, Sequence, TYPE_CHECKING, TypeAlias, TypeVar, cast
 
 import jaclang.compiler.unitree as uni
@@ -32,11 +34,16 @@ T = TypeVar("T", bound=uni.UniNode)
 class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
     """Jac Parser."""
 
-    def __init__(self, ir_in: uni.PythonModuleAst, prog: JacProgram) -> None:
+    def __init__(
+        self,
+        ir_in: uni.PythonModuleAst,
+        prog: JacProgram,
+        cancel_token: Event | None = None,
+    ) -> None:
         """Initialize parser."""
         self.mod_path = ir_in.loc.mod_path
         self.orig_src = ir_in.loc.orig_src
-        Transform.__init__(self, ir_in=ir_in, prog=prog)
+        Transform.__init__(self, ir_in=ir_in, prog=prog, cancel_token=cancel_token)
 
     def nu(self, node: T) -> T:
         """Update node."""
@@ -45,6 +52,8 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
 
     def convert(self, node: py_ast.AST) -> uni.UniNode:
         """Get python node type."""
+        if self.is_canceled():
+            raise StopIteration
         if hasattr(self, f"proc_{pascal_to_snake(type(node).__name__)}"):
             ret = getattr(self, f"proc_{pascal_to_snake(type(node).__name__)}")(node)
         else:
@@ -1160,7 +1169,7 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
         else:
             raise self.ice()
 
-    def proc_formatted_value(self, node: py_ast.FormattedValue) -> uni.ExprStmt:
+    def proc_formatted_value(self, node: py_ast.FormattedValue) -> uni.FormattedValue:
         """Process python node.
 
         class FormattedValue(expr):
@@ -1171,10 +1180,15 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
         format_spec: expr | None
         """
         value = self.convert(node.value)
+        if node.format_spec:
+            fmt_spec = cast(uni.Expr, self.convert(node.format_spec))
+        else:
+            fmt_spec = None
         if isinstance(value, uni.Expr):
-            ret = uni.ExprStmt(
-                expr=value,
-                in_fstring=True,
+            ret = uni.FormattedValue(
+                format_part=value,
+                conversion=node.conversion,
+                format_spec=fmt_spec,
                 kid=[value],
             )
         else:
@@ -1369,13 +1383,34 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
             __match_args__ = ("values",)
         values: list[expr]
         """
-        values = [self.convert(value) for value in node.values]
-        valid = [
-            value for value in values if isinstance(value, (uni.String, uni.ExprStmt))
-        ]
+        valid: list[uni.Token | uni.FormattedValue] = []
+        for i in node.values:
+            if isinstance(i, py_ast.Constant) and isinstance(i.value, str):
+                valid.append(self.operator(Tok.STRING, i.value))
+            elif isinstance(i, py_ast.FormattedValue):
+                converted = self.convert(i)
+                if isinstance(converted, uni.FormattedValue):
+                    valid.append(converted)
+            else:
+                raise self.ice("Invalid node in joined str")
+        ast_seg = py_ast.get_source_segment(self.orig_src.code, node)
+        if ast_seg is None:
+            ast_seg = 'f""'
+        match = re.match(r"(?i)(fr|rf|f)('{3}|\"{3}|'|\")", ast_seg)
+        if match:
+            prefix, quote = match.groups()
+            start = match.group(0)
+            end = quote * (3 if len(quote) == 3 else 1)
+        else:
+            start = "f'"
+            end = "'"
+        tok_start = self.operator(Tok.STRING, start)
+        tok_end = self.operator(Tok.STRING, end)
         fstr = uni.FString(
+            start=tok_start,
             parts=valid,
-            kid=[*valid] if valid else [uni.EmptyToken()],
+            end=tok_end,
+            kid=[tok_start, *valid, tok_end] if valid else [uni.EmptyToken()],
         )
         return uni.MultiString(strings=[fstr], kid=[fstr])
 

jaclang/compiler/passes/main/pyjac_ast_link_pass.py

@@ -27,6 +27,8 @@ class PyJacAstLinkPass(UniPass):
         self, jac_node: uni.UniNode, py_nodes: list[ast3.AST]
     ) -> None:
         """Link jac name ast to py ast nodes."""
+        if isinstance(jac_node, uni.ClientFacingNode) and jac_node.is_client_decl:
+            return
         jac_node.gen.py_ast = py_nodes
         for i in py_nodes:
             if isinstance(i.jac_link, list):  # type: ignore

jaclang/compiler/passes/main/sym_tab_build_pass.py

@@ -19,6 +19,7 @@ type checking, and semantic analysis throughout the compilation process.
 """
 
 import jaclang.compiler.unitree as uni
+from jaclang.compiler.constant import SymbolAccess
 from jaclang.compiler.passes import UniPass
 from jaclang.compiler.unitree import UniScopeNode
 
@@ -75,7 +76,7 @@ class SymTabBuildPass(UniPass):
     def exit_module_path(self, node: uni.ModulePath) -> None:
         if node.alias:
             node.alias.sym_tab.def_insert(node.alias, single_decl="import")
-        elif node.path:
+        elif node.path and not node.is_import_from:
             if node.parent_of_type(uni.Import) and not (
                 node.parent_of_type(uni.Import).from_loc
                 and node.parent_of_type(uni.Import).is_jac
@@ -84,9 +85,25 @@ class SymTabBuildPass(UniPass):
         else:
             pass  # Need to support pythonic import symbols with dots in it
 
+        # There will be symbols for
+        # import from math {sqrt}  <- math will have a symbol but no symtab entry
+        # import math as m  <- m will have a symbol and symtab entry
+        if node.path and (node.is_import_from or (node.alias)):
+            for n in node.path:
+                n.sym = n.create_symbol(
+                    access=SymbolAccess.PUBLIC,
+                    imported=True,
+                )
+
     def exit_module_item(self, node: uni.ModuleItem) -> None:
         sym_node = node.alias or node.name
         sym_node.sym_tab.def_insert(sym_node, single_decl="import")
+        if node.alias:
+            # create symbol for module item
+            node.name.sym = node.name.create_symbol(
+                access=SymbolAccess.PUBLIC,
+                imported=True,
+            )
 
     def enter_archetype(self, node: uni.Archetype) -> None:
         self.push_scope_and_link(node)

jaclang/compiler/passes/main/tests/fixtures/autoimpl.cl.jac

@@ -0,0 +1,7 @@
+node Button {
+    has title: str;
+}
+
+def component() {
+    print("client component");
+}

jaclang/compiler/passes/main/tests/fixtures/checker_arity.jac

@@ -15,4 +15,7 @@ with entry {
     f.with_default_args(1, 2);    # <-- Ok
     f.with_default_args(1, 2, 3); # <-- Error
     f.with_default_args();        # <-- Error
+
+    # Argument unpacking
+    f.with_default_args(*[1, 2]); # <-- Ok
 }

jaclang/compiler/passes/main/tests/fixtures/checker_class_construct.jac

@@ -0,0 +1,33 @@
+glob RAD = 5.0;
+
+obj Circle1{
+    has radius:float ,age:int;
+    def init(color:str){
+        self.color = color;
+    }
+}
+with entry {
+    c1 = Circle1(RAD);
+}
+# ---------------------------------------------------------------
+glob length = 5.0;
+
+obj Square{
+    has side_length:float ,age:int;
+}
+
+with entry {
+    c2 = Square(length);
+}
+# ---------------------------------------------------------------
+
+glob name = "John";
+
+obj Person{
+    has name:str ,age:int=90;
+}
+
+with entry {
+    c = Person(name=name, age=25);
+    c = Person();
+}

jaclang/compiler/passes/main/tests/fixtures/defuse_modpath.jac

@@ -0,0 +1,7 @@
+import from math { sqrt as square_root }
+import from concurrent.futures { ThreadPoolExecutor }
+
+with entry {
+    square_root(16);
+    # math;
+}

jaclang/compiler/passes/main/tests/fixtures/member_access_type_resolve.jac

@@ -5,7 +5,8 @@ node Foo {
   has bar: Bar;
 }
 with entry {
-  f: Foo = Foo();
+  bar_obj: Bar = Bar(23);
+  f: Foo = Foo(bar=bar_obj);
   i: int = f.bar.baz; # <-- Ok
   s: str = f.bar.baz; # <-- Error
 }

jaclang/compiler/passes/main/tests/test_checker_pass.py

@@ -1,8 +1,6 @@
 
 """Tests for typechecker pass (the pyright implementation)."""
 
-from tempfile import NamedTemporaryFile
-
 from jaclang.utils.test import TestCase
 from jaclang.compiler.passes.main import TypeCheckPass
 from jaclang.compiler.program import JacProgram
@@ -270,6 +268,37 @@ class TypeCheckerPassTests(TestCase):
         for i, expected in enumerate(expected_errors):
             self._assert_error_pretty_found(expected, program.errors_had[i].pretty_print())
 
+    def test_class_construct(self) -> None:
+        program = JacProgram()
+        path = self.fixture_abs_path("checker_class_construct.jac")
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 3)
+
+        expected_errors = [
+            """
+            Cannot assign <class float> to parameter 'color' of type <class str>
+                    with entry {
+                        c1 = Circle1(RAD);
+                                    ^^^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'age'
+                    with entry {
+                    c2 = Square(length);
+                         ^^^^^^^^^^^^^^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'name'
+                    c = Person(name=name, age=25);
+                    c = Person();
+                        ^^^^^^^^
+            """,
+        ]
+
+        for i, expected in enumerate(expected_errors):
+            self._assert_error_pretty_found(expected, program.errors_had[i].pretty_print())
+
     def test_self_type_inference(self) -> None:
         path = self.fixture_abs_path("checker_self_type.jac")
         program = JacProgram()

jaclang/compiler/passes/main/tests/test_def_use_pass.py

@@ -21,3 +21,15 @@ class DefUsePassTests(TestCase):
         self.assertEqual(len(uses[1]), 1)
         self.assertIn("output", [uses[0][0].sym_name, uses[1][0].sym_name])
         self.assertIn("message", [uses[0][0].sym_name, uses[1][0].sym_name])
+    
+    def test_def_use_modpath(self) -> None:
+        """Basic test for pass."""
+        state = JacProgram().compile(
+            file_path=self.fixture_abs_path("defuse_modpath.jac")
+        )
+        all_symbols = list(
+            state.sym_tab.names_in_scope.values()
+        )
+        self.assertEqual(len(all_symbols), 2)
+        self.assertEqual(all_symbols[0].sym_name, "square_root")
+        self.assertEqual(all_symbols[1].sym_name, "ThreadPoolExecutor")

jaclang/compiler/passes/main/tests/test_import_pass.py

@@ -30,29 +30,31 @@ class ImportPassPassTests(TestCase):
         (prog := JacProgram()).compile(self.fixture_abs_path("autoimpl.jac"))
         num_modules = len(list(prog.mod.hub.values())[0].impl_mod)
         mod_names = [i.name for i in list(prog.mod.hub.values())[0].impl_mod]
-        self.assertEqual(num_modules, 4)
+        self.assertEqual(num_modules, 5)
         self.assertIn("getme.impl", mod_names)
         self.assertIn("autoimpl.impl", mod_names)
         self.assertIn("autoimpl.something.else.impl", mod_names)
+        self.assertIn("autoimpl.cl", mod_names)
 
     def test_import_include_auto_impl(self) -> None:
         """Basic test for pass."""
         (prog := JacProgram()).build(self.fixture_abs_path("incautoimpl.jac"))
         num_modules = len(list(prog.mod.hub.values())[1].impl_mod) + 1
         mod_names = [i.name for i in list(prog.mod.hub.values())[1].impl_mod]
-        self.assertEqual(num_modules, 5)
+        self.assertEqual(num_modules, 6)
         self.assertEqual("incautoimpl", list(prog.mod.hub.values())[0].name)
         self.assertEqual("autoimpl", list(prog.mod.hub.values())[1].name)
         self.assertIn("getme.impl", mod_names)
         self.assertIn("autoimpl.impl", mod_names)
         self.assertIn("autoimpl.something.else.impl", mod_names)
+        self.assertIn("autoimpl.cl", mod_names)
 
     def test_annexalbe_by_discovery(self) -> None:
         """Basic test for pass."""
         (prog := JacProgram()).build(self.fixture_abs_path("incautoimpl.jac"))
         count = 0
         all_mods = prog.mod.hub.values()
-        self.assertEqual(len(all_mods), 6)
+        self.assertEqual(len(all_mods), 7)
         for main_mod in all_mods:
             for i in main_mod.impl_mod:
                 if i.name not in ["autoimpl", "incautoimpl"]:
@@ -60,7 +62,24 @@ class ImportPassPassTests(TestCase):
                     self.assertEqual(
                         i.annexable_by, self.fixture_abs_path("autoimpl.jac")
                     )
-        self.assertEqual(count, 4)
+        self.assertEqual(count, 5)
+
+    def test_cl_annex_marked_client(self) -> None:
+        """Ensure .cl.jac annex files are autoloaded and marked client."""
+
+        (prog := JacProgram()).compile(self.fixture_abs_path("autoimpl.jac"))
+        main_mod = list(prog.mod.hub.values())[0]
+        cl_mod = next(
+            (mod for mod in main_mod.impl_mod if mod.name.endswith(".cl")), None
+        )
+        self.assertIsNotNone(cl_mod, "Expected .cl annex module to be loaded")
+        abilities = cl_mod.get_all_sub_nodes(uni.Ability)
+        self.assertTrue(abilities, "Expected abilities in .cl annex module")
+        for ability in abilities:
+            self.assertTrue(
+                ability.is_client_decl,
+                "All client annex abilities should be marked as client declarations",
+            )
 
     @unittest.skip("TODO: Fix when we have the type checker")
     def test_py_raise_map(self) -> None:

jaclang/compiler/passes/main/tests/test_pyast_gen_pass.py

@@ -153,6 +153,31 @@ class PyastGenPassTests(TestCaseMicroSuite, AstSyncTestMixin):
 
         self.assertFalse(out.errors_had)
 
+    def test_iife_fixture_executes(self) -> None:
+        """Ensure IIFE and block lambdas lower to executable Python."""
+        fixture_path = self.lang_fixture_abs_path("iife_functions.jac")
+        code_gen = (prog := JacProgram()).compile(fixture_path)
+        self.assertFalse(prog.errors_had)
+        if code_gen.gen.py_ast and isinstance(code_gen.gen.py_ast[0], ast3.Module):
+            module_ast = code_gen.gen.py_ast[0]
+            compiled = compile(module_ast, filename="<ast>", mode="exec")
+            captured = io.StringIO()
+            original_stdout = sys.stdout
+            try:
+                sys.stdout = captured
+                module = types.ModuleType("__main__")
+                module.__dict__["__file__"] = code_gen.loc.mod_path
+                exec(compiled, module.__dict__)
+            finally:
+                sys.stdout = original_stdout
+            output = captured.getvalue()
+            self.assertIn("Test 1 - Basic IIFE: 42", output)
+            self.assertIn(
+                "Test 6 - IIFE returning function, adder(5): 15",
+                output,
+            )
+            self.assertIn("All IIFE tests completed!", output)
+
     def parent_scrub(self, node: uni.UniNode) -> bool:
         """Validate every node has parent."""
         success = True

jaclang/compiler/passes/main/type_checker_pass.py

@@ -69,10 +69,17 @@ class TypeCheckPass(UniPass):
 
     def exit_import(self, node: uni.Import) -> None:
         """Exit an import node."""
+        # import from math {sqrt, sin as s}
         if node.from_loc:
+            self.evaluator.get_type_of_module(node.from_loc)
             for item in node.items:
                 if isinstance(item, uni.ModuleItem):
                     self.evaluator.get_type_of_module_item(item)
+        else:
+            # import math as m, os, sys;
+            for item in node.items:
+                if isinstance(item, uni.ModulePath):
+                    self.evaluator.get_type_of_module(item)
 
     def exit_assignment(self, node: uni.Assignment) -> None:
         """Pyright: Checker.visitAssignment(node: AssignmentNode): boolean."""

jaclang/compiler/passes/tool/doc_ir_gen_pass.py

@@ -892,6 +892,13 @@ class DocIRGenPass(UniPass):
             parts.append(i.gen.doc_ir)
         node.gen.doc_ir = self.group(self.concat(parts))
 
+    def exit_formatted_value(self, node: uni.FormattedValue) -> None:
+        """Generate DocIR for formatted value expressions."""
+        parts: list[doc.DocType] = []
+        for i in node.kid:
+            parts.append(i.gen.doc_ir)
+        node.gen.doc_ir = self.group(self.concat(parts))
+
     def exit_if_else_expr(self, node: uni.IfElseExpr) -> None:
         """Generate DocIR for conditional expressions."""
         parts: list[doc.DocType] = []
@@ -1747,3 +1754,111 @@ class DocIRGenPass(UniPass):
             node.gen.doc_ir = self.group(
                 self.concat([self.text(node.value), self.hard_line()])
             )
+
+    def exit_jsx_element(self, node: uni.JsxElement) -> None:
+        """Generate DocIR for JSX elements - kid-centric beautiful formatting!"""
+        parts: list[doc.DocType] = []
+        prev = None
+
+        # Check if we have any JSX element children
+        # Use node.children instead of node.kid to avoid counting opening/closing tags
+        has_jsx_elem_children = any(
+            isinstance(k, uni.JsxElement) for k in node.children
+        )
+
+        # Only break/indent if we have JSX element children
+        # (simple text/expression children stay inline)
+        should_format_children = has_jsx_elem_children
+
+        for i in node.kid:
+            # Add line break between attributes (allows them to wrap nicely)
+            if (
+                prev
+                and isinstance(prev, (uni.JsxElementName, uni.JsxAttribute))
+                and isinstance(i, uni.JsxAttribute)
+            ):
+                parts.append(self.line())
+            # Add hard line between JSX element children, or before first child
+            elif (
+                prev
+                and (
+                    (
+                        isinstance(prev, (uni.JsxChild, uni.JsxElement))
+                        and isinstance(i, (uni.JsxChild, uni.JsxElement))
+                    )
+                    or (
+                        isinstance(prev, (uni.JsxElementName, uni.JsxAttribute))
+                        and isinstance(i, (uni.JsxChild, uni.JsxElement))
+                    )
+                )
+                and should_format_children
+            ):
+                parts.append(self.hard_line())
+
+            # Indent JSX element children, but not text/expression children
+            if isinstance(i, uni.JsxElement) and should_format_children:
+                parts.append(self.indent(i.gen.doc_ir))
+            else:
+                parts.append(i.gen.doc_ir)
+
+            prev = i
+
+        node.gen.doc_ir = self.group(self.concat(parts))
+
+    def exit_jsx_element_name(self, node: uni.JsxElementName) -> None:
+        """Generate DocIR for JSX element names."""
+        parts: list[doc.DocType] = []
+        for i in node.kid:
+            parts.append(i.gen.doc_ir)
+        node.gen.doc_ir = self.concat(parts)
+
+    def exit_jsx_spread_attribute(self, node: uni.JsxSpreadAttribute) -> None:
+        """Generate DocIR for JSX spread attributes."""
+        parts: list[doc.DocType] = []
+        for i in node.kid:
+            parts.append(i.gen.doc_ir)
+        node.gen.doc_ir = self.concat(parts)
+
+    def exit_jsx_normal_attribute(self, node: uni.JsxNormalAttribute) -> None:
+        """Generate DocIR for JSX normal attributes."""
+        # Normalize to ensure LBRACE/RBRACE tokens are added for expression values
+        node.normalize()
+        parts: list[doc.DocType] = []
+        for i in node.kid:
+            # Tokens created by normalize() have empty doc_ir, so regenerate it
+            if (
+                isinstance(i, uni.Token)
+                and isinstance(i.gen.doc_ir, doc.Text)
+                and not i.gen.doc_ir.text
+            ):
+                i.gen.doc_ir = self.text(i.value)
+            elif not isinstance(
+                i.gen.doc_ir,
+                (
+                    doc.Text,
+                    doc.Concat,
+                    doc.Group,
+                    doc.Indent,
+                    doc.Line,
+                    doc.Align,
+                    doc.IfBreak,
+                ),
+            ):
+                # For nodes with invalid doc_ir, generate it by visiting
+                self.enter_exit(i)
+            parts.append(i.gen.doc_ir)
+        node.gen.doc_ir = self.concat(parts)
+
+    def exit_jsx_text(self, node: uni.JsxText) -> None:
+        """Generate DocIR for JSX text."""
+        parts: list[doc.DocType] = []
+        for i in node.kid:
+            parts.append(i.gen.doc_ir)
+        node.gen.doc_ir = self.concat(parts)
+
+    def exit_jsx_expression(self, node: uni.JsxExpression) -> None:
+        """Generate DocIR for JSX expressions."""
+        parts: list[doc.DocType] = []
+        for i in node.kid:
+            parts.append(i.gen.doc_ir)
+        node.gen.doc_ir = self.concat(parts)

jaclang/compiler/passes/tool/fuse_comments_pass.py

@@ -41,19 +41,9 @@ class FuseCommentsPass(UniPass):
         except StopIteration:
             next_comment = None
 
-        # Handle possible leading comments
-        if next_comment and (not next_code or _is_before(next_comment, next_code)):
-            self.ir_out.src_terminals.insert(0, next_comment)
-
         # Merge streams in order
         while next_comment or next_code:
             if next_comment and (not next_code or _is_before(next_comment, next_code)):
-                # Add comment token
-                if merged and (last_token := merged[-1]):
-                    self.ir_out.src_terminals.insert(
-                        self.ir_out.src_terminals.index(last_token) + 1,
-                        next_comment,
-                    )
                 merged.append(next_comment)
                 try:
                     next_comment = next(comments)
@@ -66,6 +56,7 @@ class FuseCommentsPass(UniPass):
                     next_code = next(code_tokens)
                 except StopIteration:
                     next_code = None
+        self.ir_out.src_terminals[:] = merged
 
         return merged
 

jaclang/compiler/passes/tool/tests/test_jac_format_pass.py

@@ -133,6 +133,8 @@ class JacFormatPassTests(TestCaseMicroSuite):
                     self.assertEqual(tokens[i + 1], "{")
             self.assertEqual(num_test, 3)
             return
+        before = ""
+        after = ""
         try:
             before = ast3.dump(code_gen_pure.gen.py_ast[0], indent=2)
             after = ast3.dump(code_gen_jac.gen.py_ast[0], indent=2)
@@ -151,7 +153,8 @@ class JacFormatPassTests(TestCaseMicroSuite):
             print("\n+++++++++++++++++++++++++++++++++++++++\n")
             print(add_line_numbers(code_gen_format))
             print("\n+++++++++++++++++++++++++++++++++++++++\n")
-            print("\n".join(unified_diff(before.splitlines(), after.splitlines())))
+            if before and after:
+                print("\n".join(unified_diff(before.splitlines(), after.splitlines())))
             raise e
 
 

jaclang/compiler/passes/transform.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import time
 from abc import ABC, abstractmethod
+from threading import Event
 from typing import Generic, Optional, TYPE_CHECKING, Type, TypeVar
 
 from jaclang.compiler.codeinfo import CodeLocInfo
@@ -70,7 +71,9 @@ class Alert:
 class Transform(ABC, Generic[T, R]):
     """Abstract class for IR passes."""
 
-    def __init__(self, ir_in: T, prog: JacProgram) -> None:
+    def __init__(
+        self, ir_in: T, prog: JacProgram, cancel_token: Event | None = None
+    ) -> None:
         """Initialize pass."""
         self.logger = logging.getLogger(self.__class__.__name__)
         self.errors_had: list[Alert] = []
@@ -79,6 +82,7 @@ class Transform(ABC, Generic[T, R]):
         self.prog = prog
         self.time_taken = 0.0
         self.ir_in: T = ir_in
+        self.cancel_token = cancel_token
         self.pre_transform()
         self.ir_out: R = self.timed_transform(ir_in=ir_in)
         self.post_transform()
@@ -136,6 +140,10 @@ class Transform(ABC, Generic[T, R]):
         """Log info."""
         self.logger.info(msg)
 
+    def is_canceled(self) -> bool:
+        """Check if the pass has been canceled."""
+        return self.cancel_token is not None and self.cancel_token.is_set()
+
     def ice(self, msg: str = "Something went horribly wrong!") -> RuntimeError:
         """Pass Error."""
         self.log_error(f"ICE: Pass {self.__class__.__name__} - {msg}")

jaclang/compiler/passes/uni_pass.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from threading import Event
 from typing import Optional, TYPE_CHECKING, Type, TypeVar
 
 import jaclang.compiler.unitree as uni
@@ -21,11 +22,12 @@ class UniPass(Transform[uni.Module, uni.Module]):
         self,
         ir_in: uni.Module,
         prog: JacProgram,
+        cancel_token: Event | None = None,
     ) -> None:
         """Initialize parser."""
         self.term_signal = False
         self.prune_signal = False
-        Transform.__init__(self, ir_in, prog)
+        Transform.__init__(self, ir_in, prog, cancel_token=cancel_token)
 
     def before_pass(self) -> None:
         """Run once before pass."""
@@ -43,10 +45,6 @@ class UniPass(Transform[uni.Module, uni.Module]):
         if hasattr(self, f"exit_{pascal_to_snake(type(node).__name__)}"):
             getattr(self, f"exit_{pascal_to_snake(type(node).__name__)}")(node)
 
-    def terminate(self) -> None:
-        """Terminate traversal."""
-        self.term_signal = True
-
     def prune(self) -> None:
         """Prune traversal."""
         self.prune_signal = True
@@ -120,7 +118,7 @@ class UniPass(Transform[uni.Module, uni.Module]):
 
     def traverse(self, node: uni.UniNode) -> uni.UniNode:
         """Traverse tree."""
-        if self.term_signal:
+        if self.is_canceled():
             return node
         self.cur_node = node
         self.enter_node(node)
@@ -131,7 +129,7 @@ class UniPass(Transform[uni.Module, uni.Module]):
         else:
             self.prune_signal = False
         self.cur_node = node
-        if self.term_signal:
+        if self.is_canceled():
             return node
         self.exit_node(node)
         return node