WIP: Wide arithmetic

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",
@@ -218,7 +222,11 @@ __all__ = [
     "VrfVerify",
     "While",
     "WideRatio",
+    "WideUint128",
     "abi",
     "compileTeal",
+    "expW",
     "pragma",
+    "prodW",
+    "sumW",
 ]

pyteal/ast/__init__.py

@@ -109,7 +109,7 @@
 
 # ternary ops
 from pyteal.ast.ternaryexpr import (
-    Divw,
+    DivW,
     Ed25519Verify,
     Ed25519Verify_Bare,
     SetBit,
@@ -124,7 +124,7 @@
 
 # more ops
 from pyteal.ast.naryexpr import NaryExpr, Add, And, Mul, Or, Concat
-from pyteal.ast.widemath import WideRatio
+from pyteal.ast.widemath import WideRatio, sumW, prodW, expW, WideUint128
 
 # control flow
 from pyteal.ast.if_ import If
@@ -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",
@@ -266,6 +267,10 @@
     "Or",
     "Concat",
     "WideRatio",
+    "sumW",
+    "prodW",
+    "expW",
+    "WideUint128",
     "If",
     "Cond",
     "Seq",
@@ -330,4 +335,8 @@
     "EcdsaRecover",
     "JsonRef",
     "VrfVerify",
+    "AddW",
+    "MulW",
+    "ExpW",
+    "DivModW",
 ]

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.
 

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)

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(len(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)

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.
 

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

pyteal/ast/wideexpr.py

@@ -0,0 +1,121 @@
+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)
+
+
+"""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],
+    )