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

jaclang/compiler/codeloc.py

@@ -22,6 +22,13 @@ class CodeGenTarget:
     mypy_ast: list[MypyNode] = field(default_factory=lambda: [])
     py_bytecode: Optional[bytes] = None
 
+    def clean(self) -> None:
+        """Clean code generation target."""
+        self.py = ""
+        self.jac = ""
+        self.py_ast = []
+        self.mypy_ast = []
+
 
 class CodeLocInfo:
     """Code location info."""

jaclang/compiler/compile.py

@@ -16,7 +16,7 @@ def compile_jac(file_path: str, cache_result: bool = False) -> Pass:
         file_path=file_path,
         schedule=pass_schedule,
     )
-    if cache_result and isinstance(code.ir, ast.Module) and not code.errors_had:
+    if cache_result and isinstance(code.ir, ast.Module):
         print_pass = PyOutPass(input_ir=code.ir, prior=code)
         return print_pass
     else:

jaclang/compiler/constant.py

@@ -236,6 +236,23 @@ class Tokens(str, Enum):
         return self.value
 
 
+DELIM_MAP = {
+    Tokens.COMMA: ",",
+    Tokens.EQ: "=",
+    Tokens.DECOR_OP: "@",
+    Tokens.WS: "\n",
+    Tokens.SEMI: ";",
+    Tokens.COLON: ":",
+    Tokens.LBRACE: "{",
+    Tokens.RBRACE: "}",
+    Tokens.LSQUARE: "[",
+    Tokens.RSQUARE: "]",
+    Tokens.LPAREN: "(",
+    Tokens.RPAREN: ")",
+    Tokens.RETURN_HINT: "->",
+    Tokens.DOT: ".",
+}
+
 colors = [
     "#FFE9E9",
     "#F0FFF0",

jaclang/compiler/parser.py

@@ -390,6 +390,7 @@ class JacParser(Pass):
             """
             ret = ast.SubNodeList[ast.ModuleItem](
                 items=[i for i in kid if isinstance(i, ast.ModuleItem)],
+                delim=Tok.COMMA,
                 kid=kid,
             )
             return self.nu(ret)
@@ -509,6 +510,7 @@ class JacParser(Pass):
                 return self.nu(
                     ast.SubNodeList[ast.Expr](
                         items=valid_decors,
+                        delim=Tok.DECOR_OP,
                         kid=kid,
                     )
                 )
@@ -525,6 +527,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.Expr](
                     items=valid_inh,
+                    delim=Tok.COMMA,
                     kid=kid,
                 )
             )
@@ -596,6 +599,7 @@ class JacParser(Pass):
             if isinstance(kid[0], (ast.Enum, ast.EnumDef)):
                 return self.nu(kid[0])
             else:
+
                 raise self.ice()
 
         def enum_decl(self, kid: list[ast.AstNode]) -> ast.Enum:
@@ -655,37 +659,65 @@ class JacParser(Pass):
 
             enum_block: LBRACE ((enum_stmt COMMA)* enum_stmt)? RBRACE
             """
-            ret = ast.SubNodeList[ast.EnumBlockStmt](
-                items=[],
-                kid=kid,
-            )
+            ret = ast.SubNodeList[ast.EnumBlockStmt](items=[], delim=Tok.COMMA, kid=kid)
             ret.items = [i for i in kid if isinstance(i, ast.EnumBlockStmt)]
             return self.nu(ret)
 
         def enum_stmt(self, kid: list[ast.AstNode]) -> ast.EnumBlockStmt:
             """Grammar rule.
 
-            enum_stmt: NAME EQ expression
-                    | NAME
+            enum_stmt: NAME (COLON STRING)? EQ expression
+                    | NAME (COLON STRING)?
                     | py_code_block
             """
             if isinstance(kid[0], ast.PyInlineCode):
                 return self.nu(kid[0])
             if isinstance(kid[0], (ast.Name)):
-                if len(kid) == 1:
-                    kid[0].is_enum_singleton = True
-                    return self.nu(kid[0])
-                elif isinstance(kid[2], ast.Expr):
-                    targ = ast.SubNodeList[ast.Expr](items=[kid[0]], kid=[kid[0]])
+                if (
+                    len(kid) >= 3
+                    and isinstance(kid[-1], ast.Expr)
+                    and not isinstance(kid[-1], ast.String)
+                ):
+                    semstr = (
+                        kid[2]
+                        if len(kid) > 3 and isinstance(kid[2], ast.String)
+                        else None
+                    )
+                    targ = ast.SubNodeList[ast.Expr](
+                        items=[kid[0]], delim=Tok.COMMA, kid=[kid[0]]
+                    )
+                    kid[0] = targ
+                    return self.nu(
+                        ast.Assignment(
+                            target=targ,
+                            value=kid[-1],
+                            type_tag=None,
+                            kid=kid,
+                            semstr=semstr,
+                            is_enum_stmt=True,
+                        )
+                    )
+                else:
+                    semstr = (
+                        kid[2]
+                        if len(kid) == 3 and isinstance(kid[2], ast.String)
+                        else None
+                    )
+                    targ = ast.SubNodeList[ast.Expr](
+                        items=[kid[0]], delim=Tok.COMMA, kid=[kid[0]]
+                    )
                     kid[0] = targ
                     return self.nu(
                         ast.Assignment(
                             target=targ,
-                            value=kid[2],
+                            value=None,
                             type_tag=None,
                             kid=kid,
+                            semstr=semstr,
+                            is_enum_stmt=True,
                         )
                     )
+
             raise self.ice()
 
         def ability(
@@ -748,7 +780,6 @@ class JacParser(Pass):
             chomp = chomp[1:] if semstr else chomp
             name = chomp[0]
             chomp = chomp[1:]
-            is_func = isinstance(chomp[0], ast.FuncSignature)
             signature = chomp[0]
             chomp = chomp[1:]
             body = chomp[0] if isinstance(chomp[0], ast.SubNodeList) else None
@@ -758,7 +789,6 @@ class JacParser(Pass):
                 return self.nu(
                     ast.Ability(
                         name_ref=name,
-                        is_func=is_func,
                         is_async=False,
                         is_override=is_override,
                         is_static=is_static,
@@ -818,7 +848,6 @@ class JacParser(Pass):
             chomp = chomp[1:] if semstr else chomp
             name = chomp[0]
             chomp = chomp[1:]
-            is_func = isinstance(chomp[0], ast.FuncSignature)
             signature = chomp[0]
             chomp = chomp[1:]
             if isinstance(name, ast.NameSpec) and isinstance(
@@ -827,7 +856,6 @@ class JacParser(Pass):
                 return self.nu(
                     ast.Ability(
                         name_ref=name,
-                        is_func=is_func,
                         is_async=False,
                         is_override=is_override,
                         is_static=is_static,
@@ -846,7 +874,7 @@ class JacParser(Pass):
             """Grammar rule.
 
             genai_ability: KW_OVERRIDE? KW_STATIC? KW_CAN access_tag? STRING?
-            any_ref (func_decl) KW_WITH atomic_call SEMI
+            any_ref (func_decl) KW_BY atomic_call SEMI
             """
             chomp = [*kid]
             is_override = (
@@ -864,22 +892,19 @@ class JacParser(Pass):
             chomp = chomp[1:] if semstr else chomp
             name = chomp[0]
             chomp = chomp[1:]
-            is_func = isinstance(chomp[0], ast.FuncSignature)
             signature = chomp[0]
             chomp = chomp[1:]
-            has_with = isinstance(chomp[0], ast.Token) and chomp[0].name == Tok.KW_WITH
-            chomp = chomp[1:] if has_with else chomp
-            is_funccall = isinstance(chomp[0], ast.FuncCall)
+            has_by = isinstance(chomp[0], ast.Token) and chomp[0].name == Tok.KW_BY
+            chomp = chomp[1:] if has_by else chomp
             if (
                 isinstance(name, ast.NameSpec)
                 and isinstance(signature, (ast.FuncSignature, ast.EventSignature))
-                and is_funccall
-                and has_with
+                and isinstance(chomp[0], ast.FuncCall)
+                and has_by
             ):
                 return self.nu(
                     ast.Ability(
                         name_ref=name,
-                        is_func=is_func,
                         is_async=False,
                         is_override=is_override,
                         is_static=is_static,
@@ -887,7 +912,7 @@ class JacParser(Pass):
                         access=access,
                         semstr=semstr,
                         signature=signature,
-                        body=chomp[0],  # type: ignore
+                        body=chomp[0],
                         kid=kid,
                     )
                 )
@@ -959,6 +984,7 @@ class JacParser(Pass):
             """
             ret = ast.SubNodeList[ast.ParamVar](
                 items=[i for i in kid if isinstance(i, ast.ParamVar)],
+                delim=Tok.COMMA,
                 kid=kid,
             )
             return self.nu(ret)
@@ -1014,9 +1040,12 @@ class JacParser(Pass):
             """
             ret = ast.SubNodeList[ast.ArchBlockStmt](
                 items=[],
+                delim=Tok.WS,
                 kid=kid,
             )
             ret.items = [i for i in kid if isinstance(i, ast.ArchBlockStmt)]
+            ret.left_enc = kid[0] if isinstance(kid[0], ast.Token) else None
+            ret.right_enc = kid[-1] if isinstance(kid[-1], ast.Token) else None
             return self.nu(ret)
 
         def member_stmt(self, kid: list[ast.AstNode]) -> ast.ArchBlockStmt:
@@ -1092,6 +1121,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.HasVar](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=new_kid,
                 )
             )
@@ -1161,7 +1191,6 @@ class JacParser(Pass):
                         col_end=kid[0].loc.col_end,
                         pos_start=kid[0].pos_start,
                         pos_end=kid[0].pos_end,
-                        kid=kid[0].kid,
                     )
                 )
             else:
@@ -1172,32 +1201,18 @@ class JacParser(Pass):
         ) -> ast.SubNodeList[ast.CodeBlockStmt]:
             """Grammar rule.
 
-            code_block: LBRACE statement_list* RBRACE
-            """
-            if isinstance(kid[1], ast.SubNodeList):
-                kid[1].add_kids_left([kid[0]])
-                kid[1].add_kids_right([kid[2]])
-                return self.nu(kid[1])
-            else:
-                return self.nu(
-                    ast.SubNodeList[ast.CodeBlockStmt](
-                        items=[],
-                        kid=kid,
-                    )
-                )
-
-        def statement_list(
-            self, kid: list[ast.AstNode]
-        ) -> ast.SubNodeList[ast.CodeBlockStmt]:
-            """Grammar rule.
-
-            statement_list: statement+
+            code_block: LBRACE statement* RBRACE
             """
+            left_enc = kid[0] if isinstance(kid[0], ast.Token) else None
+            right_enc = kid[-1] if isinstance(kid[-1], ast.Token) else None
             valid_stmt = [i for i in kid if isinstance(i, ast.CodeBlockStmt)]
-            if len(valid_stmt) == len(kid):
+            if len(valid_stmt) == len(kid) - 2:
                 return self.nu(
                     ast.SubNodeList[ast.CodeBlockStmt](
                         items=valid_stmt,
+                        delim=Tok.WS,
+                        left_enc=left_enc,
+                        right_enc=right_enc,
                         kid=kid,
                     )
                 )
@@ -1367,6 +1382,7 @@ class JacParser(Pass):
                 return self.nu(
                     ast.SubNodeList[ast.Except](
                         items=valid_kid,
+                        delim=Tok.WS,
                         kid=kid,
                     )
                 )
@@ -1515,6 +1531,7 @@ class JacParser(Pass):
             """
             ret = ast.SubNodeList[ast.ExprAsItem](
                 items=[i for i in kid if isinstance(i, ast.ExprAsItem)],
+                delim=Tok.COMMA,
                 kid=kid,
             )
             return self.nu(ret)
@@ -1737,70 +1754,68 @@ class JacParser(Pass):
         def assignment(self, kid: list[ast.AstNode]) -> ast.Assignment:
             """Grammar rule.
 
-            assignment: KW_LET? (atomic_chain EQ)+ (yield_stmt | expression)
-                    | atomic_chain type_tag (EQ (yield_stmt | expression))?
-                    | atomic_chain aug_op (yield_stmt | expression)
+            assignment: KW_LET? (atomic_chain EQ)+ (yield_expr | expression)
+                    | atomic_chain (COLON STRING)? type_tag (EQ (yield_expr | expression))?
+                    | atomic_chain aug_op (yield_expr | expression)
             """
             chomp = [*kid]
             is_frozen = isinstance(chomp[0], ast.Token) and chomp[0].name == Tok.KW_LET
             is_aug = None
             assignees = []
             chomp = chomp[1:] if is_frozen else chomp
-            if (
-                len(chomp) > 1
-                and isinstance(chomp[1], ast.Token)
-                and chomp[1].name != Tok.EQ
-            ):
-                assignees += [chomp[0]]
-                is_aug = chomp[1]
-                chomp = chomp[2:]
-            elif (
-                len(chomp) > 1
-                and isinstance(chomp[1], ast.Token)
-                and chomp[1].name == Tok.EQ
-            ):
-                while (
-                    isinstance(chomp[0], ast.Expr)
-                    and len(chomp) > 1
-                    and isinstance(chomp[1], ast.Token)
-                    and chomp[1].name == Tok.EQ
-                ):
-                    assignees += [chomp[0], chomp[1]]
-                    chomp = chomp[2:]
-            elif isinstance(chomp[0], ast.Expr):
-                assignees += [chomp[0]]
-                chomp = chomp[1:]
-            else:
-                raise self.ice()
+            value = chomp[-1] if isinstance(chomp[-1], ast.Expr) else None
+            chomp = (
+                chomp[:-2]
+                if value and isinstance(chomp[-3], ast.SubTag)
+                else chomp[:-1] if value else chomp
+            )
+            type_tag = chomp[-1] if isinstance(chomp[-1], ast.SubTag) else None
+            if not value:
+                semstr = chomp[2] if len(chomp) > 2 else None
+                chomp = chomp[:-2] if semstr else chomp
+            else:
+                if type_tag:
+                    chomp = chomp[:-1]
+                    semstr = (
+                        chomp[-1]
+                        if len(chomp) > 1 and isinstance(chomp[-1], ast.String)
+                        else None
+                    )
+                    chomp = chomp[:-2] if semstr else chomp
+                else:
+                    semstr = None
+                    if (
+                        isinstance(chomp[1], ast.Token)
+                        and chomp[1].name != Tok.EQ
+                        and chomp[1].name != Tok.COLON
+                    ):
+                        assignees += [chomp[0]]
+                        is_aug = chomp[1]
+                        chomp = chomp[2:]
+                    else:
+                        while (
+                            len(chomp) > 1
+                            and isinstance(chomp[0], ast.Expr)
+                            and isinstance(chomp[1], ast.Token)
+                            and chomp[1].name == Tok.EQ
+                        ):
+                            assignees += [chomp[0], chomp[1]]
+                            chomp = chomp[2:]
+
+            assignees += chomp
             valid_assignees = [i for i in assignees if isinstance(i, (ast.Expr))]
             new_targ = ast.SubNodeList[ast.Expr](
                 items=valid_assignees,
+                delim=Tok.EQ,
                 kid=assignees,
             )
             kid = [x for x in kid if x not in assignees]
             kid.insert(1, new_targ) if is_frozen else kid.insert(0, new_targ)
-            type_tag = (
-                chomp[0]
-                if len(chomp) > 0 and isinstance(chomp[0], ast.SubTag)
-                else None
-            )
-            chomp = chomp[1:] if type_tag else chomp
-            if (
-                len(chomp) > 0
-                and isinstance(chomp[0], ast.Token)
-                and chomp[0].name == Tok.EQ
-            ):
-                chomp = chomp[1:]
-            value = (
-                chomp[0]
-                if len(chomp) > 0 and isinstance(chomp[0], (ast.YieldExpr, ast.Expr))
-                else None
-            )
             if is_aug:
                 return self.nu(
                     ast.Assignment(
                         target=new_targ,
-                        type_tag=type_tag,
+                        type_tag=type_tag if isinstance(type_tag, ast.SubTag) else None,
                         value=value,
                         mutable=is_frozen,
                         aug_op=is_aug,
@@ -1810,10 +1825,11 @@ class JacParser(Pass):
             return self.nu(
                 ast.Assignment(
                     target=new_targ,
-                    type_tag=type_tag,
+                    type_tag=type_tag if isinstance(type_tag, ast.SubTag) else None,
                     value=value,
                     mutable=is_frozen,
                     kid=kid,
+                    semstr=semstr if isinstance(semstr, ast.String) else None,
                 )
             )
 
@@ -2280,7 +2296,7 @@ class JacParser(Pass):
                         target=target,
                         right=chomp[0],
                         is_null_ok=is_null_ok,
-                        is_attr=None,
+                        is_attr=False,
                         kid=kid,
                     )
                 )
@@ -2295,7 +2311,7 @@ class JacParser(Pass):
                         target=(target if chomp[0].name != Tok.DOT_BKWD else chomp[1]),
                         right=(chomp[1] if chomp[0].name != Tok.DOT_BKWD else target),
                         is_null_ok=is_null_ok,
-                        is_attr=chomp[0],
+                        is_attr=True,
                         kid=kid,
                     )
                 )
@@ -2333,7 +2349,7 @@ class JacParser(Pass):
                         expr = index.values.items[0] if index.values else None
                     else:
                         sublist = ast.SubNodeList[ast.Expr | ast.KWPair](
-                            items=[*index.values.items], kid=index.kid
+                            items=[*index.values.items], delim=Tok.COMMA, kid=index.kid
                         )
                         expr = ast.TupleVal(values=sublist, kid=[sublist])
                         kid = [expr]
@@ -2392,7 +2408,7 @@ class JacParser(Pass):
                     and isinstance(kid[1], (ast.Expr, ast.YieldExpr))
                     and isinstance(kid[2], ast.Token)
                 ):
-                    ret = ast.AtomUnit(value=kid[1], is_paren=True, kid=kid)
+                    ret = ast.AtomUnit(value=kid[1], kid=kid)
                     return self.nu(ret)
                 else:
                     raise self.ice()
@@ -2509,6 +2525,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.String | ast.ExprStmt](
                     items=valid_parts,
+                    delim=None,
                     kid=valid_parts,
                 )
             )
@@ -2532,6 +2549,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.String | ast.ExprStmt](
                     items=valid_parts,
+                    delim=None,
                     kid=valid_parts,
                 )
             )
@@ -2622,6 +2640,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.Expr](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=new_kid,
                 )
             )
@@ -2646,6 +2665,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.KWPair](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=new_kid,
                 )
             )
@@ -2698,6 +2718,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.Name](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=new_kid,
                 )
             )
@@ -2732,6 +2753,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.Expr | ast.KWPair](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=kid,
                 )
             )
@@ -2847,7 +2869,7 @@ class JacParser(Pass):
         def inner_compr(self, kid: list[ast.AstNode]) -> ast.InnerCompr:
             """Grammar rule.
 
-            inner_compr: KW_ASYNC? KW_FOR atomic_chain KW_IN walrus_assign (KW_IF expression)?
+            inner_compr: KW_ASYNC? KW_FOR atomic_chain KW_IN pipe_call (KW_IF walrus_assign)*
             """
             chomp = [*kid]
             is_async = bool(
@@ -2862,7 +2884,7 @@ class JacParser(Pass):
                         target=chomp[0],
                         collection=chomp[2],
                         conditional=(
-                            chomp[4]
+                            [i for i in chomp[4:] if isinstance(i, ast.Expr)]
                             if len(chomp) > 4 and isinstance(chomp[4], ast.Expr)
                             else None
                         ),
@@ -2898,6 +2920,7 @@ class JacParser(Pass):
                     return self.nu(
                         ast.SubNodeList[ast.Expr | ast.KWPair](
                             items=valid_kid,
+                            delim=Tok.COMMA,
                             kid=kid,
                         )
                     )
@@ -2927,6 +2950,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.Assignment](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=new_kid,
                 )
             )
@@ -3181,8 +3205,7 @@ class JacParser(Pass):
         def edge_ref_chain(self, kid: list[ast.AstNode]) -> ast.EdgeRefTrailer:
             """Grammar rule.
 
-            edge_ref_chain: (EDGE_OP|NODE_OP)? LSQUARE expression?
-                (edge_op_ref (NODE_OP? expression)?)+ RSQUARE
+            (EDGE_OP|NODE_OP)? LSQUARE expression? (edge_op_ref (filter_compr | expression)?)+ RSQUARE
             """
             valid_chain = [i for i in kid if isinstance(i, (ast.Expr, ast.FilterCompr))]
             return self.nu(
@@ -3249,7 +3272,7 @@ class JacParser(Pass):
         def connect_op(self, kid: list[ast.AstNode]) -> ast.ConnectOp:
             """Grammar rule.
 
-            connect_op: EDGE_OP? (connect_from | connect_to)
+            connect_op: connect_from | connect_to | connect_any
             """
             if len(kid) < 2 and isinstance(kid[0], ast.ConnectOp):
                 return self.nu(kid[0])
@@ -3394,6 +3417,7 @@ class JacParser(Pass):
             return self.nu(
                 ast.SubNodeList[ast.CompareExpr](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=new_kid,
                 )
             )
@@ -3434,7 +3458,7 @@ class JacParser(Pass):
         def assign_compr(self, kid: list[ast.AstNode]) -> ast.AssignCompr:
             """Grammar rule.
 
-            filter_compr: LPAREN STAR_MUL kw_expr_list RPAREN
+            filter_compr: LPAREN EQ kw_expr_list RPAREN
             """
             if isinstance(kid[2], ast.SubNodeList):
                 return self.nu(
@@ -3470,11 +3494,9 @@ class JacParser(Pass):
             """
             pattern = kid[1]
             guard = kid[3] if len(kid) > 4 else None
-            stmts = kid[-1]
-            if (
-                isinstance(pattern, ast.MatchPattern)
-                and isinstance(guard, (ast.Expr, type(None)))
-                and isinstance(stmts, ast.SubNodeList)
+            stmts = [i for i in kid if isinstance(i, ast.CodeBlockStmt)]
+            if isinstance(pattern, ast.MatchPattern) and isinstance(
+                guard, (ast.Expr, type(None))
             ):
                 return self.nu(
                     ast.MatchCase(
@@ -3737,6 +3759,7 @@ class JacParser(Pass):
             valid_kid = [i for i in new_kid if isinstance(i, ast.MatchPattern)]
             return ast.SubNodeList[ast.MatchPattern](
                 items=valid_kid,
+                delim=Tok.COMMA,
                 kid=kid,
             )
 
@@ -3782,6 +3805,7 @@ class JacParser(Pass):
                 valid_kid = [i for i in new_kid if isinstance(i, ast.MatchKVPair)]
                 return ast.SubNodeList[ast.MatchKVPair](
                     items=valid_kid,
+                    delim=Tok.COMMA,
                     kid=new_kid,
                 )
             else:
@@ -3802,7 +3826,6 @@ class JacParser(Pass):
                         pos_start=token.start_pos if token.start_pos is not None else 0,
                         pos_end=token.end_pos if token.end_pos is not None else 0,
                         is_kwesc=True,
-                        kid=[],
                     )
                 )
             elif token.type == Tok.NAME:
@@ -3841,6 +3864,5 @@ class JacParser(Pass):
                     col_end=token.end_column if token.end_column is not None else 0,
                     pos_start=token.start_pos if token.start_pos is not None else 0,
                     pos_end=token.end_pos if token.end_pos is not None else 0,
-                    kid=[],
                 )
             )

jaclang/compiler/passes/ir_pass.py

@@ -68,6 +68,24 @@ class Pass(Transform[ast.T]):
                     result.extend(Pass.get_all_sub_nodes(i, typ, brute_force))
         return result
 
+    @staticmethod
+    def has_parent_of_type(node: ast.AstNode, typ: Type[ast.T]) -> Optional[ast.T]:
+        """Check if node has parent of type."""
+        while node.parent:
+            if isinstance(node.parent, typ):
+                return node.parent
+            node = node.parent
+        return None
+
+    @staticmethod
+    def has_parent_of_node(node: ast.AstNode, parent: ast.AstNode) -> bool:
+        """Check if node has parent of type."""
+        while node.parent:
+            if node.parent == parent:
+                return True
+            node = node.parent
+        return False
+
     def recalculate_parents(self, node: ast.AstNode) -> None:
         """Recalculate parents."""
         if not node:

jaclang/compiler/passes/main/__init__.py

@@ -10,6 +10,7 @@ from .pyast_load_pass import PyastBuildPass  # type: ignore # noqa: I100
 from .pyast_gen_pass import PyastGenPass  # noqa: I100
 from .schedules import py_code_gen  # noqa: I100
 from .type_check_pass import JacTypeCheckPass  # noqa: I100
+from .registry_pass import RegistryPass  # noqa: I100
 
 
 pass_schedule = py_code_gen
@@ -24,4 +25,5 @@ __all__ = [
     "PyastBuildPass",
     "PyastGenPass",
     "JacTypeCheckPass",
+    "RegistryPass",
 ]