From f05aceac2d3ad1aca6a491cd80ce80068bf14e92 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=20L=C3=A4ufer?= <laeufer@cornell.edu>
Date: Wed, 18 Sep 2024 14:24:16 -0400
Subject: [PATCH] calyx-py: add line numbers to group enables, similar to
 Dhalia

---
 calyx-py/calyx/py_ast.py | 46 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/calyx-py/calyx/py_ast.py b/calyx-py/calyx/py_ast.py
index 6f0d661d29..7655c74b22 100644
--- a/calyx-py/calyx/py_ast.py
+++ b/calyx-py/calyx/py_ast.py
@@ -2,6 +2,8 @@
 from dataclasses import dataclass, field
 from typing import List, Any, Tuple, Optional
 from calyx.utils import block
+import inspect
+import os
 
 
 @dataclass
@@ -461,6 +463,44 @@ def doc(self) -> str:
         return f"({self.left.doc()} >= {self.right.doc()})"
 
 
+@dataclass
+class SourceLoc:
+    line: int
+    file: str
+
+
+def frame_to_source_loc(frame: inspect.FrameInfo) -> SourceLoc:
+    """builds a SourceLoc object from a Python frame"""
+    return SourceLoc(frame.lineno, os.path.basename(frame.filename))
+
+
+def determine_source_loc() -> Optional[SourceLoc]:
+    """Inspects the call stack to determine the first call site outside the calyx-py library."""
+    stacktrace = inspect.stack()
+
+    # inspect top frame to determine the path to the calyx-py library
+    top = stacktrace[0]
+    assert top.function == "determine_source_loc"
+    library_path = os.path.dirname(top.filename)
+    assert os.path.join(library_path, "py_ast.py") == top.filename
+
+    # find first stack frame that is not part of the library
+    user = None
+    for frame in stacktrace:
+        print(frame.filename, frame.function, frame)
+        # skip frames that do not have a real filename
+        if frame.filename == "<string>":
+            continue
+        if not frame.filename.startswith(library_path):
+            user = frame
+            break
+    if user is None:
+        return None
+
+    # build source locator from frame
+    return frame_to_source_loc(user)
+
+
 # Control
 @dataclass
 class Control(Emittable):
@@ -470,9 +510,13 @@ class Control(Emittable):
 @dataclass
 class Enable(Control):
     stmt: str
+    loc: Optional[SourceLoc] = field(default_factory=determine_source_loc)
 
     def doc(self) -> str:
-        return f"{self.stmt};"
+        if self.loc is None:
+            return f"{self.stmt};"
+        else:
+            return f"@pos({self.loc.line}) {self.stmt};"
 
 
 @dataclass