diff --git a/pyteal/__init__.pyi b/pyteal/__init__.pyi index 4431f2116..e4bf54ada 100644 --- a/pyteal/__init__.pyi +++ b/pyteal/__init__.pyi @@ -37,6 +37,7 @@ __all__ = [ "AccountParam", "AccountParamObject", "Add", + "AddW", "Addr", "And", "App", @@ -90,7 +91,8 @@ __all__ = [ "DEFAULT_PROGRAM_VERSION", "DEFAULT_TEAL_VERSION", "Div", - "Divw", + "DivModW", + "DivW", "DynamicScratchVar", "EcdsaCurve", "EcdsaDecompress", @@ -102,6 +104,7 @@ __all__ = [ "Eq", "Err", "Exp", + "ExpW", "Expr", "Extract", "ExtractUint16", @@ -151,6 +154,7 @@ __all__ = [ "Mod", "Mode", "Mul", + "MulW", "MultiValue", "NUM_SLOTS", "NaryExpr", diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py index 77f66da5a..91e233e96 100644 --- a/pyteal/ast/__init__.py +++ b/pyteal/ast/__init__.py @@ -109,7 +109,7 @@ # ternary ops from pyteal.ast.ternaryexpr import ( - Divw, + DivW, Ed25519Verify, Ed25519Verify_Bare, SetBit, @@ -157,6 +157,7 @@ from pyteal.ast.scratchvar import DynamicScratchVar, ScratchVar from pyteal.ast.maybe import MaybeValue from pyteal.ast.multi import MultiValue +from pyteal.ast.wideexpr import AddW, MulW, ExpW, DivModW from pyteal.ast.opup import OpUp, OpUpMode from pyteal.ast.ecdsa import EcdsaCurve, EcdsaVerify, EcdsaDecompress, EcdsaRecover from pyteal.ast.router import ( @@ -240,7 +241,7 @@ "Div", "Mod", "Exp", - "Divw", + "DivW", "BitwiseAnd", "BitwiseOr", "BitwiseXor", @@ -330,4 +331,8 @@ "EcdsaRecover", "JsonRef", "VrfVerify", + "AddW", + "MulW", + "ExpW", + "DivModW", ] diff --git a/pyteal/ast/maybe.py b/pyteal/ast/maybe.py index 43846384d..9d06a593c 100644 --- a/pyteal/ast/maybe.py +++ b/pyteal/ast/maybe.py @@ -16,8 +16,8 @@ def __init__( op: Op, type: TealType, *, - immediate_args: List[Union[int, str]] = None, - args: List[Expr] = None + immediate_args: List[Union[int, str]] | None = None, + args: List[Expr] | None = None ): """Create a new MaybeValue. diff --git a/pyteal/ast/multi.py b/pyteal/ast/multi.py index e61af4b4a..1996386a6 100644 --- a/pyteal/ast/multi.py +++ b/pyteal/ast/multi.py @@ -19,8 +19,8 @@ def __init__( op: Op, types: List[TealType], *, - immediate_args: List[Union[int, str]] = None, - args: List[Expr] = None, + immediate_args: List[Union[int, str]] | None = None, + args: List[Expr] | None = None, compile_check: Callable[["CompileOptions"], None] = lambda _: None, ): """Create a new MultiValue. @@ -41,8 +41,8 @@ def __init__( self.output_slots = [ScratchSlot() for _ in self.types] def outputReducer(self, reducer: Callable[..., Expr]) -> Expr: - input = [slot.load(self.types[i]) for i, slot in enumerate(self.output_slots)] - return Seq(self, reducer(*input)) + inputs = [slot.load(self.types[i]) for i, slot in enumerate(self.output_slots)] + return Seq(self, reducer(*inputs)) def __str__(self): ret_str = "(({}".format(self.op) diff --git a/pyteal/ast/multi_test.py b/pyteal/ast/multi_test.py index a85cebed7..7e511e667 100644 --- a/pyteal/ast/multi_test.py +++ b/pyteal/ast/multi_test.py @@ -6,6 +6,10 @@ options = pt.CompileOptions() +def TOO_MANY_ARGS(arg_len): + return ValueError("Too many arguments. Expected 0-2, got {}".format(arg_len)) + + def __test_single(expr: pt.MultiValue): assert expr.output_slots[0] != expr.output_slots[1] @@ -54,6 +58,8 @@ def __test_single_conditional( arg_2, after_arg_2 = args[1].__teal__(options) after_arg_1.setNextBlock(arg_2) after_arg_2.setNextBlock(expected_call) + else: + raise TOO_MANY_ARGS(len(args)) expected.addIncoming() expected = pt.TealBlock.NormalizeBlocks(expected) @@ -94,6 +100,8 @@ def __test_single_assert(expr: pt.MultiValue, op, args: List[pt.Expr], iargs, re arg_2, after_arg_2 = args[1].__teal__(options) after_arg_1.setNextBlock(arg_2) after_arg_2.setNextBlock(expected_call) + else: + raise TOO_MANY_ARGS(len(args)) expected.addIncoming() expected = pt.TealBlock.NormalizeBlocks(expected) @@ -136,6 +144,8 @@ def __test_single_with_vars( arg_2, after_arg_2 = args[1].__teal__(options) after_arg_1.setNextBlock(arg_2) after_arg_2.setNextBlock(expected_call) + else: + raise TOO_MANY_ARGS(len(args)) expected.addIncoming() expected = pt.TealBlock.NormalizeBlocks(expected) diff --git a/pyteal/ast/ternaryexpr.py b/pyteal/ast/ternaryexpr.py index a8a526cff..3ad219ed1 100644 --- a/pyteal/ast/ternaryexpr.py +++ b/pyteal/ast/ternaryexpr.py @@ -146,7 +146,7 @@ def SetByte(value: Expr, index: Expr, newByteValue: Expr) -> TernaryExpr: ) -def Divw(hi: Expr, lo: Expr, y: Expr) -> TernaryExpr: +def DivW(hi: Expr, lo: Expr, y: Expr) -> TernaryExpr: """ Performs wide division by interpreting `hi` and `lo` as a uint128 value. diff --git a/pyteal/ast/ternaryexpr_test.py b/pyteal/ast/ternaryexpr_test.py index 0530ec6b6..025b8905f 100644 --- a/pyteal/ast/ternaryexpr_test.py +++ b/pyteal/ast/ternaryexpr_test.py @@ -176,7 +176,7 @@ def test_set_byte_invalid(): def test_divw(): args = [pt.Int(0), pt.Int(90), pt.Int(30)] - expr = pt.Divw(args[0], args[1], args[2]) + expr = pt.DivW(args[0], args[1], args[2]) assert expr.type_of() == pt.TealType.uint64 expected = pt.TealSimpleBlock( @@ -197,10 +197,15 @@ def test_divw(): def test_divw_invalid(): with pytest.raises(pt.TealTypeError): - pt.Divw(pt.Bytes("10"), pt.Int(0), pt.Int(1)) + pt.DivW(pt.Bytes("10"), pt.Int(0), pt.Int(1)) with pytest.raises(pt.TealTypeError): - pt.Divw(pt.Int(10), pt.Bytes("0"), pt.Int(1)) + pt.DivW(pt.Int(10), pt.Bytes("0"), pt.Int(1)) with pytest.raises(pt.TealTypeError): - pt.Divw(pt.Int(10), pt.Int(0), pt.Bytes("1")) + pt.DivW(pt.Int(10), pt.Int(0), pt.Bytes("1")) + + +def test_divw_invalid_version(): + with pytest.raises(pt.TealInputError): + pt.DivW(pt.Int(2), pt.Int(2), pt.Int(2)).__teal__(avm5Options) # needs >=6 diff --git a/pyteal/ast/wideexpr.py b/pyteal/ast/wideexpr.py new file mode 100644 index 000000000..7b73d0080 --- /dev/null +++ b/pyteal/ast/wideexpr.py @@ -0,0 +1,123 @@ +from typing import TYPE_CHECKING, List + +from pyteal.ast.expr import Expr +from pyteal.ast.multi import MultiValue +from pyteal.errors import verifyProgramVersion +from pyteal.ir import Op +from pyteal.types import TealType, require_type + +if TYPE_CHECKING: + from pyteal.compiler import CompileOptions + + +class WideExpr(MultiValue): + """Base class for WideInt Operations + + This type of expression produces WideInt(MultiValue). + """ + + def __init__( + self, + op: Op, + args: List[Expr], + ): + """Create a new WideExpr, whose returned type is always a MultiValue of [TealType.uint64, TealType.uint64]. + + Args: + op: The operation that returns values. + args: Stack arguments for the op. + min_version: The minimum TEAL version required to use this expression. + """ + + super().__init__( + op=op, + types=[TealType.uint64, TealType.uint64], + args=args, + immediate_args=None, + ) + + for arg in args: + require_type(arg, TealType.uint64) + + def __teal__(self, options: "CompileOptions"): + verifyProgramVersion( + self.op.min_version, + options.version, + "Program version too low to use op {}".format(self.op), + ) + + return super().__teal__(options) + + +WideExpr.__module__ = "pyteal" + +"""Binary MultiValue operations""" + + +def AddW(adder: Expr, adder_: Expr) -> MultiValue: + """Add two 64-bit integers. + + Produces a MultiValue with two outputs: the sum and the carry-bit. + + Args: + adder: Must evaluate to uint64. + adder_: Must evaluate to uint64. + """ + return WideExpr(Op.addw, [adder, adder_]) + + +def MulW(factor: Expr, factor_: Expr) -> MultiValue: + """Multiply two 64-bit integers. + + Produces a MultiValue with two outputs: the product and the carry-bit. + + Args: + factor: Must evaluate to uint64. + factor_: Must evaluate to uint64. + """ + + return WideExpr(Op.mulw, [factor, factor_]) + + +def ExpW(base: Expr, exponent: Expr) -> MultiValue: + """Raise a 64-bit integer to a power. + + Produces a MultiValue with two outputs: the result and the carry-bit. + + Args: + base: Must evaluate to uint64. + exponent: Must evaluate to uint64. + """ + + return WideExpr(Op.expw, [base, exponent]) + + +def DivModW( + dividendHigh: Expr, dividendLow: Expr, divisorHigh: Expr, divisorLow: Expr +) -> MultiValue: + """Divide two wide-64-bit integers. + + Produces a MultiValue with four outputs: the quotient and its carry-bit, the remainder and its carry-bit. + + Stack: + ..., A: uint64, B: uint64, C: uint64, D: uint64 --> ..., W: uint64, X: uint64, Y: uint64, Z: uint64 + Where W,X = (A,B / C,D); Y,Z = (A,B modulo C,D) + + Example: + All ints should be initialized with Int(). For readability, we didn't use Int() in the example. + DivModW(0, 10, 0, 5) = (0, 2, 0, 0) # 10 / 5 = 2, 10 % 5 = 0 + DivModW(0, 10, 0, 3) = (0, 3, 0, 1) # 10 / 3 = 3, 10 % 3 = 1 + DivModW(5, 14, 0, 5) = (1, 2, 0, 4) # ((5<<64)+14) / 5 = (1<<64)+2, ((5<<64)+14) % 5 = 4 + DivModW(7, 29, 1, 3) = (0, 7, 0, 8) # ((7<<64)+29) / ((1<<64)+3) = 7, ((7<<64)+29) % ((1<<64)+3) = 8 + + Args: + dividendHigh: Must evaluate to uint64. + dividendLow: Must evaluate to uint64. + divisorHigh: Must evaluate to uint64. + divisorLow: Must evaluate to uint64. + """ + + return WideExpr( + Op.divmodw, + [dividendHigh, dividendLow, divisorHigh, divisorLow], + ) diff --git a/pyteal/ast/wideexpr_test.py b/pyteal/ast/wideexpr_test.py new file mode 100644 index 000000000..7acc9571b --- /dev/null +++ b/pyteal/ast/wideexpr_test.py @@ -0,0 +1,161 @@ +import pytest + +import pyteal as pt + +avm2Options = pt.CompileOptions(version=2) +avm3Options = pt.CompileOptions(version=3) +avm4Options = pt.CompileOptions(version=4) +avm5Options = pt.CompileOptions(version=5) + + +def test_addw(): + args = [pt.Int(2), pt.Int(3)] + expr = pt.AddW(pt.Int(1), pt.Int(2)) + + expected = pt.TealSimpleBlock( + [ + pt.TealOp(args[0], pt.Op.int, 1), + pt.TealOp(args[1], pt.Op.int, 2), + pt.TealOp(expr, pt.Op.addw), + pt.TealOp(expr.output_slots[1].store(), pt.Op.store, expr.output_slots[1]), + pt.TealOp(expr.output_slots[0].store(), pt.Op.store, expr.output_slots[0]), + ] + ) + + actual, _ = expr.__teal__(avm2Options) + actual.addIncoming() + actual = pt.TealBlock.NormalizeBlocks(actual) + + with pt.TealComponent.Context.ignoreExprEquality(): + assert actual == expected + + +# TODO: test: test_addw_overload() + + +def test_addw_invalid(): + with pytest.raises(pt.TealTypeError): + pt.AddW(pt.Int(2), pt.Txn.receiver()) + + with pytest.raises(pt.TealTypeError): + pt.AddW(pt.Txn.sender(), pt.Int(2)) + + +def test_mulw(): + args = [pt.Int(3), pt.Int(8)] + expr = pt.MulW(args[0], args[1]) + assert expr.type_of() == pt.TealType.none + + expected = pt.TealSimpleBlock( + [ + pt.TealOp(args[0], pt.Op.int, 3), + pt.TealOp(args[1], pt.Op.int, 8), + pt.TealOp(expr, pt.Op.mulw), + pt.TealOp(expr.output_slots[1].store(), pt.Op.store, expr.output_slots[1]), + pt.TealOp(expr.output_slots[0].store(), pt.Op.store, expr.output_slots[0]), + ] + ) + + actual, _ = expr.__teal__(avm2Options) + actual.addIncoming() + actual = pt.TealBlock.NormalizeBlocks(actual) + + with pt.TealComponent.Context.ignoreExprEquality(): + assert actual == expected + + +# TODO: test: test_mulw_overload() + + +def test_mulw_invalid(): + with pytest.raises(pt.TealTypeError): + pt.MulW(pt.Int(2), pt.Txn.receiver()) + + with pytest.raises(pt.TealTypeError): + pt.MulW(pt.Txn.sender(), pt.Int(2)) + + +# TODO: test: def test_modw_overload(): +# TODO: test: def test_modw_invalid(): + + +def test_expw(): + args = [pt.Int(2), pt.Int(9)] + expr = pt.ExpW(args[0], args[1]) + assert expr.type_of() == pt.TealType.none + + expected = pt.TealSimpleBlock( + [ + pt.TealOp(args[0], pt.Op.int, 2), + pt.TealOp(args[1], pt.Op.int, 9), + pt.TealOp(expr, pt.Op.expw), + pt.TealOp(expr.output_slots[1].store(), pt.Op.store, expr.output_slots[1]), + pt.TealOp(expr.output_slots[0].store(), pt.Op.store, expr.output_slots[0]), + ] + ) + + actual, _ = expr.__teal__(avm4Options) + actual.addIncoming() + actual = pt.TealBlock.NormalizeBlocks(actual) + + with pt.TealComponent.Context.ignoreExprEquality(): + assert actual == expected + + +# TODO: test: def test_expw_overload(): + + +def test_expw_invalid(): + with pytest.raises(pt.TealTypeError): + pt.ExpW(pt.Txn.receiver(), pt.Int(2)) + + with pytest.raises(pt.TealTypeError): + pt.ExpW(pt.Int(2), pt.Txn.sender()) + + +def test_expw_invalid_version(): + with pytest.raises(pt.TealInputError): + pt.ExpW(pt.Int(2), pt.Int(2)).__teal__(avm3Options) # needs >=4 + + +def test_divmodw(): + args = [pt.Int(7), pt.Int(29), pt.Int(1), pt.Int(3)] + expr = pt.DivModW(args[0], args[1], args[2], args[3]) + assert expr.type_of() == pt.TealType.none + + expected = pt.TealSimpleBlock( + [ + pt.TealOp(args[0], pt.Op.int, 7), + pt.TealOp(args[1], pt.Op.int, 29), + pt.TealOp(args[2], pt.Op.int, 1), + pt.TealOp(args[3], pt.Op.int, 3), + pt.TealOp(expr, pt.Op.divmodw), + pt.TealOp(expr.output_slots[1].store(), pt.Op.store, expr.output_slots[1]), + pt.TealOp(expr.output_slots[0].store(), pt.Op.store, expr.output_slots[0]), + ] + ) + + actual, _ = expr.__teal__(avm5Options) + actual.addIncoming() + actual = pt.TealBlock.NormalizeBlocks(actual) + + with pt.TealComponent.Context.ignoreExprEquality(): + assert actual == expected + + +def test_divmodw_invalid(): + with pytest.raises(pt.TealTypeError): + pt.DivModW(pt.Int(2), pt.Txn.receiver(), pt.Int(2), pt.Int(2)) + + with pytest.raises(pt.TealTypeError): + pt.DivModW(pt.Int(2), pt.Int(2), pt.Txn.sender(), pt.Int(2)) + + with pytest.raises(pt.TealTypeError): + pt.DivModW(pt.Int(2), pt.Int(2), pt.Int(2), pt.Txn.sender()) + + +def test_divmodw_invalid_version(): + with pytest.raises(pt.TealInputError): + pt.DivModW(pt.Int(2), pt.Int(2), pt.Int(2), pt.Int(2)).__teal__( + avm3Options + ) # needs >=4