Repository: algorand/pyteal (language: python)

Issue:
Support for wide (128-bit) arithmetic
## Summary
TEAL has the opcodes `mulw`, `addw`, and `divmodw`, which can be used to multiply, add, and divide 64-bit unsigned integers without worrying about possible overflow. They work by using two 64-bit integers to represent a single 128-bit integer on the stack.

PyTeal could expose these opcodes in order to create convenience methods for math operations whose input and output are 64-bit integers, but whose intermediate values might be as large as a 128-bit integer. (By this I mean we could make it easier to express `A * B / C` where `A * B` might overflow a 64-bit int, but the final expression will not overflow. If you're looking for ways to directly handle integers larger than 64-bits, consider [byteslice arithmetic](https://pyteal.readthedocs.io/en/latest/arithmetic_expression.html#byteslice-arithmetic) instead.)

## Scope
Create a set of functions which internally use `mulw`, `addw`, and `divmodw` to perform multi-step arithmetic expressions on 64-bit unsigned integers. These could be:
* `WideRatio(numeratorFactors: List[Expr], denominatorFactors: List[Expr]) -> Expr`: calculate the product of `numeratorFactors` divided by the product of `denominatorFactors`, e.g. `(N_1 * N_2 * N_3 * ...) / (D_1 * D_2 * D_3 * ...)`
* `WideSummation(positiveTerms: List[Expr], negativeTerms: List[Expr]) -> Expr`: calculate the summation of all `positiveTerms` minus the summation of all `negativeTerms`, e.g. `(P_1 + P_2 + P_3 + ...) - (N_1 + N_2 + N_3 + ...)`

Maybe there's a need for expressions mixing addition and division too, possibly of the form `(A + B) / C`?

## Priority

Not urgent because it's possible to use byteslice arithmetic in place of this (e.g. `Btoi(BytesDiv(BytesMul(Itob(Int(1)), Itob(Int(2))), Itob(Int(3))))` expresses `1 * 2 / 3` in a way that's safe from overflow). However, exposing support for `divmodw` and friends would be more efficient in terms of opcode cost.



Discussion / hints:
bumping this; @jasonpaulos, as @cusma says, judging from discord discussions a fair amount of projects would stand to benefit from an official, safe and efficient library leveraging the 128-bit math opcodes

Using byteslice is wasteful

Fix patch:
diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py
index 157623b99..1983924db 100644
--- a/pyteal/ast/__init__.py
+++ b/pyteal/ast/__init__.py
@@ -92,6 +92,7 @@
 
 # more ops
 from .naryexpr import NaryExpr, And, Or, Concat
+from .safemath import SafeRatio
 
 # control flow
 from .if_ import If
@@ -194,6 +195,7 @@
     "And",
     "Or",
     "Concat",
+    "SafeRatio",
     "If",
     "Cond",
     "Seq",
