Better Router Mapping

pyteal/ast/abi/method_return.py

@@ -20,7 +20,6 @@ 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:
@@ -30,7 +29,7 @@ def __teal__(self, options: "CompileOptions") -> Tuple[TealBlock, TealSimpleBloc
         start, end = Log(Concat(Bytes(RETURN_HASH_PREFIX), self.arg.encode())).__teal__(
             options
         )
-        NatalStackFrame.reframe_ops_in_blocks(self._sframes_container or self, start)
+        NatalStackFrame.reframe_ops_in_blocks(self, start)
         return start, end
 
     def __str__(self) -> str:

pyteal/ast/router.py

@@ -24,6 +24,7 @@
     OutputKwArgInfo,
     Subroutine,
     SubroutineCall,
+    SubroutineDefinition,
     SubroutineFnWrapper,
 )
 from pyteal.ast.txn import Txn
@@ -237,8 +238,7 @@ def approval_construction(self) -> Optional[Expr]:
             if oca.is_empty():
                 continue
             wrapped_handler = ASTBuilder.wrap_handler(
-                False,
-                cast(ActionType, oca.action),
+                False, cast(ActionType, oca.action)
             )
             match oca.call_config:
                 case CallConfig.ALL:
@@ -256,13 +256,12 @@ def approval_construction(self) -> Optional[Expr]:
                     raise TealInternalError(
                         f"Unexpected CallConfig: {oca.call_config!r}"
                     )
