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

jaclang/compiler/parser.py

@@ -1772,6 +1772,7 @@ class JacParser(Pass):
             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]
@@ -1782,6 +1783,7 @@ class JacParser(Pass):
                     )
                     chomp = chomp[:-2] if semstr else chomp
                 else:
+                    semstr = None
                     if (
                         isinstance(chomp[1], ast.Token)
                         and chomp[1].name != Tok.EQ
@@ -1827,6 +1829,7 @@ class JacParser(Pass):
                     value=value,
                     mutable=is_frozen,
                     kid=kid,
+                    semstr=semstr if isinstance(semstr, ast.String) else None,
                 )
             )
 

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",
 ]

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -11,6 +11,7 @@ from typing import Optional, Sequence, TypeVar
 import jaclang.compiler.absyntree as ast
 from jaclang.compiler.constant import Constants as Con, EdgeDir, Tokens as Tok
 from jaclang.compiler.passes import Pass
+from jaclang.core.utils import get_sem_scope
 
 T = TypeVar("T", bound=ast3.AST)
 
@@ -208,8 +209,16 @@ class PyastGenPass(Pass):
             if isinstance(i, ast3.AST):
                 i.lineno = jac_node.loc.first_line
                 i.col_offset = jac_node.loc.col_start
-                i.end_lineno = jac_node.loc.last_line
-                i.end_col_offset = jac_node.loc.col_end
+                i.end_lineno = (
+                    jac_node.loc.last_line
+                    if jac_node.loc.last_line
+                    else jac_node.loc.first_line
+                )
+                i.end_col_offset = (
+                    jac_node.loc.col_end
+                    if jac_node.loc.col_end
+                    else jac_node.loc.col_start
+                )
                 i.jac_link: list[ast3.AST] = [jac_node]  # type: ignore
         return py_node
 
@@ -248,14 +257,17 @@ class PyastGenPass(Pass):
         doc: Optional[ast.String] = None,
     ) -> list[ast3.AST]:
         """Unwind codeblock."""
