Use deferred subroutine expression for ABI returns

pyteal/ast/subroutine.py

@@ -11,7 +11,6 @@
 
 from pyteal.ast import abi
 from pyteal.ast.expr import Expr
-from pyteal.ast.return_ import Return
 from pyteal.ast.seq import Seq
 from pyteal.ast.scratchvar import DynamicScratchVar, ScratchVar
 from pyteal.errors import TealInputError, verifyTealVersion
@@ -302,10 +301,16 @@ def __hash__(self):
 
 
 class SubroutineDeclaration(Expr):
-    def __init__(self, subroutine: SubroutineDefinition, body: Expr) -> None:
+    def __init__(
+        self,
+        subroutine: SubroutineDefinition,
+        body: Expr,
+        deferred_expr: Optional[Expr] = None,
+    ) -> None:
         super().__init__()
         self.subroutine = subroutine
         self.body = body
+        self.deferred_expr = deferred_expr
 
     def __teal__(self, options: "CompileOptions"):
         return self.body.__teal__(options)
@@ -720,24 +725,21 @@ def var_n_loaded(
         raise TealInputError(
             f"Subroutine function does not return a PyTeal expression. Got type {type(subroutine_body)}."
         )
+
+    deferred_expr: Optional[Expr] = None
+
     # if there is an output keyword argument for ABI, place the storing on the stack
     if output_carrying_abi:
-        if subroutine_body.has_return():
-            raise TealInputError(
-                "ABI returning subroutine definition should have no return"
-            )
         if subroutine_body.type_of() != TealType.none:
             raise TealInputError(
                 f"ABI returning subroutine definition should evaluate to TealType.none, "
                 f"while evaluate to {subroutine_body.type_of()}."
             )
-        subroutine_body = Seq(
-            subroutine_body, Return(output_carrying_abi.stored_value.load())
-        )
+        deferred_expr = output_carrying_abi.stored_value.load()
 
     # Arg usage "A" to be pick up and store in scratch parameters that have been placed on the stack
     # need to reverse order of argumentVars because the last argument will be on top of the stack
     body_ops = [var.slot.store() for var in arg_vars[::-1]]
     body_ops.append(subroutine_body)
 
-    return SubroutineDeclaration(subroutine, Seq(body_ops))
+    return SubroutineDeclaration(subroutine, Seq(body_ops), deferred_expr)

pyteal/compiler/compiler.py

@@ -10,7 +10,7 @@
     SubroutineDefinition,
     SubroutineDeclaration,
 )
-from pyteal.ir import Mode, TealComponent, TealOp, TealBlock, TealSimpleBlock
+from pyteal.ir import Mode, Op, TealComponent, TealOp, TealBlock, TealSimpleBlock
 from pyteal.errors import TealInputError, TealInternalError
 
 from pyteal.compiler.sort import sortBlocks