-            conditions_n_branches.append(
-                CondNode(
-                    Txn.on_completion() == oc,
-                    cond_body,
-                )
-            )
-        return Cond(*[[n.condition, n.branch] for n in conditions_n_branches])
+            cn = CondNode(Txn.on_completion() == oc, cond_body)
+            cn.reframe_asts(oca.stack_frames)
+            conditions_n_branches.append(cn)
+        cond = Cond(*[[n.condition, n.branch] for n in conditions_n_branches])
+        cond.stack_frames = self.stack_frames
+        return cond
 
     def get_method_config(self) -> MethodConfig:
         return MethodConfig(
@@ -288,6 +287,9 @@ class CondNode:
     condition: Expr
     branch: Expr
 
+    def reframe_asts(self, stack_frames: NatalStackFrame) -> None:
+        stack_frames.reframe(self.condition, self.branch)
+
 
 CondNode.__module__ = "pyteal"
 
@@ -328,21 +330,27 @@ def to_cond_node(self, use_frame_pt: bool = False) -> CondNode:
         if not (isinstance(self.condition, Expr) or self.condition == 1):
             raise TealInputError("Invalid condition input for CondWithMethod")
 
-        res = ASTBuilder.wrap_handler(True, self.method, use_frame_pt=use_frame_pt)
+        user_frames_holder: list[NatalStackFrame] = []
+        res = ASTBuilder.wrap_handler(
+            True,
+            self.method,
+            use_frame_pt=use_frame_pt,
+            handler_stack_frames_container=user_frames_holder,
+        )
+        assert (
+            ufhlen := len(user_frames_holder)
+        ) == 1, f"Unexpected length for user_frames_holder: {ufhlen}"
+        user_frames: NatalStackFrame = user_frames_holder[0]
+
         if isinstance(self.condition, Expr):
             res = Seq(Assert(self.condition), res)
-            NatalStackFrame.reframe_asts(self.condition.stack_frames, res)
-        return CondNode(walk_in_cond, res)
 
-
-CondWithMethod.__module__ = "pyteal"
+        cn = CondNode(walk_in_cond, res)
+        cn.reframe_asts(user_frames)
+        return cn
 
 
-def _smap_friendly_approve():
-    # TODO: Consider replacing _smap_friendly_approve() with a reframe_asts()
-    a = Approve()
-    a.stack_frames._compiler_gen = True
-    return a
+CondWithMethod.__module__ = "pyteal"
 
 
 @dataclass
@@ -513,6 +521,7 @@ def wrap_handler(
         *,
         wrap_to_name: str | None = None,
         use_frame_pt: bool = False,
+        handler_stack_frames_container: list[NatalStackFrame] | None = None,
     ) -> Expr:
         """This is a helper function that handles transaction arguments passing in bare-app-call/abi-method handlers.
         If `is_method_call` is True, then it can only be `ABIReturnSubroutine`,
@@ -531,6 +540,13 @@ def wrap_handler(
                 - for abi-method it returns the txn args correctly decomposed into ABI variables,
                   passed in ABIReturnSubroutine and logged, then approve.
         """
+
+        def scavenge(frames_holder) -> None:
+            if handler_stack_frames_container is not None:
+                handler_stack_frames_container.append(frames_holder.stack_frames)
+
+        handler_evald: abi.ReturnedValue | SubroutineCall
+
         if not is_method_call:
             wrap_to_name = wrap_to_name or "bare appcall"
 
@@ -540,11 +556,8 @@ def wrap_handler(
                         raise TealInputError(
                             f"{wrap_to_name} handler should be TealType.none not {handler.type_of()}."
                         )
-                    return (
-                        handler
-                        if handler.has_return()
-                        else Seq(handler, _smap_friendly_approve())
-                    )
+                    scavenge(handler)
+                    return handler if handler.has_return() else Seq(handler, Approve())
                 case SubroutineFnWrapper():
                     if handler.type_of() != TealType.none:
                         raise TealInputError(
@@ -555,7 +568,11 @@ def wrap_handler(
                             f"subroutine call should take 0 arg for {wrap_to_name}. "
                             f"this subroutine takes {handler.subroutine.argument_count()}."
                         )
-                    return Seq(handler(), _smap_friendly_approve())
+                    handler_evald = handler()
+                    if isinstance(handler_evald, abi.ReturnedValue):
+                        handler_evald = handler_evald.computation
+                    scavenge(cast(SubroutineCall, handler_evald).subroutine)
+                    return Seq(handler_evald, Approve())
                 case ABIReturnSubroutine():
                     if handler.type_of() != "void":
                         raise TealInputError(
@@ -566,7 +583,12 @@ def wrap_handler(
                             f"abi-returning subroutine call should take 0 arg for {wrap_to_name}. "
                             f"this abi-returning subroutine takes {handler.subroutine.argument_count()}."
                         )
-                    return Seq(cast(Expr, handler()), _smap_friendly_approve())
+                    handler_evald = handler()
+                    if isinstance(handler_evald, abi.ReturnedValue):
+                        handler_evald = handler_evald.computation
+
+                    scavenge(handler_evald)
+                    return Seq(cast(Expr, handler_evald), Approve())
                 case _:
                     raise TealInputError(
                         f"{wrap_to_name} can only accept: none type Expr, or Subroutine/ABIReturnSubroutine with none return and no arg"
@@ -583,13 +605,19 @@ def wrap_handler(
                 f"{wrap_to_name} ABIReturnSubroutine is not routable "
                 f"got {handler.subroutine.argument_count()} args with {len(handler.subroutine.abi_args)} ABI args."
             )
-        if not use_frame_pt:
-            return ASTBuilder.__de_abify_subroutine_vanilla(handler)
-        else:
-            return ASTBuilder.__de_abify_subroutine_frame_pointers(handler)
+
+        ret_expr, subdef = (
+            ASTBuilder.__de_abify_subroutine_frame_pointers(handler)
+            if use_frame_pt
+            else ASTBuilder.__de_abify_subroutine_vanilla(handler)
+        )
+        scavenge(subdef)
+        return ret_expr
 
     @staticmethod
-    def __de_abify_subroutine_vanilla(handler: ActionType) -> Expr:
+    def __de_abify_subroutine_vanilla(
+        handler: ActionType,
+    ) -> tuple[Expr, SubroutineDefinition]:
         """This private function retains the previous (pre-frame-pointer) logic of handling ABIReturnSubroutine method's IO.
 
         This function can be roughly separated into following 4 parts:
@@ -638,28 +666,35 @@ def __de_abify_subroutine_vanilla(handler: ActionType) -> Expr:
             arg_vals, app_arg_vals, txn_arg_vals, handler
         )
 
+        ret_expr: Expr
+        handler_evald: abi.ReturnedValue | SubroutineCall
         if handler.type_of() == sdk_abi.Returns.VOID:
-            return Seq(
+            ret_expr = Seq(
                 *decode_instructions,
-                cast(SubroutineCall, handler(*arg_vals)),
-                _smap_friendly_approve(),
+                cast(SubroutineCall, handler_evald := handler(*arg_vals)),
+                Approve(),
             )
         else:
             output_temp: abi.BaseType = cast(
                 OutputKwArgInfo, handler.output_kwarg_info
             ).abi_type.new_instance()
-            returned_val: abi.ReturnedValue = cast(
-                abi.ReturnedValue, handler(*arg_vals)
-            )
-            return Seq(
+            handler_evald = cast(abi.ReturnedValue, handler(*arg_vals))
+            ret_expr = Seq(
                 *decode_instructions,
-                returned_val.store_into(output_temp),
+                handler_evald.store_into(output_temp),
                 abi.MethodReturn(output_temp),
-                _smap_friendly_approve(),
+                Approve(),
             )
 
+        if isinstance(handler_evald, abi.ReturnedValue):
+            handler_evald = handler_evald.computation
+
+        return ret_expr, cast(SubroutineCall, handler_evald).subroutine
+
     @staticmethod
-    def __de_abify_subroutine_frame_pointers(handler: ActionType) -> Expr:
+    def __de_abify_subroutine_frame_pointers(
+        handler: ActionType,
+    ) -> tuple[Expr, SubroutineDefinition]:
         """This private function implements the frame-pointer-based logic of handling ABIReturnSubroutine method's IO.
 
         This function can be roughly separated into following 4 parts:
@@ -737,15 +772,16 @@ def __de_abify_subroutine_frame_pointers(handler: ActionType) -> Expr:
         ]
         returning_steps: list[Expr]
 
+        handler_evald: abi.ReturnedValue | SubroutineCall
         if handler.type_of() == sdk_abi.Returns.VOID:
-            returning_steps = [cast(SubroutineCall, handler(*arg_vals))]
+            returning_steps = [cast(Expr, handler_evald := handler(*arg_vals))]
         else:
             output_temp: abi.BaseType = cast(
                 OutputKwArgInfo, handler.output_kwarg_info
             ).abi_type.new_instance()
             output_temp._stored_value = FrameVar(proto, 0)
             returned_val: abi.ReturnedValue = cast(
-                abi.ReturnedValue, handler(*arg_vals)
+                abi.ReturnedValue, handler_evald := handler(*arg_vals)
             )
             returning_steps = [
                 returned_val.store_into(output_temp),
@@ -755,7 +791,13 @@ def __de_abify_subroutine_frame_pointers(handler: ActionType) -> Expr:
         def declaration():
             return Seq(*decoding_steps, *returning_steps)
 
-        return Seq(subroutine_caster(declaration)(), _smap_friendly_approve())
+        if isinstance(handler_evald, abi.ReturnedValue):
+            handler_evald = handler_evald.computation
+
+        return (
+            Seq(subroutine_caster(declaration)(), Approve()),
+            cast(SubroutineCall, handler_evald).subroutine,
+        )
 
     def add_method_to_ast(
         self, method_signature: str, cond: Expr | int, handler: ABIReturnSubroutine
@@ -801,19 +843,21 @@ class _RouterBundle:
     input: Optional["_RouterCompileInput"] = None
 
     def get_results(self) -> RouterResults:
-        approval_sm: PyTealSourceMap | None = None
-        clear_sm: PyTealSourceMap | None = None
+        approval_sourcemap: PyTealSourceMap | None = None
+        clear_sourcemap: PyTealSourceMap | None = None
         if self.approval_sourcemapper:
-            approval_sm = self.approval_sourcemapper.get_sourcemap(self.approval_teal)
+            approval_sourcemap = self.approval_sourcemapper.get_sourcemap(
+                self.approval_teal
+            )
         if self.clear_sourcemapper:
-            clear_sm = self.clear_sourcemapper.get_sourcemap(self.clear_teal)
+            clear_sourcemap = self.clear_sourcemapper.get_sourcemap(self.clear_teal)
 
         return RouterResults(
-            self.approval_teal,
-            self.clear_teal,
-            self.abi_contract,
-            approval_sm,
-            clear_sm,
+            approval_teal=self.approval_teal,
+            clear_teal=self.clear_teal,
+            abi_contract=self.abi_contract,
+            approval_sourcemap=approval_sourcemap,
+            clear_sourcemap=clear_sourcemap,
         )
 
 
@@ -1106,7 +1150,7 @@ def _build_program(
                         act := cast(Expr, bare_call_approval),
                     )
                 ]
-                NatalStackFrame.reframe_asts(bare_call_approval.stack_frames, cond)
+                bare_call_approval.stack_frames.reframe(cond)
                 act.stack_frames = bare_call_approval.stack_frames
 
         optimize = optimize or OptimizeOptions()
@@ -1134,8 +1178,6 @@ def compile_program(
         optimize: Optional[OptimizeOptions] = None,
     ) -> tuple[str, str, sdk_abi.Contract]:
         """
-        DEPRECATED BUT KEPT FOR BACKWARDS COMPATIBILITY. PREFER Router.compile().
-
         Constructs and compiles approval and clear-state programs from the registered methods and
         bare app calls in the router, and also generates a Contract object to allow client read and call
         the methods easily.
@@ -1151,6 +1193,8 @@ def compile_program(
             * approval_program: compiled approval program string
             * clear_state_program: compiled clear-state program string
             * contract: a Python SDK Contract object to allow clients to make off-chain calls
+
+        NOTE: For generating a source map, please refer to the `Router.compile` method.
         """
         input = _RouterCompileInput(
             version=version,

pyteal/ast/router_test.py

@@ -382,6 +382,17 @@ def test_bare_call_config_clear_state_failure():
     )
 
 
+def equal_ocas(oca1, oca2):
+    assert oca1.call_config == oca2.call_config
+    if oca1.action is None:
+        assert oca2.action is None
+    else:
+        assert isinstance(oca1.action, pt.Int)
+        assert isinstance(oca2.action, pt.Int)
+
+    return True
+
+
 def test_BareCallActions_asdict():
     no_action = pt.OnCompleteAction()
     del_action = pt.OnCompleteAction(action=pt.Int(1), call_config=pt.CallConfig.ALL)
@@ -392,16 +403,7 @@ def test_BareCallActions_asdict():
         close_out=close_action,
     )
 
-    bcad = bca.asdict()
-    assert set(bcad.keys()) == {
-        "clear_state",
-        "close_out",
-        "delete_application",
-        "no_op",
-        "opt_in",
-        "update_application",
-    }
-    assert bcad == {
+    expected = {
         "clear_state": no_action,
         "close_out": close_action,
         "delete_application": del_action,
@@ -410,6 +412,13 @@ def test_BareCallActions_asdict():
         "update_application": no_action,
     }
 
+    bcad = bca.asdict()
+    assert set(bcad.keys()) == set(expected.keys())
+
+    for k, oca1 in bcad.items():
+        oca2 = expected[k]
+        equal_ocas(oca1, oca2)
+
 
 def test_BareCallActions_aslist():
     no_action = pt.OnCompleteAction()
@@ -423,15 +432,18 @@ def test_BareCallActions_aslist():
         opt_in=optin_action,
     )
 
-    bcal = bca.aslist()
-    assert bcal == [
+    expected = [
         no_action,
         no_action,
         no_action,
         no_action,
         optin_action,
         update_action,
     ]
+    bcal = bca.aslist()
+    assert len(expected) == len(bcal)
+
+    assert all([equal_ocas(expected[i], actual) for i, actual in enumerate(bcal)])
 
 
 def test_BareCallActions_get_method_config():