Best Guess Sourcemapping when Partially Enabled (#688)

pyteal/compiler/sourcemap.py

@@ -461,6 +461,7 @@ def __init__(
         teal_line: str,
         teal_component: "pt.TealComponent",
         pcs: list[int] | None = None,
+        is_sentinel: bool = False,
     ):
         super().__init__(
             frame_info=pt_frame.frame_info,
@@ -474,6 +475,7 @@ def __init__(
         self.teal_line: Final[str] = teal_line
         self.teal_component: Final[pt.TealComponent] = teal_component
         self.pcs: Final[list[int] | None] = pcs if pcs else None
+        self.is_sentinel: Final[bool] = is_sentinel
 
     def pcs_repr(self, prefix: str = "") -> str:
         if not self.pcs:
@@ -482,7 +484,7 @@ def pcs_repr(self, prefix: str = "") -> str:
 
     def __repr__(self) -> str:
         P = " // "
-        return f"TealLine({self.teal_lineno}: {self.teal_line}{self.pcs_repr(prefix=P)} // PyTeal: {self._hybrid_w_offset()[0]}"
+        return f"TealLine({self.teal_lineno}: {self.teal_line}{self.pcs_repr(prefix=P)} // PyTeal: {self._hybrid_w_offset()[0]})"
 
     def teal_column(self) -> int:
         """Always returns 0 as the 0-index STARTING column offset"""
@@ -700,7 +702,7 @@ def __init__(
         self.teal_filename: str | None = teal_filename
         self.verbose: bool = verbose
 
-        self._best_frames: list[PyTealFrame] = []
+        self._best_frames: list[PyTealFrame | None] = []
         self._cached_r3sourcemap: R3SourceMap | None = None
 
         self._cached_tmis: list[TealMapItem] = []
@@ -793,14 +795,20 @@ def build(self) -> None:
         #   TO:   a PyTealFrame
         # overall resulting in a list[PyTealFrame]
         self._best_frames = [
-            tc.stack_frames().best().as_pyteal_frame() for tc in self.components
+            tc.stack_frames()._best_frame_as_pyteal_frame() for tc in self.components
         ]
 
         if not self._best_frames:
             raise self._unexpected_error(
                 f"This shouldn't have happened as we already checked! Check again: {len(self.components)=}"
             )
 
+        # stand-in in the case of missing frames
+        sentinel_frame = self._best_frames[0]
+        assert (
+            sentinel_frame
+        ), "Abort source mapping as even the very first best frame is missing"
+
         # PASS II. Attempt to fill any "gaps" by inferring from adjacent BFC's
         self._best_frames, inferred = self._infer(self._best_frames)
         if inferred:
@@ -816,11 +824,12 @@ def build(self) -> None:
                     pcs = pcsm.line_to_pc.get(lineno - 1, [])
                 self._cached_tmis.append(
                     TealMapItem(
-                        pt_frame=best_frame,
+                        pt_frame=best_frame or sentinel_frame,
                         teal_lineno=lineno,
                         teal_line=line,  # type: ignore
                         teal_component=self.components[i],
                         pcs=pcs,
+                        is_sentinel=(best_frame is None),
                     )
                 )
                 lineno += 1
@@ -857,6 +866,16 @@ def _validate_build(self):
                     f"TealMapItem's don't match R3SourceMap.file_lines at index {i}. ('{tmi_line}' v. '{target_line}')"
                 )
 
+        for tmi in self._cached_tmis:
+            if not tmi.is_sentinel:
+                continue
+            print(
+                f"""-----------------
+WARNING: Source mapping is unknown for the following:
+{tmi!r}
+"""
+            )
+
     def _build_r3sourcemap(self):
         assert self._cached_tmis, "Unexpected error: no cached TealMapItems found"
 
@@ -919,8 +938,8 @@ def as_r3sourcemap(self) -> R3SourceMap | None:
 
     @classmethod
     def _infer(
-        cls, best_frames: list[PyTealFrame]
-    ) -> tuple[list[PyTealFrame], list[int]]:
+        cls, best_frames: list[PyTealFrame | None]
+    ) -> tuple[list[PyTealFrame | None], list[int]]:
         inferred = []
         frames = list(best_frames)
         N = len(frames)

pyteal/stack_frame.py

@@ -12,6 +12,14 @@
 from executing import Source  # type: ignore
 
 
+class SourceMapStackFramesError(RuntimeError):
+    def __init__(self, msg: str):
+        self.msg = msg
+
+    def __str__(self):
+        return self.msg
+
+
 @dataclass(frozen=True)
 class StackFrame:
     """
@@ -96,11 +104,7 @@ def _init_or_drop(
     ]
     _internal_paths_re = re.compile("|".join(_internal_paths))
 
-    def as_pyteal_frame(
-        self,
-        rel_paths: bool = True,
-        parent: "PyTealFrame | None" = None,
-    ) -> "PyTealFrame":
+    def as_pyteal_frame(self) -> "PyTealFrame":
         """
         Downcast one level in the class hierarchy
         """
@@ -109,8 +113,8 @@ def as_pyteal_frame(
             node=self.node,
             creator=self.creator,
             full_stack=self.full_stack,
-            rel_paths=rel_paths,
-            parent=parent,
+            rel_paths=True,
+            parent=None,
         )
 
     @classmethod
@@ -284,16 +288,35 @@ def _make_stack_frames(fis):
     def __len__(self) -> int:
         return len(self._frames)
 
-    def best(self) -> StackFrame:
+    def _best_frame(self) -> StackFrame:
         """
-        Return the best guess as to the user-authored birthplace of the
-        associated StackFrame's
+        Return the best guess as to the user-authored birthplace of the associated StackFrame's.
+        If self._frames is non-empty return the last frame in the stack trace.
+        Otherwise, raise a SourceMapStackFramesError.
         """
-        assert (
-            self._frames
-        ), f"expected to have some frames but currently {self._frames=}"
+        if not self._frames:
+            raise SourceMapStackFramesError(
+                f"expected to have some frames but currently {self._frames=}"
+            )
+
         return self._frames[-1]
 
+    def _best_frame_as_pyteal_frame(self) -> "PyTealFrame | None":
+        """
+        Return the best frame converted to a PyTealFrame, or None if there was an error.
+        """
+        try:
+            return self._best_frame().as_pyteal_frame()
+        except SourceMapStackFramesError as smsfe:
+            print(
+                f"""-------------------
+WARNING: Error retrieving the best frame for source mapping.
+This may occur because FeatureGates was imported and `FeatureGates.set_sourcemap_enabled(True)` called _AFTER_ pyteal was imported.
+error: {smsfe}
+"""
+            )
+        return None
+
     def __repr__(self) -> str:
         return f"{'C' if self._compiler_gen else 'U'}{self._frames}"
 
@@ -347,9 +370,11 @@ def _debug_asts(cls, *exprs: "Expr") -> None:  # type: ignore
             return
 
         def dbg(e: Expr):
-            print(
-                type(e), ": ", e.stack_frames.best().as_pyteal_frame().hybrid_unparsed()
-            )
+            try:
+                finfo = e.stack_frames._best_frame().as_pyteal_frame().hybrid_unparsed()
+            except SourceMapStackFramesError as smsfe:
+                finfo = str(smsfe)
+            print(type(e), ": ", finfo)
 
         cls._walk_asts(dbg, *exprs)
 

tests/unit/sourcemap_test.py

@@ -378,6 +378,49 @@ def mock_R3SourceMap(lines):
     ptsm._validate_annotated(omit_headers, teal, annotated_w_headers)
 
 
+def test_examples_sourcemap():
+    """
+    Test to ensure that examples/application/teal/sourcemap.py doesn't go stale
+    """
+    from examples.application.sourcemap import Compilation, Mode, program
+
+    examples = Path.cwd() / "examples" / "application" / "teal"
+
+    approval_program = program()
+
+    results = Compilation(approval_program, mode=Mode.Application, version=8).compile(
+        with_sourcemap=True, annotate_teal=True, annotate_teal_headers=True
+    )
+
+    teal = examples / "sourcemap.teal"
+    annotated = examples / "sourcemap_annotated.teal"
+
+    with open(teal) as f:
+        assert f.read() == results.teal
+
+    with open(annotated) as f:
+        fixture = f.read().splitlines()
+        annotated = results.sourcemap.annotated_teal.splitlines()
+        for i, (f, a) in enumerate(zip(fixture, annotated)):
+            f_cols = f.split()
+            a_cols = a.split()
+            if f_cols == a_cols:
+                continue
+
+            if f_cols[-1] == "annotate_teal_headers=True)":
+                assert f_cols[:2] == a_cols[:2], f"index {i} doesn't match"
+                assert f_cols[-4:] == a_cols[-4:], f"index {i} doesn't match"
+                continue
+
+            # must differ because fixture repeats PYTEAL PATH so omits it
+            assert len(f_cols) == len(a_cols) - 1, f"index {i} doesn't match"
+
+            a_comment = a_cols.index("//")
+            assert f_cols == (
+                a_cols[: a_comment + 1] + a_cols[a_comment + 2 :]
+            ), f"index {i} doesn't match"
+
+
 @pytest.mark.skip(
     reason="""Supressing this flaky test as 
 router_test::test_router_compile_program_idempotence is similar in its goals