Add handler for allowing Txna in InnerTxnBuilder (#384)

Allow SetField(s) to accept the foreign reference array directly

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 ## Added
+* Add the ability to pass foreign reference arrays directly into inner transactions ([#384](https://github.com/algorand/pyteal/pull/384))
 
 ## Fixed
 

pyteal/ast/abi/__init__.py

@@ -69,6 +69,7 @@
     make,
     size_of,
     type_spec_from_annotation,
+    type_specs_from_signature,
     contains_type_spec,
 )
 
@@ -140,6 +141,7 @@
     "KeyRegisterTransactionTypeSpec",
     "TransactionTypeSpecs",
     "type_spec_from_annotation",
+    "type_specs_from_signature",
     "make",
     "size_of",
     "algosdk_from_annotation",

pyteal/ast/itxn.py

@@ -1,15 +1,25 @@
 from enum import Enum
-from typing import TYPE_CHECKING, cast
 import algosdk
 
-from pyteal.ast.abi.util import type_specs_from_signature
+from typing import TYPE_CHECKING, cast
 from pyteal.ast.int import EnumInt
+from pyteal.ast.for_ import For
+from pyteal.ast.int import Int
+from pyteal.ast.scratchvar import ScratchVar
+
 from pyteal.ast.methodsig import MethodSignature
 from pyteal.types import TealType, require_type
 from pyteal.errors import TealInputError, TealTypeError, verifyProgramVersion
 from pyteal.ir import TealOp, Op, TealBlock
 from pyteal.ast.expr import Expr
-from pyteal.ast.txn import TxnField, TxnExprBuilder, TxnType, TxnaExprBuilder, TxnObject
+from pyteal.ast.txn import (
+    TxnType,
+    TxnArray,
+    TxnField,
+    TxnExprBuilder,
+    TxnaExprBuilder,
+    TxnObject,
+)
 from pyteal.ast.seq import Seq
 from pyteal.ast.bytes import Bytes
 from pyteal.ast import abi
@@ -149,6 +159,7 @@ def SetField(cls, field: TxnField, value: Expr | list[Expr]) -> Expr:
 
         Note: For non-array field (e.g., note), setting it twice will overwrite the original value.
               While for array field (e.g., accounts), setting it multiple times will append the values.
+              It is also possible to pass the entire array field if desired (e.g., Txn.accounts) to pass all the references.
 
         Requires program version 5 or higher. This operation is only permitted in application mode.
 
@@ -158,33 +169,43 @@ def SetField(cls, field: TxnField, value: Expr | list[Expr]) -> Expr:
                 compatible with the field being set.
         """
         if not field.is_array:
-            if type(value) is list:
+            if type(value) is list or isinstance(value, TxnArray):
                 raise TealInputError(
                     "inner transaction set field {} does not support array value".format(
                         field
                     )
                 )
             return InnerTxnFieldExpr(field, cast(Expr, value))
         else:
-            if type(value) is not list:
+            if type(value) is not list and not isinstance(value, TxnArray):
                 raise TealInputError(
                     "inner transaction set array field {} with non-array value".format(
                         field
                     )
                 )
-            for valueIter in value:
-                if not isinstance(valueIter, Expr):
-                    raise TealInputError(
-                        "inner transaction set array field {} with non PyTeal expression array element {}".format(
-                            field, valueIter
+
+            if type(value) is list:
+                for valueIter in value:
+                    if not isinstance(valueIter, Expr):
+                        raise TealInputError(
+                            "inner transaction set array field {} with non PyTeal expression array element {}".format(
+                                field, valueIter
+                            )
                         )
-                    )
-            return Seq(
-                *[
-                    InnerTxnFieldExpr(field, cast(Expr, valueIter))
-                    for valueIter in value
-                ]
-            )
+
+                return Seq(
+                    *[
+                        InnerTxnFieldExpr(field, cast(Expr, valueIter))
+                        for valueIter in value
+                    ]
+                )
+            else:
+                arr = cast(TxnArray, value)
+                return For(
+                    (i := ScratchVar()).store(Int(0)),
+                    i.load() < arr.length(),
+                    i.store(i.load() + Int(1)),
+                ).Do(InnerTxnFieldExpr(field, arr[i.load()]))
 
     @classmethod
     def Execute(cls, fields: dict[TxnField, Expr | list[Expr]]) -> Expr:
@@ -216,6 +237,7 @@ def SetFields(cls, fields: dict[TxnField, Expr | list[Expr]]) -> Expr:
 
         Note: For non-array field (e.g., note), setting it twice will overwrite the original value.
               While for array field (e.g., accounts), setting it multiple times will append the values.
+              It is also possible to pass the entire array field if desired (e.g., Txn.accounts) to pass all the references.
 
         Requires program version 5 or higher. This operation is only permitted in application mode.
 
@@ -273,7 +295,7 @@ def MethodCall(
 
         # We only care about the args
         arg_type_specs: list[abi.TypeSpec]
-        arg_type_specs, _ = type_specs_from_signature(method_signature)
+        arg_type_specs, _ = abi.type_specs_from_signature(method_signature)
 
         if len(args) != len(arg_type_specs):
             raise TealInputError(

pyteal/ast/itxn_test.py

@@ -1,7 +1,7 @@
 import pytest
 
 import pyteal as pt
-from pyteal.ast.txn import TxnField, TxnType
+from pyteal.ast.txn import Txn, TxnField, TxnType
 from pyteal.types import types_match
 
 avm4Options = pt.CompileOptions(version=4)
@@ -111,6 +111,16 @@ def test_InnerTxnBuilder_SetField():
             ),
         ),
     ),
+    (
+        {pt.TxnField.accounts: pt.Txn.accounts},
+        pt.For(
+            (i := pt.ScratchVar()).store(pt.Int(0)),
+            i.load() < pt.Txn.accounts.length(),
+            i.store(i.load() + pt.Int(1)),
+        ).Do(
+            pt.InnerTxnBuilder.SetField(pt.TxnField.accounts, [Txn.accounts[i.load()]])
+        ),
+    ),
 ]
 
 
@@ -128,9 +138,14 @@ def test_InnerTxnBuilder_SetFields():
         actual.addIncoming()
         actual = pt.TealBlock.NormalizeBlocks(actual)
 
-        with pt.TealComponent.Context.ignoreExprEquality():
+        with pt.TealComponent.Context.ignoreScratchSlotEquality(), pt.TealComponent.Context.ignoreExprEquality():
             assert actual == expected
 
+        assert pt.TealBlock.MatchScratchSlotReferences(
+            pt.TealBlock.GetReferencedScratchSlots(actual),
+            pt.TealBlock.GetReferencedScratchSlots(expected),
+        )
+
         if len(fields) != 0:
             with pytest.raises(pt.TealInputError):
                 expr.__teal__(avm4Options)
@@ -152,9 +167,14 @@ def test_InnerTxnBuilder_Execute():
         actual.addIncoming()
         actual = pt.TealBlock.NormalizeBlocks(actual)
 
-        with pt.TealComponent.Context.ignoreExprEquality():
+        with pt.TealComponent.Context.ignoreScratchSlotEquality(), pt.TealComponent.Context.ignoreExprEquality():
             assert actual == expected
 
+        assert pt.TealBlock.MatchScratchSlotReferences(
+            pt.TealBlock.GetReferencedScratchSlots(actual),
+            pt.TealBlock.GetReferencedScratchSlots(expected),
+        )
+
         with pytest.raises(pt.TealInputError):
             expr.__teal__(avm4Options)