Merge pull request #237 from OpShin/fix/missing_locations

Fix/missing locations

opshin/compiler.py

@@ -1007,6 +1007,35 @@ def generic_visit(self, node: AST) -> plt.AST:
         raise NotImplementedError(f"Can not compile {node}")
 
 
+def custom_fix_missing_locations(node):
+    """
+    Works like ast.fix_missing_location but forces it onto everything
+    """
+
+    def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
+        if not hasattr(node, "lineno"):
+            node.lineno = lineno
+        else:
+            lineno = node.lineno
+        if getattr(node, "end_lineno", None) is None:
+            node.end_lineno = end_lineno
+        else:
+            end_lineno = node.end_lineno
+        if not hasattr(node, "col_offset"):
+            node.col_offset = col_offset
+        else:
+            col_offset = node.col_offset
+        if getattr(node, "end_col_offset", None) is None:
+            node.end_col_offset = end_col_offset
+        else:
+            end_col_offset = node.end_col_offset
+        for child in iter_child_nodes(node):
+            _fix(child, lineno, col_offset, end_lineno, end_col_offset)
+
+    _fix(node, 1, 0, 1, 0)
+    return node
+
+
 def compile(
     prog: AST,
     filename=None,
@@ -1015,7 +1044,7 @@ def compile(
     constant_folding=False,
     validator_function_name="validator",
 ):
-    rewrite_steps = [
+    compile_pipeline = [
         # Important to call this one first - it imports all further files
         RewriteImport(filename=filename),
         # Rewrites that simplify the python code
@@ -1037,13 +1066,6 @@ def compile(
         RewriteImportUPLCBuiltins(),
         RewriteInjectBuiltinsConstr(),
         RewriteRemoveTypeStuff(),
-    ]
-    for s in rewrite_steps:
-        prog = s.visit(prog)
-        prog = fix_missing_locations(prog)
-
-    # from here on raw uplc may occur, so we dont attempt to fix locations
-    compile_pipeline = [
         # Save the original names of variables
         RewriteOrigName(),
         RewriteScoping(),
@@ -1052,13 +1074,16 @@ def compile(
         OptimizeVarlen(),
         OptimizeRemoveDeadconstants(),
         OptimizeRemovePass(),
-        # the compiler runs last
-        UPLCCompiler(
-            force_three_params=force_three_params,
-            validator_function_name=validator_function_name,
-        ),
     ]
     for s in compile_pipeline:
         prog = s.visit(prog)
+        prog = custom_fix_missing_locations(prog)
+
+    # the compiler runs last
+    s = UPLCCompiler(
+        force_three_params=force_three_params,
+        validator_function_name=validator_function_name,
+    )
+    prog = s.visit(prog)
 
     return prog

opshin/typed_ast.py

@@ -178,3 +178,6 @@ class TypedAssert(typedstmt, Assert):
 class RawPlutoExpr(typedexpr):
     typ: Type
     expr: plt.AST
+
+    _attributes = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
+    _fields = []