diff --git a/pyteal/ast/safemath.py b/pyteal/ast/safemath.py
new file mode 100644
index 000000000..c08544016
--- /dev/null
+++ b/pyteal/ast/safemath.py
@@ -0,0 +1,159 @@
+from typing import List, Tuple, TYPE_CHECKING
+
+from ..types import TealType, require_type
+from ..errors import TealInputError, TealInternalError, TealCompileError
+from ..ir import TealOp, Op, TealSimpleBlock, TealBlock
+from .expr import Expr
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+
+def multiplyFactors(
+    expr: Expr, factors: List[Expr], options: "CompileOptions"
+) -> Tuple[TealSimpleBlock, TealSimpleBlock]:
+    if len(factors) == 0:
+        raise TealInternalError("Received 0 factors")
+
+    start = TealSimpleBlock([])
+
+    fac0Start, fac0End = factors[0].__teal__(options)
+
+    if len(factors) == 1:
+        # need to use 0 as high word
+        highword = TealSimpleBlock([TealOp(expr, Op.int, 0)])
+
+        start.setNextBlock(highword)
+        highword.setNextBlock(fac0Start)
+
+        end = fac0End
+    else:
+        start.setNextBlock(fac0Start)
+
+        fac1Start, fac1End = factors[1].__teal__(options)
+        fac0End.setNextBlock(fac1Start)
+
+        multiplyFirst2 = TealSimpleBlock([TealOp(expr, Op.mulw)])
+        fac1End.setNextBlock(multiplyFirst2)
+
+        end = multiplyFirst2
+        for factor in factors[2:]:
+            facXStart, facXEnd = factor.__teal__(options)
+            end.setNextBlock(facXStart)
+
+            # stack is [..., A, B, C], where C is current factor
+            # need to pop all A,B,C from stack and push X,Y, where X and Y are:
+            #       X * 2**64 + Y = (A * 2**64 + B) * C
+            # <=>   X * 2**64 + Y = A * C * 2**64 + B * C
+            # <=>   X = A * C + highword(B * C)
+            #       Y = lowword(B * C)
+            multiply = TealSimpleBlock(
+                [
+                    TealOp(expr, Op.uncover, 2),  # stack: [..., B, C, A]
+                    TealOp(expr, Op.dig, 1),  # stack: [..., B, C, A, C]
+                    TealOp(expr, Op.mul),  # stack: [..., B, C, A*C]
+                    TealOp(expr, Op.cover, 2),  # stack: [..., A*C, B, C]
+                    TealOp(
+                        expr, Op.mulw
+                    ),  # stack: [..., A*C, highword(B*C), lowword(B*C)]
+                    TealOp(
+                        expr, Op.cover, 2
+                    ),  # stack: [..., lowword(B*C), A*C, highword(B*C)]
+                    TealOp(
+                        expr, Op.add
+                    ),  # stack: [..., lowword(B*C), A*C+highword(B*C)]
+                    TealOp(
+                        expr, Op.swap
+                    ),  # stack: [..., A*C+highword(B*C), lowword(B*C)]
+                ]
+            )
+
+            facXEnd.setNextBlock(multiply)
+            end = multiply
+
+    return start, end
+
+
+class SafeRatio(Expr):
+    """A class used to calculate expressions of the form :code:`(N_1 * N_2 * N_3 * ...) / (D_1 * D_2 * D_3 * ...)`
+
+    Use this class if all inputs to the expression are uint64s, the output fits in a uint64, and all
+    intermediate values fit in a uint128.
+    """
+
+    def __init__(
+        self, numeratorFactors: List[Expr], denominatorFactors: List[Expr]
+    ) -> None:
+        """Create a new SafeRatio expression with the given numerator and denominator factors.
+
+        This will calculate :code:`(N_1 * N_2 * N_3 * ...) / (D_1 * D_2 * D_3 * ...)`, where each
+        :code:`N_i` represents an element in :code:`numeratorFactors` and each :code:`D_i`
+        represents an element in :code:`denominatorFactors`.
+
+        Requires TEAL version 5 or higher.
+
+        Args:
+            numeratorFactors: The factors in the numerator of the ratio. This list must have at
+                least 1 element. If this list has exactly 1 element, then denominatorFactors must
+                have more than 1 element (otherwise basic division should be used).
+            denominatorFactors: The factors in the denominator of the ratio. This list must have at
+                least 1 element.
+        """
+        super().__init__()
+        if len(numeratorFactors) == 0 or len(denominatorFactors) == 0:
+            raise TealInternalError(
+                "At least 1 factor must be present in the numerator and denominator"
+            )
+        if len(numeratorFactors) == 1 and len(denominatorFactors) == 1:
+            raise TealInternalError(
+                "There is only a single factor in the numerator and denominator. Use basic division instead."
+            )
+        self.numeratorFactors = numeratorFactors
+        self.denominatorFactors = denominatorFactors
+
+    def __teal__(self, options: "CompileOptions"):
+        if options.version < Op.cover.min_version:
+            raise TealCompileError(
+                "SafeRatio requires TEAL version {} or higher".format(
+                    Op.cover.min_version
+                ),
+                self,
+            )
+
+        numStart, numEnd = multiplyFactors(self, self.numeratorFactors, options)
+        denomStart, denomEnd = multiplyFactors(self, self.denominatorFactors, options)
+        numEnd.setNextBlock(denomStart)
+
+        combine = TealSimpleBlock(
+            [
+                TealOp(self, Op.divmodw),
+                TealOp(self, Op.pop),  # pop remainder low word
+                TealOp(self, Op.pop),  # pop remainder high word
+                TealOp(self, Op.swap),  # swap quotient high and low words
+                TealOp(self, Op.logic_not),
+                TealOp(self, Op.assert_),  # assert quotient high word is 0
+                # end with quotient low word remaining on the stack
+            ]
+        )
+        denomEnd.setNextBlock(combine)
+
+        return numStart, combine
+
+    def __str__(self):
+        ret_str = "(SafeRatio (*"
+        for f in self.numeratorFactors:
+            ret_str += " " + str(f)
+        ret_str += ") (*"
+        for f in self.denominatorFactors:
+            ret_str += " " + str(f)
+        ret_str += ")"
+        return ret_str
+
+    def type_of(self):
+        return TealType.uint64
+
+    def has_return(self):
+        return False
+
+
+SafeRatio.__module__ = "pyteal"


