unify subroutine case and abi-returning-subroutine

algorand · Jun 1, 2022 · 10ce036 · 10ce036
1 parent c2bbd39
commit 10ce036
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 5 deletions.
diff --git a/tests/blackbox.py b/tests/blackbox.py
@@ -70,9 +70,7 @@ def _fill(
         self, input_types: list[TealType | None]
     ) -> list[TealType | abi.TypeSpec | None]:
         match self.subroutine:
-            case SubroutineFnWrapper():
-                return cast(list[TealType | abi.TypeSpec | None], input_types)
-            case ABIReturnSubroutine():
+            case SubroutineFnWrapper() | ABIReturnSubroutine():
                 args = self.subroutine.subroutine.arguments()
                 abis = self.subroutine.subroutine.abi_args
                 return [(x if x else abis[args[i]]) for i, x in enumerate(input_types)]
@@ -288,17 +286,19 @@ def _handle_SubroutineFnWrapper(self):
 
         def arg_prep_n_call(i, p):
             name = arg_names[i]
-            by_ref = name in subdef.by_ref_args
             arg_expr = (
                 Txn.application_args[i] if self.mode == Mode.Application else Arg(i)
             )
             if p == TealType.uint64:
                 arg_expr = Btoi(arg_expr)
             prep = None
             arg_var = arg_expr
-            if by_ref:
+            if name in subdef.by_ref_args:
                 arg_var = ScratchVar(p)
                 prep = arg_var.store(arg_expr)
+            elif name in subdef.abi_args:
+                arg_var = p.new_instance()
+                prep = arg_var.decode(arg_expr)
             return prep, arg_var
 
         def subr_caller():

diff --git a/tests/integration/graviton_test.py b/tests/integration/graviton_test.py
@@ -829,6 +829,33 @@ def report(kind):
     report("lsig")
 
 
+def blackbox_pyteal_example5():
+    from pyteal import (
+        abi,
+        Sqrt,
+        Seq,
+        Subroutine,
+        TealType,
+        Expr,
+        ScratchVar,
+        Assert,
+        Int,
+        For,
+    )
+
+    @Blackbox([TealType.uint64])
+    @Subroutine(TealType.uint64)
+    def primality_from_sieve(n: Expr):
+        i = ScratchVar(TealType.uint64)
+        return Seq(
+            Assert(n > Int(1)),
+            For(i.store(Int(2)), i.load() <= Sqrt(n), i.store(i.load() + Int(1))).Do(),
+        )
+        pass
+
+    pass
+
+
 def blackbox_pyteal_while_continue_test():
     from tests.blackbox import Blackbox
     from pyteal import (