@@ -133,26 +133,60 @@ def compileSubroutine(
             ast = Return(ast)
 
     options.setSubroutine(currentSubroutine)
+
     start, end = ast.__teal__(options)
     start.addIncoming()
     start.validateTree()
 
-    start = TealBlock.NormalizeBlocks(start)
-    start.validateTree()
+    if (
+        currentSubroutine is not None
+        and currentSubroutine.get_declaration().deferred_expr is not None
+    ):
+        # this represents code that should be inserted before each retsub op
+        deferred_expr = cast(Expr, currentSubroutine.get_declaration().deferred_expr)
+
+        for block in TealBlock.Iterate(start):
+            if not any(op.getOp() == Op.retsub for op in block.ops):
+                continue
+
+            if len(block.ops) != 1:
+                # we expect all retsub ops to be in their own block at this point since
+                # TealBlock.NormalizeBlocks has not yet been used
+                raise TealInternalError(
+                    f"Expected retsub to be the only op in the block, but there are {len(block.ops)} ops"
+                )
 
-    order = sortBlocks(start, end)
-    teal = flattenBlocks(order)
+            # we invoke __teal__ here and not outside of this loop because the same block cannot be
+            # added in multiple places to the control flow graph
+            deferred_start, deferred_end = deferred_expr.__teal__(options)
+            deferred_start.addIncoming()
+            deferred_start.validateTree()
 
-    verifyOpsForVersion(teal, options.version)
-    verifyOpsForMode(teal, options.mode)
+            # insert deferred blocks between the previous block(s) and this one
+            deferred_start.incoming = block.incoming
+            block.incoming = [deferred_end]
+            deferred_end.nextBlock = block
+
+            for prev in deferred_start.incoming:
+                prev.replaceOutgoing(block, deferred_start)
+
+            if block is start:
+                # this is the start block, replace start
+                start = deferred_start
+
+    start.validateTree()
+
+    start = TealBlock.NormalizeBlocks(start)
+    start.validateTree()
 
     subroutine_start_blocks[currentSubroutine] = start
     subroutine_end_blocks[currentSubroutine] = end
 
     referencedSubroutines: Set[SubroutineDefinition] = set()
-    for stmt in teal:
-        for subroutine in stmt.getSubroutines():
-            referencedSubroutines.add(subroutine)
+    for block in TealBlock.Iterate(start):
+        for stmt in block.ops:
+            for subroutine in stmt.getSubroutines():
+                referencedSubroutines.add(subroutine)
 
     if currentSubroutine is not None:
         subroutineGraph[currentSubroutine] = referencedSubroutines
@@ -256,6 +290,9 @@ def compileTeal(
     subroutineLabels = resolveSubroutines(subroutineMapping)
     teal = flattenSubroutines(subroutineMapping, subroutineLabels)
 
+    verifyOpsForVersion(teal, options.version)
+    verifyOpsForMode(teal, options.mode)
+
     if assembleConstants:
         if version < 3:
             raise TealInternalError(

pyteal/compiler/compiler_test.py

@@ -1627,6 +1627,194 @@ def storeValue(key: pt.Expr, t1: pt.Expr, t2: pt.Expr, t3: pt.Expr) -> pt.Expr:
     assert actual == expected
 
 
+def test_compile_subroutine_deferred_expr():
+    @pt.Subroutine(pt.TealType.none)
+    def deferredExample(value: pt.Expr) -> pt.Expr:
+        return pt.Seq(
+            pt.If(value == pt.Int(0)).Then(pt.Return()),
+            pt.If(value == pt.Int(1)).Then(pt.Approve()),
+            pt.If(value == pt.Int(2)).Then(pt.Reject()),
+            pt.If(value == pt.Int(3)).Then(pt.Err()),
+        )
+
+    program = pt.Seq(deferredExample(pt.Int(10)), pt.Approve())
+
+    expected_no_deferred = """#pragma version 6
+int 10
+callsub deferredExample_0
+int 1
+return
+
+// deferredExample
+deferredExample_0:
+store 0
+load 0
+int 0
+==
+bnz deferredExample_0_l7
+load 0
+int 1
+==
+bnz deferredExample_0_l6
+load 0
+int 2
+==
+bnz deferredExample_0_l5
+load 0
+int 3
+==
+bz deferredExample_0_l8
+err
+deferredExample_0_l5:
+int 0
+return
+deferredExample_0_l6:
+int 1
+return
+deferredExample_0_l7:
+retsub
+deferredExample_0_l8:
+retsub
+    """.strip()
+    actual_no_deferred = pt.compileTeal(
+        program, pt.Mode.Application, version=6, assembleConstants=False
+    )
+    assert actual_no_deferred == expected_no_deferred
+
+    # manually add deferred expression to SubroutineDefinition
+    declaration = deferredExample.subroutine.get_declaration()
+    declaration.deferred_expr = pt.Pop(pt.Bytes("deferred"))
+
+    expected_deferred = """#pragma version 6
+int 10
+callsub deferredExample_0
+int 1
+return
+
+// deferredExample
+deferredExample_0:
+store 0
+load 0
+int 0
+==
+bnz deferredExample_0_l7
+load 0
+int 1
+==
+bnz deferredExample_0_l6
+load 0
+int 2
+==
+bnz deferredExample_0_l5
+load 0
+int 3
+==
+bz deferredExample_0_l8
+err
+deferredExample_0_l5:
+int 0
+return
+deferredExample_0_l6:
+int 1
+return
+deferredExample_0_l7:
+byte "deferred"
+pop
+retsub
+deferredExample_0_l8:
+byte "deferred"
+pop
+retsub
+    """.strip()
+    actual_deferred = pt.compileTeal(
+        program, pt.Mode.Application, version=6, assembleConstants=False
+    )
+    assert actual_deferred == expected_deferred
+
+
+def test_compile_subroutine_deferred_expr_empty():
+    @pt.Subroutine(pt.TealType.none)
+    def empty() -> pt.Expr:
+        return pt.Return()
+
+    program = pt.Seq(empty(), pt.Approve())
+
+    expected_no_deferred = """#pragma version 6
+callsub empty_0
+int 1
+return
+
+// empty
+empty_0:
+retsub
+    """.strip()
+    actual_no_deferred = pt.compileTeal(
+        program, pt.Mode.Application, version=6, assembleConstants=False
+    )
+    assert actual_no_deferred == expected_no_deferred
+
+    # manually add deferred expression to SubroutineDefinition
+    declaration = empty.subroutine.get_declaration()
+    declaration.deferred_expr = pt.Pop(pt.Bytes("deferred"))
+
+    expected_deferred = """#pragma version 6
+callsub empty_0
+int 1
+return
+
+// empty
+empty_0:
+byte "deferred"
+pop
+retsub
+    """.strip()
+    actual_deferred = pt.compileTeal(
+        program, pt.Mode.Application, version=6, assembleConstants=False
+    )
+    assert actual_deferred == expected_deferred
+
+
+def test_compileSubroutine_deferred_block_malformed():
+    class BadRetsub(pt.Expr):
+        def type_of(self) -> pt.TealType:
+            return pt.TealType.none
+
+        def has_return(self) -> bool:
+            return True
+
+        def __str__(self) -> str:
+            return "(BadRetsub)"
+
+        def __teal__(
+            self, options: pt.CompileOptions
+        ) -> tuple[pt.TealBlock, pt.TealSimpleBlock]:
+            block = pt.TealSimpleBlock(
+                [
+                    pt.TealOp(self, pt.Op.int, 1),
+                    pt.TealOp(self, pt.Op.pop),
+                    pt.TealOp(self, pt.Op.retsub),
+                ]
+            )
+
+            return block, block
+
+    @pt.Subroutine(pt.TealType.none)
+    def bad() -> pt.Expr:
+        return BadRetsub()
+
+    program = pt.Seq(bad(), pt.Approve())
+
+    # manually add deferred expression to SubroutineDefinition
+    declaration = bad.subroutine.get_declaration()
+    declaration.deferred_expr = pt.Pop(pt.Bytes("deferred"))
+
+    with pytest.raises(
+        pt.TealInternalError,
+        match=r"^Expected retsub to be the only op in the block, but there are 3 ops$",
+    ):
+        pt.compileTeal(program, pt.Mode.Application, version=6, assembleConstants=False)
+
+
 def test_compile_wide_ratio():
     cases = (
         (