Add opcode support for itxnas and gitxnas (#193)

michaeldiamant · web-flow · commit 2a682332bcec · 2022-02-17T10:51:29.000-05:00
* Add opcode support for itxnas and gitxnas

* Update stale reference to inner transaction limit

* Fix allowed types for GitxnaExpr txnIndex

* Remove obsolete logic for handling GitxnaExpr.teal construction

* Remove unnecessary cast and fix gitxna runtime type checking

* Move type validation to constructors for gtxn and gitxn variants

* Add missed tests from prior commit

* Fix duplicate test case

* Move index validation from subclasses to TxnaExpr

* Inline validation functions per PR feedback

* Remove unused imports

* Refactor to isinstance tupled check

* Remove TEAL v1 min version test per PR feedback

* Fix constructor type checking for GtxnExpr

* Refactor to remove duplicate type check function
diff --git a/pyteal/ast/gitxn.py b/pyteal/ast/gitxn.py
@@ -1,10 +1,12 @@
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, cast, Union
 
 from pyteal.config import MAX_GROUP_SIZE
 
 from ..errors import TealInputError, verifyFieldVersion, verifyTealVersion
 from ..ir import TealOp, Op, TealBlock
+from .expr import Expr
 from .txn import TxnExpr, TxnField, TxnObject, TxnaExpr
+from ..types import require_type, TealType
 
 if TYPE_CHECKING:
     from ..compiler import CompileOptions
@@ -15,6 +17,15 @@ class GitxnExpr(TxnExpr):
 
     def __init__(self, txnIndex: int, field: TxnField) -> None:
         super().__init__(Op.gitxn, "Gitxn", field)
+
+        # Currently we do not have gitxns. Only gitxn with immediate transaction index supported.
+        if type(txnIndex) is not int:
+            raise TealInputError(
+                "Invalid gitxn syntax with immediate transaction field {} and transaction index {}".format(
+                    field, txnIndex
+                )
+            )
+
         self.txnIndex = txnIndex
 
     def __str__(self):
@@ -23,14 +34,6 @@ def __str__(self):
     def __teal__(self, options: "CompileOptions"):
         verifyFieldVersion(self.field.arg_name, self.field.min_version, options.version)
 
-        # currently we do not have gitxns, only gitxn with immediate transaction index supported
-        if type(self.txnIndex) is not int:
-            raise TealInputError(
-                "Invalid gitxn syntax with immediate transaction field {} and transaction index {}".format(
-                    self.field, self.txnIndex
-                )
-            )
-
         verifyTealVersion(
             Op.gitxn.min_version,
             options.version,
@@ -46,8 +49,14 @@ def __teal__(self, options: "CompileOptions"):
 class GitxnaExpr(TxnaExpr):
     """An expression that accesses an inner transaction array field from an inner transaction in the last inner group."""
 
-    def __init__(self, txnIndex: int, field: TxnField, index: int) -> None:
-        super().__init__(Op.gitxna, None, "Gitxna", field, index)
+    def __init__(self, txnIndex: int, field: TxnField, index: Union[int, Expr]) -> None:
+        super().__init__(Op.gitxna, Op.gitxnas, "Gitxna", field, index)
+
+        if type(txnIndex) is not int:
+            raise TealInputError(
+                f"Invalid txnIndex type:  Expected int, but received {txnIndex}."
+            )
+
         self.txnIndex = txnIndex
 
     def __str__(self):
@@ -57,18 +66,23 @@ def __str__(self):
 
     def __teal__(self, options: "CompileOptions"):
         verifyFieldVersion(self.field.arg_name, self.field.min_version, options.version)
-        if type(self.txnIndex) is not int or type(self.index) is not int:
-            raise TealInputError(
-                "Invalid gitxna syntax with immediate transaction index {}, transaction field {}, array index {}".format(
-                    self.txnIndex, self.field, self.index
-                )
-            )
+
+        if type(self.index) is int:
+            opToUse = Op.gitxna
+        else:
+            opToUse = Op.gitxnas
 
         verifyTealVersion(
-            Op.gitxna.min_version, options.version, "TEAL version too low to use gitxna"
+            opToUse.min_version,
+            options.version,
+            "TEAL version too low to use op {}".format(opToUse),
         )
-        op = TealOp(self, Op.gitxna, self.txnIndex, self.field.arg_name, self.index)
-        return TealBlock.FromOp(options, op)
+
+        if type(self.index) is int:
+            op = TealOp(self, opToUse, self.txnIndex, self.field.arg_name, self.index)
+            return TealBlock.FromOp(options, op)
+        op = TealOp(self, opToUse, self.txnIndex, self.field.arg_name)
+        return TealBlock.FromOp(options, op, cast(Expr, self.index))
 
 
 GitxnaExpr.__module__ = "pyteal"
@@ -85,14 +99,14 @@ def __getitem__(self, txnIndex: int) -> TxnObject:
 
         if txnIndex < 0 or txnIndex >= MAX_GROUP_SIZE:
             raise TealInputError(
-                "Invalid Gtxn index {}, shoud be in [0, {})".format(
+                "Invalid Gtxn index {}, should be in [0, {})".format(
                     txnIndex, MAX_GROUP_SIZE
                 )
             )
 
         return TxnObject(
             lambda field: GitxnExpr(txnIndex, field),
-            lambda field, index: GitxnaExpr(txnIndex, field, cast(int, index)),
+            lambda field, index: GitxnaExpr(txnIndex, field, index),
         )
 
 
diff --git a/pyteal/ast/gitxn_test.py b/pyteal/ast/gitxn_test.py
@@ -7,24 +7,75 @@
 
 
 def test_gitxn_invalid():
-    with pytest.raises(TealInputError):
-        GitxnExpr(0, TxnField.sender).__teal__(teal5Options)
+    for ctor, e in [
+        (
+            lambda: Gitxn[MAX_GROUP_SIZE],
+            TealInputError,
+        ),
+        (
+            lambda: Gitxn[-1],
+            TealInputError,
+        ),
+    ]:
+        with pytest.raises(e):
+            ctor()
 
-    with pytest.raises(TealInputError):
-        Gitxn[MAX_GROUP_SIZE].sender()
 
-    with pytest.raises(TealInputError):
-        Gitxn[-1].asset_sender()
+def test_gitxn_valid():
+    for i in range(MAX_GROUP_SIZE):
+        Gitxn[i].sender()
 
-    with pytest.raises(TealInputError):
-        Gitxn[Bytes("first")].sender()
 
+def test_gitxn_expr_invalid():
+    for f, e in [
+        (
+            lambda: GitxnExpr(Int(1), TxnField.sender),
+            TealInputError,
+        ),
+        (
+            lambda: GitxnExpr(1, TxnField.sender).__teal__(teal5Options),
+            TealInputError,
+        ),
+    ]:
+        with pytest.raises(e):
+            f()
 
-def test_gitxn_valid():
-    GitxnaExpr(0, TxnField.application_args, 1).__teal__(teal6Options)
 
-    for i in range(MAX_GROUP_SIZE):
-        Gitxn[i].sender()
+def test_gitxn_expr_valid():
+    GitxnExpr(1, TxnField.sender).__teal__(teal6Options)
+
+
+def test_gitxna_expr_invalid():
+    for f, e in [
+        (
+            lambda: GitxnaExpr("Invalid_type", TxnField.application_args, 1),
+            TealInputError,
+        ),
+        (
+            lambda: GitxnaExpr(1, TxnField.application_args, "Invalid_type"),
+            TealInputError,
+        ),
+        (
+            lambda: GitxnaExpr(0, TxnField.application_args, Assert(Int(1))),
+            TealTypeError,
+        ),
+        (
+            lambda: GitxnaExpr(0, TxnField.application_args, 0).__teal__(teal5Options),
+            TealInputError,
+        ),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gitxna_valid():
+    [
+        e.__teal__(teal6Options)
+        for e in [
+            GitxnaExpr(0, TxnField.application_args, 1),
+            GitxnaExpr(0, TxnField.application_args, Int(1)),
+        ]
+    ]
 
 
 # txn_test.py performs additional testing
diff --git a/pyteal/ast/gtxn.py b/pyteal/ast/gtxn.py
@@ -11,11 +11,21 @@
     from ..compiler import CompileOptions
 
 
+def validate_txn_index_or_throw(txnIndex: Union[int, Expr]):
+    if not isinstance(txnIndex, (int, Expr)):
+        raise TealInputError(
+            f"Invalid txnIndex type:  Expected int or Expr, but received {txnIndex}"
+        )
+    if isinstance(txnIndex, Expr):
+        require_type(txnIndex, TealType.uint64)
+
+
 class GtxnExpr(TxnExpr):
     """An expression that accesses a transaction field from a transaction in the current group."""
 
     def __init__(self, txnIndex: Union[int, Expr], field: TxnField) -> None:
         super().__init__(Op.gtxn, "Gtxn", field)
+        validate_txn_index_or_throw(txnIndex)
         self.txnIndex = txnIndex
 
     def __str__(self):
@@ -51,6 +61,7 @@ def __init__(
         self, txnIndex: Union[int, Expr], field: TxnField, index: Union[int, Expr]
     ) -> None:
         super().__init__(Op.gtxna, Op.gtxnas, "Gtxna", field, index)
+        validate_txn_index_or_throw(txnIndex)
         self.txnIndex = txnIndex
 
     def __str__(self):
diff --git a/pyteal/ast/gtxn_test.py b/pyteal/ast/gtxn_test.py
@@ -2,19 +2,57 @@
 
 from .. import *
 
+teal6Options = CompileOptions(version=6)
 
-def test_gtxn_invalid():
-    with pytest.raises(TealInputError):
-        Gtxn[-1].fee()
-
-    with pytest.raises(TealInputError):
-        Gtxn[MAX_GROUP_SIZE + 1].sender()
 
-    with pytest.raises(TealTypeError):
-        Gtxn[Pop(Int(0))].sender()
-
-    with pytest.raises(TealTypeError):
-        Gtxn[Bytes("index")].sender()
+def test_gtxn_invalid():
+    for f, e in [
+        (lambda: Gtxn[-1], TealInputError),
+        (lambda: Gtxn[MAX_GROUP_SIZE + 1], TealInputError),
+        (lambda: Gtxn[Pop(Int(0))], TealTypeError),
+        (lambda: Gtxn[Bytes("index")], TealTypeError),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gtxn_expr_invalid():
+    for f, e in [
+        (lambda: GtxnExpr(Assert(Int(1)), TxnField.sender), TealTypeError),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gtxn_expr_valid():
+    [
+        e.__teal__(teal6Options)
+        for e in [
+            GtxnExpr(1, TxnField.sender),
+            GtxnExpr(Int(1), TxnField.sender),
+        ]
+    ]
+
+
+def test_gtxna_expr_invalid():
+    for f, e in [
+        (lambda: GtxnaExpr("Invalid_type", TxnField.assets, 1), TealInputError),
+        (lambda: GtxnaExpr(1, TxnField.assets, "Invalid_type"), TealInputError),
+        (lambda: GtxnaExpr(Assert(Int(1)), TxnField.assets, 1), TealTypeError),
+        (lambda: GtxnaExpr(1, TxnField.assets, Assert(Int(1))), TealTypeError),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gtxna_expr_valid():
+    [
+        e.__teal__(teal6Options)
+        for e in [
+            GtxnaExpr(1, TxnField.assets, 1),
+            GtxnaExpr(Int(1), TxnField.assets, Int(1)),
+        ]
+    ]
 
 
 # txn_test.py performs additional testing
diff --git a/pyteal/ast/itxn.py b/pyteal/ast/itxn.py
@@ -120,7 +120,7 @@ def Submit(cls) -> Expr:
         :any:`InnerTxnBuilder.Begin` and :any:`InnerTxnBuilder.SetField` must be called before
         submitting an inner transaction.
 
-        This will fail fail if 16 inner transactions have already been executed, or if the
+        This will fail if 256 inner transactions have already been executed, or if the
         inner transaction itself fails. Upon failure, the current program will immediately exit and
         fail as well.
 
@@ -204,7 +204,8 @@ def SetFields(cls, fields: Dict[TxnField, Union[Expr, List[Expr]]]) -> Expr:
 InnerTxnBuilder.__module__ = "pyteal"
 
 InnerTxn: TxnObject = TxnObject(
-    TxnExprBuilder(Op.itxn, "InnerTxn"), TxnaExprBuilder(Op.itxna, None, "InnerTxna")
+    TxnExprBuilder(Op.itxn, "InnerTxn"),
+    TxnaExprBuilder(Op.itxna, Op.itxnas, "InnerTxna"),
 )
 
 InnerTxn.__module__ = "pyteal"
diff --git a/pyteal/ast/txn.py b/pyteal/ast/txn.py
@@ -161,6 +161,15 @@ def type_of(self):
 class TxnaExpr(LeafExpr):
     """An expression that accesses a transaction array field from the current transaction."""
 
+    @staticmethod
+    def __validate_index_or_throw(index: Union[int, Expr]):
+        if not isinstance(index, (int, Expr)):
+            raise TealInputError(
+                f"Invalid index type:  Expected int or Expr, but received {index}."
+            )
+        if isinstance(index, Expr):
+            require_type(index, TealType.uint64)
+
     def __init__(
         self,
         staticOp: Op,
@@ -172,6 +181,8 @@ def __init__(
         super().__init__()
         if not field.is_array:
             raise TealInputError("Unexpected non-array field: {}".format(field))
+        self.__validate_index_or_throw(index)
+
         self.staticOp = staticOp
         self.dynamicOp = dynamicOp
         self.name = name
diff --git a/pyteal/ast/txn_test.py b/pyteal/ast/txn_test.py
@@ -96,9 +96,9 @@ def test_txn_fields():
             [],
             [TealOp(dynamicGtxnArg, Op.int, 0)],
         ),
-        (InnerTxn, Op.itxn, Op.itxna, None, [], []),
+        (InnerTxn, Op.itxn, Op.itxna, Op.itxnas, [], []),
         *[
-            (Gitxn[i], Op.gitxn, Op.gitxna, None, [i], [])
+            (Gitxn[i], Op.gitxn, Op.gitxna, Op.gitxnas, [i], [])
             for i in range(MAX_GROUP_SIZE)
         ],
     ]
