From a01594e3b691faa57498ad21d9924a3065a4bd51 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Niels=20M=C3=BCndler?= <n.muendler@posteo.de>
Date: Sat, 12 Aug 2023 00:58:38 +0200
Subject: [PATCH 1/3] Support lineno for RawPlutoExpressions

---
 opshin/compiler.py  | 22 +++++++++-------------
 opshin/typed_ast.py |  3 +++
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/opshin/compiler.py b/opshin/compiler.py
index edbe989c..cc8c3a94 100644
--- a/opshin/compiler.py
+++ b/opshin/compiler.py
@@ -1035,13 +1035,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(),
@@ -1050,13 +1043,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:
+    for s in rewrite_steps:
         prog = s.visit(prog)
+        prog = 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
diff --git a/opshin/typed_ast.py b/opshin/typed_ast.py
index aaf75a53..d87022ec 100644
--- a/opshin/typed_ast.py
+++ b/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 = []

From c2f16e27f0e1ed42e0204d7b432ac6b931c00b52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Niels=20M=C3=BCndler?= <n.muendler@posteo.de>
Date: Sat, 12 Aug 2023 01:02:00 +0200
Subject: [PATCH 2/3] Force the fixed locations onto everything

---
 opshin/compiler.py | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/opshin/compiler.py b/opshin/compiler.py
index cc8c3a94..53cc15c7 100644
--- a/opshin/compiler.py
+++ b/opshin/compiler.py
@@ -1006,6 +1006,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,
@@ -1046,7 +1075,7 @@ def compile(
     ]
     for s in rewrite_steps:
         prog = s.visit(prog)
-        prog = fix_missing_locations(prog)
+        prog = custom_fix_missing_locations(prog)
 
     # the compiler runs last
     s = UPLCCompiler(

From 602f8b846f819e26d68afc84ba2648d9a3500ef0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Niels=20M=C3=BCndler?= <n.muendler@posteo.de>
Date: Sat, 12 Aug 2023 01:03:28 +0200
Subject: [PATCH 3/3] Rename for clarity

---
 opshin/compiler.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/opshin/compiler.py b/opshin/compiler.py
index 53cc15c7..4df8bd55 100644
--- a/opshin/compiler.py
+++ b/opshin/compiler.py
@@ -1043,7 +1043,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
@@ -1073,7 +1073,7 @@ def compile(
         OptimizeRemoveDeadconstants(),
         OptimizeRemovePass(),
     ]
-    for s in rewrite_steps:
+    for s in compile_pipeline:
         prog = s.visit(prog)
         prog = custom_fix_missing_locations(prog)