+        valid_stmts = (
+            [i for i in node.items if not isinstance(i, ast.Semi)] if node else []
+        )
         ret: list[ast3.AST] = (
             [self.sync(ast3.Pass(), node)]
-            if isinstance(node, ast.SubNodeList) and not node.items
+            if isinstance(node, ast.SubNodeList) and not valid_stmts
             else (
                 self.flatten(
                     [
                         x.gen.py_ast
-                        for x in node.items
+                        for x in valid_stmts
                         if not isinstance(x, ast.AstImplOnlyNode)
                     ]
                 )
@@ -926,21 +938,88 @@ class PyastGenPass(Pass):
         if isinstance(node.body, ast.AbilityDef):
             self.link_jac_py_nodes(jac_node=node.body, py_nodes=node.gen.py_ast)
 
+    def bfs_collect_type(self, node: ast.AstNode) -> list[str]:
+        """Collect type information in assignment using bfs."""
+        extracted_type = []
+        if isinstance(node, (ast.BuiltinType, ast.Token)):
+            extracted_type.append(node.value)
+        for child in node.kid:
+            extracted_type.extend(self.bfs_collect_type(child))
+        return extracted_type
+
     def gen_llm_body(self, node: ast.Ability) -> list[ast3.AST]:
         """Generate llm body."""
         if isinstance(node.body, ast.FuncCall):
-            model_params = (
-                {
-                    param.key: param.value
-                    for param in node.body.params.items
-                    if isinstance(param, ast.KWPair)
-                }
-                if node.body.params
-                else {}
+            model_params = {}
+            include_info = []
+            exclude_info = []
+            if node.body.params:
+                for param in node.body.params.items:
+                    if isinstance(param, ast.KWPair) and isinstance(
+                        param.key, ast.Name
+                    ):
+                        key = param.key.value
+                        value = param.value
+                        if key not in ["incl_info", "excl_info"]:
+                            model_params[key] = value
+                        elif key == "incl_info":
+                            if isinstance(value, ast.AtomUnit):
+                                var_name = (
+                                    value.value.right.value
+                                    if isinstance(value.value, ast.AtomTrailer)
+                                    and isinstance(value.value.right, ast.Name)
+                                    else (
+                                        value.value.value
+                                        if isinstance(value.value, ast.Name)
+                                        else ""
+                                    )
+                                )
+                                include_info.append((var_name, value.gen.py_ast[0]))
+                            elif isinstance(value, ast.TupleVal) and value.values:
+                                for i in value.values.items:
+                                    var_name = (
+                                        i.right.value
+                                        if isinstance(i, ast.AtomTrailer)
+                                        and isinstance(i.right, ast.Name)
+                                        else (
+                                            i.value if isinstance(i, ast.Name) else ""
+                                        )
+                                    )
+                                    include_info.append((var_name, i.gen.py_ast[0]))
+                        elif key == "excl_info":
+                            if isinstance(value, ast.AtomUnit):
+                                var_name = (
+                                    value.value.right.value
+                                    if isinstance(value.value, ast.AtomTrailer)
+                                    and isinstance(value.value.right, ast.Name)
+                                    else (
+                                        value.value.value
+                                        if isinstance(value.value, ast.Name)
+                                        else ""
+                                    )
+                                )
+                                exclude_info.append((var_name, value.gen.py_ast[0]))
+                            elif isinstance(value, ast.TupleVal) and value.values:
+                                for i in value.values.items:
+                                    var_name = (
+                                        i.right.value
+                                        if isinstance(i, ast.AtomTrailer)
+                                        and isinstance(i.right, ast.Name)
+                                        else (
+                                            i.value if isinstance(i, ast.Name) else ""
+                                        )
+                                    )
+                                    exclude_info.append((var_name, i.gen.py_ast[0]))
+            extracted = (
+                "".join(self.bfs_collect_type(node.signature.return_type))
+                if isinstance(node.signature, ast.FuncSignature)
+                and node.signature.return_type
+                else None
             )
             return [
                 self.sync(
-                    ast3.Return(
+                    ast3.Assign(
+                        targets=[self.sync(ast3.Name(id="output", ctx=ast3.Store()))],
                         value=self.sync(
                             ast3.Call(
                                 func=self.sync(
@@ -957,6 +1036,16 @@ class PyastGenPass(Pass):
                                 ),
                                 args=[],
                                 keywords=[
+                                    self.sync(
+                                        ast3.keyword(
+                                            arg="file_loc",
+                                            value=self.sync(
+                                                ast3.Name(
+                                                    id="__file__", ctx=ast3.Load()
+                                                )
+                                            ),
+                                        )
+                                    ),
                                     self.sync(
                                         ast3.keyword(
                                             arg="model",
@@ -969,7 +1058,9 @@ class PyastGenPass(Pass):
                                             value=self.sync(
                                                 ast3.Dict(
                                                     keys=[
-                                                        self.sync(ast3.Constant(value=key.value))  # type: ignore
+                                                        self.sync(
+                                                            ast3.Constant(value=key)
+                                                        )
                                                         for key in model_params.keys()
                                                     ],
                                                     values=[
@@ -980,21 +1071,69 @@ class PyastGenPass(Pass):
                                             ),
                                         )
                                     ),
+                                    self.sync(
+                                        ast3.keyword(
+                                            arg="scope",
+                                            value=(
+                                                self.sync(
+                                                    ast3.Constant(
+                                                        value=str(get_sem_scope(node))
+                                                    )
+                                                )
+                                            ),
+                                        )
+                                    ),
                                     self.sync(
                                         ast3.keyword(
                                             arg="incl_info",
                                             value=self.sync(
-                                                ast3.Constant(value=None)
-                                            ),  # TODO: Add incl_info
+                                                ast3.List(
+                                                    elts=[
+                                                        self.sync(
+                                                            ast3.Tuple(
+                                                                elts=[
+                                                                    self.sync(
+                                                                        ast3.Constant(
+                                                                            value=key
+                                                                        )
+                                                                    ),
+                                                                    value,
+                                                                ],
+                                                                ctx=ast3.Load(),
+                                                            )
+                                                        )
+                                                        for key, value in include_info
+                                                    ],
+                                                    ctx=ast3.Load(),
+                                                )
+                                            ),
                                         )
                                     ),
                                     self.sync(
                                         ast3.keyword(
                                             arg="excl_info",
                                             value=self.sync(
-                                                ast3.Constant(value=None)
-                                            ),  # TODO: Add excl_info
-                                        )
+                                                ast3.List(
+                                                    elts=[
+                                                        self.sync(
+                                                            ast3.Tuple(
+                                                                elts=[
+                                                                    self.sync(
+                                                                        ast3.Constant(
+                                                                            value=key
+                                                                        )
+                                                                    ),
+                                                                    value,
+                                                                ],
+                                                                ctx=ast3.Load(),
+                                                            )
+                                                        )
+                                                        for key, value in exclude_info
+                                                    ],
+                                                    ctx=ast3.Load(),
+                                                )
+                                            ),
+                                        ),
                                     ),
                                     self.sync(
                                         ast3.keyword(
@@ -1007,11 +1146,15 @@ class PyastGenPass(Pass):
                                                                 ast3.Tuple(
                                                                     elts=[
                                                                         (
-                                                                            param.semstr.gen.py_ast[
-                                                                                0
-                                                                            ]
-                                                                            if param.semstr
-                                                                            else None
+                                                                            self.sync(
+                                                                                ast3.Constant(
+                                                                                    value=(
+                                                                                        param.semstr.lit_value
+                                                                                        if param.semstr
+                                                                                        else None
+                                                                                    )
+                                                                                )
+                                                                            )
                                                                         ),
                                                                         (
                                                                             param.type_tag.tag.gen.py_ast[
@@ -1057,11 +1200,15 @@ class PyastGenPass(Pass):
                                                     elts=(
                                                         [
                                                             (
-                                                                node.signature.semstr.gen.py_ast[
-                                                                    0
-                                                                ]
-                                                                if node.signature.semstr
-                                                                else None
+                                                                self.sync(
+                                                                    ast3.Constant(
+                                                                        value=(
+                                                                            node.signature.semstr.lit_value
+                                                                            if node.signature.semstr
+                                                                            else None
+                                                                        )
+                                                                    )
+                                                                )
                                                             ),
                                                             (
                                                                 node.signature.return_type.gen.py_ast[
@@ -1070,6 +1217,15 @@ class PyastGenPass(Pass):
                                                                 if node.signature.return_type
                                                                 else None
                                                             ),
+                                                            (
+                                                                self.sync(
+                                                                    ast3.Constant(
+                                                                        value=(
+                                                                            extracted
+                                                                        )
+                                                                    )
+                                                                )
+                                                            ),
                                                         ]
                                                         if isinstance(
                                                             node.signature,
@@ -1081,7 +1237,7 @@ class PyastGenPass(Pass):
                                                 )
                                             ),
                                         )
-                                    ),  # TODO: Add Meaning Types of Outputs
+                                    ),
                                     self.sync(
                                         ast3.keyword(
                                             arg="action",
@@ -1094,9 +1250,55 @@ class PyastGenPass(Pass):
                                     ),
                                 ],
                             )
-                        )
+                        ),
                     )
-                )
+                ),
+                self.sync(
+                    ast3.Try(
+                        body=[
+                            self.sync(
+                                ast3.Return(
+                                    value=self.sync(
+                                        ast3.Call(
+                                            func=self.sync(
+                                                ast3.Name(id="eval", ctx=ast3.Load())
+                                            ),
+                                            args=[
+                                                self.sync(
+                                                    ast3.Name(
+                                                        id="output", ctx=ast3.Load()
+                                                    )
+                                                )
+                                            ],
+                                            keywords=[],
+                                        )
+                                    )
+                                )
+                            )
+                        ],
+                        handlers=[
+                            self.sync(
+                                ast3.ExceptHandler(
+                                    type=None,
+                                    name=None,
+                                    body=[
+                                        self.sync(
+                                            ast3.Return(
+                                                value=self.sync(
+                                                    ast3.Name(
+                                                        id="output", ctx=ast3.Load()
+                                                    )
+                                                )
+                                            )
+                                        )
+                                    ],
+                                )
+                            )
+                        ],
+                        orelse=[],
+                        finalbody=[],
+                    )
+                ),
             ]
         else:
             return []
@@ -1939,7 +2141,7 @@ class PyastGenPass(Pass):
                     )
                 )
                 if node.is_enum_stmt
-                else self.ice()
+                else None if node.type_tag else self.ice()
             )
         )
         if node.type_tag:
@@ -2269,7 +2471,10 @@ class PyastGenPass(Pass):
         elif node.op.name in [Tok.STAR_MUL]:
             node.gen.py_ast = [
                 self.sync(
-                    ast3.Starred(value=node.operand.gen.py_ast[0], ctx=ast3.Load())
+                    ast3.Starred(
+                        value=node.operand.gen.py_ast[0],
+                        ctx=ast3.Load(),
+                    )
                 )
             ]
         elif node.op.name in [Tok.STAR_POW]:

jaclang/compiler/passes/main/registry_pass.py

@@ -0,0 +1,126 @@
+"""Jac Registry Pass.
+
+This pass is responsible for creating object containing the senstring, scope,
+type of different important nodes in the AST as we loose access to the
+semstrings after PyASTGen pass. So we create those as a pickled file for
+each module
+"""
+
+import os
+import pickle
+
+import jaclang.compiler.absyntree as ast
+from jaclang.compiler.constant import Constants as Con
+from jaclang.compiler.passes import Pass
+from jaclang.core.registry import SemInfo, SemRegistry
+from jaclang.core.utils import get_sem_scope
+
+
+class RegistryPass(Pass):
+    """Creates a registry for each module."""
+
+    modules_visited: list[ast.Module] = []
+
+    def enter_module(self, node: ast.Module) -> None:
+        """Create registry for each module."""
+        node.registry = SemRegistry()
+        self.modules_visited.append(node)
+
+    def exit_module(self, node: ast.Module) -> None:
+        """Save registry for each module."""
+        module_dir = os.path.join(
+            os.path.abspath(os.path.dirname(node.source.file_path)), Con.JAC_GEN_DIR
+        )
+        module_name = node.name
+        os.makedirs(module_dir, exist_ok=True)
+        with open(os.path.join(module_dir, f"{module_name}.registry.pkl"), "wb") as f:
+            pickle.dump(node.registry, f)
+        self.modules_visited.pop()
+
+    def exit_architype(self, node: ast.Architype) -> None:
+        """Save architype information."""
+        scope = get_sem_scope(node)
+        seminfo = SemInfo(
+            node.name.value,
+            node.arch_type.value,
+            node.semstr.lit_value if node.semstr else None,
+        )
+        if (
+            len(self.modules_visited)
+            and self.modules_visited[-1].registry
+            and scope.parent
+        ):
+            self.modules_visited[-1].registry.add(scope.parent, seminfo)
+
+    def exit_enum(self, node: ast.Enum) -> None:
+        """Save enum information."""
+        scope = get_sem_scope(node)
+        seminfo = SemInfo(
+            node.name.value, "Enum", node.semstr.lit_value if node.semstr else None
+        )
+        if (
+            len(self.modules_visited)
+            and self.modules_visited[-1].registry
+            and scope.parent
+        ):
+            self.modules_visited[-1].registry.add(scope.parent, seminfo)
+
+    def exit_has_var(self, node: ast.HasVar) -> None:
+        """Save variable information."""
+        extracted_type = (
+            "".join(self.extract_type(node.type_tag.kid[1:][0]))
+            if node.type_tag
+            else None
+        )
+        scope = get_sem_scope(node)
+        seminfo = SemInfo(
+            node.name.value,
+            extracted_type,
+            node.semstr.lit_value if node.semstr else None,
+        )
+        if len(self.modules_visited) and self.modules_visited[-1].registry:
+            self.modules_visited[-1].registry.add(scope, seminfo)
+
+    def exit_assignment(self, node: ast.Assignment) -> None:
+        """Save assignment information."""
+        if node.aug_op:
+            return
+
+        extracted_type = (
+            "".join(self.extract_type(node.type_tag.kid[1:][0]))
+            if node.type_tag
+            else None
+        )
+        scope = get_sem_scope(node)
+        seminfo = SemInfo(
+            (
+                node.target.items[0].value
+                if isinstance(node.target.items[0], ast.Name)
+                else ""
+            ),
+            extracted_type,
+            node.semstr.lit_value if node.semstr else None,
+        )
+        if len(self.modules_visited) and self.modules_visited[-1].registry:
+            self.modules_visited[-1].registry.add(scope, seminfo)
+
+    def exit_name(self, node: ast.Name) -> None:
+        """Save name information. for enum stmts."""
+        if (
+            node.parent
+            and node.parent.parent
+            and node.parent.parent.__class__.__name__ == "Enum"
+        ):
+            scope = get_sem_scope(node)
+            seminfo = SemInfo(node.value, None, None)
+            if len(self.modules_visited) and self.modules_visited[-1].registry:
+                self.modules_visited[-1].registry.add(scope, seminfo)
+
+    def extract_type(self, node: ast.AstNode) -> list[str]:
+        """Collect type information in assignment using bfs."""
+        extracted_type = []
+        if isinstance(node, (ast.BuiltinType, ast.Token)):
+            extracted_type.append(node.value)
+        for child in node.kid:
+            extracted_type.extend(self.extract_type(child))
+        return extracted_type

jaclang/compiler/passes/main/schedules.py

@@ -16,6 +16,7 @@ from .pybc_gen_pass import PyBytecodeGenPass  # noqa: I100
 from .pyast_gen_pass import PyastGenPass  # noqa: I100
 from .type_check_pass import JacTypeCheckPass  # noqa: I100
 from .fuse_typeinfo_pass import FuseTypeInfoPass  # noqa: I100
+from .registry_pass import RegistryPass  # noqa: I100
 
 py_code_gen = [
     SubNodeTabPass,
@@ -23,6 +24,7 @@ py_code_gen = [
     SymTabBuildPass,
     DeclDefMatchPass,
     DefUsePass,
+    RegistryPass,
     PyastGenPass,
     PyBytecodeGenPass,
 ]

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

@@ -0,0 +1,39 @@
+"""Test registry pass."""
+
+import os
+
+from jaclang.compiler.compile import jac_file_to_pass
+from jaclang.compiler.passes.main import RegistryPass
+from jaclang.utils.test import TestCase
+
+
+class RegistryPassTests(TestCase):
+    """Test pass module."""
+
+    def setUp(self) -> None:
+        """Set up test."""
+        return super().setUp()
+
+    def test_registry_pass(self) -> None:
+        """Basic test for pass."""
+        state = jac_file_to_pass(self.fixture_abs_path("registry.jac"), RegistryPass)
+        self.assertFalse(state.errors_had)
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(
+                    os.path.dirname(self.fixture_abs_path("registry.jac")),
+                    "__jac_gen__",
+                    "registry.registry.pkl",
+                )
+            )
+        )
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(
+                    os.path.dirname(self.fixture_abs_path("registry.jac")),
+                    "__jac_gen__",
+                    "fstrings.registry.pkl",
+                )
+            )
+        )
+        self.assertIn("109", str(state.ir.to_dict()))