diff --git a/pyteal/ast/gitxn.py b/pyteal/ast/gitxn.py index 168899945..4108d86ff 100644 --- 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 index 677cfc09d..ab950067b 100644 --- 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 index c4b0c14f7..b12f8f9d5 100644 --- 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 index 168090918..9b895d100 100644 --- 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 index 6736c7113..ff857c808 100644 --- 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 index c1a6f9a14..abb7d3c52 100644 --- 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 index efe81ce08..458003f7d 100644 --- 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) ], ]