Test patch:
diff --git a/pyteal/compiler/compiler_test.py b/pyteal/compiler/compiler_test.py
index a9a56975e..ce39681fc 100644
--- a/pyteal/compiler/compiler_test.py
+++ b/pyteal/compiler/compiler_test.py
@@ -1337,3 +1337,157 @@ def storeValue(key: Expr, t1: Expr, t2: Expr, t3: Expr) -> Expr:
     """.strip()
     actual = compileTeal(program, Mode.Application, version=4, assembleConstants=True)
     assert actual == expected
+
+
+def test_compile_safe_ratio():
+    cases = (
+        (
+            SafeRatio([Int(2), Int(100)], [Int(5)]),
+            """#pragma version 5
+int 2
+int 100
+mulw
+int 0
+int 5
+divmodw
+pop
+pop
+swap
+!
+assert
+return
+""",
+        ),
+        (
+            SafeRatio([Int(2), Int(100)], [Int(10), Int(5)]),
+            """#pragma version 5
+int 2
+int 100
+mulw
+int 10
+int 5
+mulw
+divmodw
+pop
+pop
+swap
+!
+assert
+return
+""",
+        ),
+        (
+            SafeRatio([Int(2), Int(100), Int(3)], [Int(10), Int(5)]),
+            """#pragma version 5
+int 2
+int 100
+mulw
+int 3
+uncover 2
+dig 1
+*
+cover 2
+mulw
+cover 2
++
+swap
+int 10
+int 5
+mulw
+divmodw
+pop
+pop
+swap
+!
+assert
+return
+""",
+        ),
+        (
+            SafeRatio([Int(2), Int(100), Int(3)], [Int(10), Int(5), Int(6)]),
+            """#pragma version 5
+int 2
+int 100
+mulw
+int 3
+uncover 2
+dig 1
+*
+cover 2
+mulw
+cover 2
++
+swap
+int 10
+int 5
+mulw
+int 6
+uncover 2
+dig 1
+*
+cover 2
+mulw
+cover 2
++
+swap
+divmodw
+pop
+pop
+swap
+!
+assert
+return
+""",
+        ),
+        (
+            SafeRatio([Int(2), Int(100), Int(3), Int(4)], [Int(10), Int(5), Int(6)]),
+            """#pragma version 5
+int 2
+int 100
+mulw
+int 3
+uncover 2
+dig 1
+*
+cover 2
+mulw
+cover 2
++
+swap
+int 4
+uncover 2
+dig 1
+*
+cover 2
+mulw
+cover 2
++
+swap
+int 10
+int 5
+mulw
+int 6
+uncover 2
+dig 1
+*
+cover 2
+mulw
+cover 2
++
+swap
+divmodw
+pop
+pop
+swap
+!
+assert
+return
+""",
+        ),
+    )
+
+    for program, expected in cases:
+        actual = compileTeal(
+            program, Mode.Application, version=5, assembleConstants=False
+        )
+        assert actual == expected.strip()