Teal to PyTeal Source Mapper (#650)

.flake8

@@ -23,8 +23,14 @@ per-file-ignores =
     examples/signature/recurring_swap.py:           F403, F405
     examples/signature/split.py:                    F403, F405
     pyteal/__init__.py:                             F401, F403
+    pyteal/compiler/flatten.py:                     F821
     pyteal/compiler/optimizer/__init__.py:          F401
     pyteal/ir/ops.py:                               E221
+    pyteal/ir/tealconditionalblock.py:              F821
+    pyteal/ir/tealsimpleblock.py:                   F821
+    pyteal/stack_frame.py:                          F821
+    tests/teal/rps_helpers/program.py:              F401, F403, F405
+    tests/teal/rps.py:                              F403, F405, I252
     tests/unit/module_test.py:                      F401, F403
 
 # from flake8-tidy-imports

.github/workflows/build.yml

@@ -1,4 +1,4 @@
-name: "Build workflow"
+name: "On Commit Workflow"
 on:
   pull_request:
   push:
@@ -49,7 +49,7 @@ jobs:
     runs-on: ubuntu-20.04
     strategy:
       matrix:
-        python: [ "3.10" ]
+        python: ["3.10"]
     steps:
       - name: Check out code
         uses: actions/checkout@v2

.github/workflows/nightly.yml

@@ -1,7 +1,7 @@
 name: "Nightly Build"
 on:
   schedule:
-    - cron: "10 17 * * *"
+    - cron: "10 1 * * *"
 
 jobs:
   nightly-slow-unit-tests:

CHANGELOG.md

@@ -4,7 +4,8 @@
 
 * Clarify that `Approve` and `Reject` always exit in the documentation. ([#660](https://github.com/algorand/pyteal/pull/660))
 * Added frame pointer support for router. ([#600](https://github.com/algorand/pyteal/pull/600))
- * NOTE: a backwards incompatable change was imposed in this PR: previous `build_program` method in `Router` was exported and public, now this method is hidden. Use `compile_program` only.
+  * NOTE: a backwards incompatable change was imposed in this PR: previous `build_program` method in `Router` was exported and public, now this method is hidden. Use `compile_program` only.
+* Experimental source mapping capability. ([#650](https://github.com/algorand/pyteal/pull/650))
 
 ## Fixed
 

Makefile

@@ -21,6 +21,8 @@ bdist-wheel:
 
 bundle-docs-clean:
 	rm -rf docs/pyteal.docset
+	rm -rf docs/_build
+	rm -f docs/pyteal.docset.tar.gz
 
 bundle-docs: bundle-docs-clean
 	cd docs && \
@@ -51,14 +53,19 @@ lint: black flake8 mypy sdist-check
 
 # ---- Unit Tests (no algod) ---- #
 
-# TODO: add blackbox_test.py to multithreaded tests when following issue has been fixed https://github.com/algorand/pyteal/issues/199
+# Slow test which are fast enough on python 3.11+
+test-unit-slow: 
+	pytest tests/unit/sourcemap_constructs311_test.py -m serial
+
+test-unit-very-slow: 
+	pytest tests/unit/sourcemap_constructs_allpy_test.py -m serial
 
-NUM_PROCS = auto
 test-unit-async:
-	pytest -n $(NUM_PROCS) --durations=10 -sv pyteal tests/unit -m "not serial"
+	pytest -n auto --durations=10 pyteal tests/unit -m "not slow" -m "not serial"
 
-test-unit-sync:
-	pytest --durations=10 -sv pyteal tests/unit -m serial
+# Run tests w/ @pytest.mark.serial under ~/tests/unit each in its own proc:
+test-unit-sync: test-unit-slow
+	find tests/unit -name '*_test.py' | sort | xargs -t -I {} pytest --suppress-no-test-exit-code --dist=no --durations=10 {} -m serial -m "not slow"
 
 test-unit: test-unit-async test-unit-sync
 
@@ -76,13 +83,19 @@ algod-start-report: algod-start algod-version
 
 algod-stop:
 	docker compose stop algod
-integration-run:
-	pytest -n $(NUM_PROCS) --durations=10 -sv tests/integration -m "not serial"
-	pytest --durations=10 -sv tests/integration -m serial
 
-test-integration: integration-run
+test-integ-async:
+	pytest -n auto --durations=10 -sv tests/integration -m "not serial"
+
+# Run tests w/ @pytest.mark.serial under ~/tests/integration each in its own proc:
+test-integ-sync:
+	find tests/integration -name '*_test.py' | sort | xargs -t -I {} pytest --suppress-no-test-exit-code --dist=no --durations=10 {} -m serial
+
+test-integration: test-integ-async test-integ-sync
 
-all-tests: lint-and-test test-integration
+all-sync: test-unit-sync test-integ-sync
+
+all-lint-unit-integ: lint-and-test test-integration
 
 # ---- Github Actions 1-Liners ---- #
 
@@ -94,8 +107,7 @@ check-code-changes:
 	git config --global --add safe.directory /__w/pyteal/pyteal
 	[ -n "$$(git log --since='24 hours ago')" ] && (echo "should_run=true" >> $(GITHUB_ENV)) || (echo "should_run=false" >> $(GITHUB_ENV))
 
-nightly-slow:
-	echo "TODO - this is a stub for a very slow test"
+nightly-slow: test-unit-very-slow
 
 # ---- Local Github Actions Simulation via `act` ---- #
 # assumes act is installed, e.g. via `brew install act`
@@ -111,3 +123,6 @@ local-gh-simulate:
 
 coverage:
 	pytest --cov-report html --cov=pyteal
+
+sourcemap-coverage:
+	pytest --cov-report html --cov=pyteal.stack_frame --cov=pyteal.compiler.sourcemap --cov=pyteal.compiler.compiler --dist=no tests/unit/sourcemap_monkey_unit_test.py -m serial

docs/requirements.txt

@@ -1,6 +1,8 @@
 sphinx==5.1.1
 sphinx-rtd-theme==1.0.0
 # dependencies from setup.py
+docstring-parser==0.14.1
+executing==1.2.0
 py-algorand-sdk>=2.0.0,<3.0.0
 semantic-version>=2.9.0,<3.0.0
-docstring-parser==0.14.1
+tabulate>=0.9.0,<0.10.0

examples/application/abi/algobank.py

@@ -112,11 +112,11 @@ def withdraw(amount: abi.Uint64, recipient: abi.Account) -> Expr:
     )
 
 
-approval_program, clear_state_program, contract = router.compile_program(
-    version=6, optimize=OptimizeOptions(scratch_slots=True)
-)
-
 if __name__ == "__main__":
+    approval_program, clear_state_program, contract = router.compile_program(
+        version=6, optimize=OptimizeOptions(scratch_slots=True)
+    )
+
     with open("algobank_approval.teal", "w") as f:
         f.write(approval_program)
 

pyteal.ini

@@ -0,0 +1,5 @@
+[pyteal]
+
+[pyteal-source-mapper]
+enabled = False
+debug = False

pyteal/__init__.py

@@ -11,11 +11,15 @@
     MIN_PROGRAM_VERSION,
     DEFAULT_PROGRAM_VERSION,
     CompileOptions,
+    Compilation,
     compileTeal,
     OptimizeOptions,
+    PyTealSourceMap,
 )
 from pyteal.types import TealType
 from pyteal.errors import (
+    AlgodClientError,
+    SourceMapDisabledError,
     TealInternalError,
     TealTypeError,
     TealSeqError,
@@ -43,9 +47,13 @@
         "DEFAULT_PROGRAM_VERSION",
         "CompileOptions",
         "pragma",
+        "Compilation",
         "compileTeal",
         "OptimizeOptions",
+        "PyTealSourceMap",
         "TealType",
+        "AlgodClientError",
+        "SourceMapDisabledError",
         "TealInternalError",
         "TealTypeError",
         "TealSeqError",

pyteal/__init__.pyi

@@ -14,11 +14,15 @@ from pyteal.compiler import (
     MIN_PROGRAM_VERSION,
     DEFAULT_PROGRAM_VERSION,
     CompileOptions,
+    Compilation,
     compileTeal,
     OptimizeOptions,
+    PyTealSourceMap,
 )
 from pyteal.types import TealType
 from pyteal.errors import (
+    AlgodClientError,
+    SourceMapDisabledError,
     TealInternalError,
     TealTypeError,
     TealSeqError,
@@ -39,6 +43,7 @@ __all__ = [
     "AccountParamObject",
     "Add",
     "Addr",
+    "AlgodClientError",
     "And",
     "App",
     "AppField",
@@ -91,6 +96,7 @@ __all__ = [
     "BytesZero",
     "CallConfig",
     "Comment",
+    "Compilation",
     "CompileOptions",
     "Concat",
     "Cond",
@@ -175,6 +181,7 @@ __all__ = [
     "Or",
     "Pop",
     "Pragma",
+    "PyTealSourceMap",
     "RETURN_HASH_PREFIX",
     "Reject",
     "Replace",
@@ -194,6 +201,7 @@ __all__ = [
     "Sha512_256",
     "ShiftLeft",
     "ShiftRight",
+    "SourceMapDisabledError",
     "Sqrt",
     "Subroutine",
     "SubroutineCall",
@@ -210,6 +218,7 @@ __all__ = [
     "TealInternalError",
     "TealLabel",
     "TealOp",
+    "TealPragma",
     "TealPragmaError",
     "TealSeqError",
     "TealSimpleBlock",

pyteal/ast/abi/method_return.py

@@ -6,8 +6,9 @@
 from pyteal.ast.unaryexpr import Log
 from pyteal.ast.naryexpr import Concat
 from pyteal.ast.bytes import Bytes
-from pyteal.ir import TealBlock, TealSimpleBlock, Op
 from pyteal.config import RETURN_HASH_PREFIX
+from pyteal.ir import TealBlock, TealSimpleBlock, Op
+from pyteal.stack_frame import NatalStackFrame
 
 if TYPE_CHECKING:
     from pyteal.compiler import CompileOptions
@@ -19,15 +20,18 @@ def __init__(self, arg: BaseType):
         if not isinstance(arg, BaseType):
             raise TealInputError(f"Expecting an ABI type argument but get {arg}")
         self.arg = arg
+        self._sframes_container: Expr | None = None
 
     def __teal__(self, options: "CompileOptions") -> Tuple[TealBlock, TealSimpleBlock]:
         if options.version < Op.log.min_version:
             raise TealInputError(
                 f"current version {options.version} is lower than log's min version {Op.log.min_version}"
             )
-        return Log(Concat(Bytes(RETURN_HASH_PREFIX), self.arg.encode())).__teal__(
+        start, end = Log(Concat(Bytes(RETURN_HASH_PREFIX), self.arg.encode())).__teal__(
             options
         )
+        NatalStackFrame.reframe_ops_in_blocks(self._sframes_container or self, start)
+        return start, end
 
     def __str__(self) -> str:
         return f"(MethodReturn {self.arg.type_spec()})"

pyteal/ast/app.py

@@ -25,12 +25,12 @@
 class OnComplete:
     """An enum of values that :any:`TxnObject.on_completion()` may return."""
 
-    NoOp = EnumInt("NoOp")
-    OptIn = EnumInt("OptIn")
-    CloseOut = EnumInt("CloseOut")
-    ClearState = EnumInt("ClearState")
-    UpdateApplication = EnumInt("UpdateApplication")
-    DeleteApplication = EnumInt("DeleteApplication")
+    NoOp = EnumInt("NoOp")  # T2PT8
+    OptIn = EnumInt("OptIn")  # T2PT8
+    CloseOut = EnumInt("CloseOut")  # T2PT8
+    ClearState = EnumInt("ClearState")  # T2PT8
+    UpdateApplication = EnumInt("UpdateApplication")  # T2PT8
+    DeleteApplication = EnumInt("DeleteApplication")  # T2PT8
 
 
 OnComplete.__module__ = "pyteal"

pyteal/ast/assert_.py

@@ -30,13 +30,15 @@ def __init__(
 
         self.comment = comment
         self.cond = [cond] + list(additional_conds)
+        self._sframes_container: Expr | None = None
 
     def __teal__(self, options: "CompileOptions"):
         if len(self.cond) > 1:
             asserts: list[Expr] = []
             for cond in self.cond:
                 asrt = Assert(cond, comment=self.comment)
                 asrt.trace = cond.trace
+                asrt._sframes_container = cond
                 asserts.append(asrt)
             return Seq(*asserts).__teal__(options)
 

pyteal/ast/cond.py

@@ -77,12 +77,12 @@ def __init__(self, *argv: List[Expr]):
     def __teal__(self, options: "CompileOptions"):
         start = None
         end = TealSimpleBlock([])
-        prevBranch = None
+        prevBranch: TealConditionalBlock | None = None
         for i, (cond, pred) in enumerate(self.args):
             condStart, condEnd = cond.__teal__(options)
             predStart, predEnd = pred.__teal__(options)
 
-            branchBlock = TealConditionalBlock([])
+            branchBlock = TealConditionalBlock([], root_expr=cond)
             branchBlock.setTrueBlock(predStart)
 
             condEnd.setNextBlock(branchBlock)
@@ -91,6 +91,7 @@ def __teal__(self, options: "CompileOptions"):
                 start = condStart
             else:
                 cast(TealConditionalBlock, prevBranch).setFalseBlock(condStart)
+
             prevBranch = branchBlock
 
         errBlock = TealSimpleBlock([TealOp(self, Op.err)])

pyteal/ast/expr.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
-from pyteal.types import TealType
 from pyteal.ir import TealBlock, TealSimpleBlock
+from pyteal.stack_frame import NatalStackFrame
+from pyteal.types import TealType
 
 if TYPE_CHECKING:
     from pyteal.compiler import CompileOptions
@@ -15,6 +16,7 @@ def __init__(self):
         import traceback
 
         self.trace = traceback.format_stack()[0:-1]
+        self.stack_frames: NatalStackFrame = NatalStackFrame()
 
     def getDefinitionTrace(self) -> list[str]:
         return self.trace

pyteal/ast/for_.py

@@ -55,9 +55,10 @@ def __teal__(self, options: "CompileOptions"):
 
         stepStart, stepEnd = self.step.__teal__(options)
         stepEnd.setNextBlock(condStart)
+        stepEnd._sframes_container = self
         doEnd.setNextBlock(stepStart)
 
-        branchBlock = TealConditionalBlock([])
+        branchBlock = TealConditionalBlock([], root_expr=self)
         branchBlock.setTrueBlock(doStart)
         branchBlock.setFalseBlock(end)
 

pyteal/ast/frame.py

@@ -8,6 +8,7 @@
 from pyteal.types import TealType, require_type
 from pyteal.errors import TealInputError, TealInternalError, verifyProgramVersion
 from pyteal.ir import TealBlock, TealSimpleBlock, TealOp, Op
+from pyteal.stack_frame import NatalStackFrame
 
 if TYPE_CHECKING:
     from pyteal.compiler import CompileOptions
@@ -190,6 +191,7 @@ def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBloc
         proto_srt, proto_end = TealBlock.FromOp(options, op)
         local_srt, local_end = self.mem_layout.__teal__(options)
         proto_end.setNextBlock(local_srt)
+        NatalStackFrame.reframe_ops_in_blocks(self, proto_srt)
         return proto_srt, local_end
 
     def __str__(self) -> str: