MultiValue expression implemented to support opcodes that return multiple values

pyteal/__init__.pyi

@@ -121,6 +121,7 @@ __all__ = [
     "ScratchStackStore",
     "ScratchVar",
     "MaybeValue",
+    "MultiValue",
     "BytesAdd",
     "BytesMinus",
     "BytesDiv",

pyteal/ast/__init__.py

@@ -123,6 +123,7 @@
 from .scratch import ScratchSlot, ScratchLoad, ScratchStore, ScratchStackStore
 from .scratchvar import ScratchVar
 from .maybe import MaybeValue
+from .multi import MultiValue
 
 __all__ = [
     "Expr",
@@ -229,6 +230,7 @@
     "ScratchStackStore",
     "ScratchVar",
     "MaybeValue",
+    "MultiValue",
     "BytesAdd",
     "BytesMinus",
     "BytesDiv",

pyteal/ast/maybe.py

@@ -1,16 +1,17 @@
 from typing import List, Union, TYPE_CHECKING
 
+from pyteal.ast.multi import MultiValue
+
 from ..types import TealType
-from ..ir import TealOp, Op, TealBlock
+from ..ir import Op
 from .expr import Expr
-from .leafexpr import LeafExpr
-from .scratch import ScratchSlot, ScratchLoad
+from .scratch import ScratchLoad, ScratchSlot
 
 if TYPE_CHECKING:
     from ..compiler import CompileOptions
 
 
-class MaybeValue(LeafExpr):
+class MaybeValue(MultiValue):
     """Represents a get operation returning a value that may not exist."""
 
     def __init__(
@@ -29,62 +30,42 @@ def __init__(
             immediate_args (optional): Immediate arguments for the op. Defaults to None.
             args (optional): Stack arguments for the op. Defaults to None.
         """
-        super().__init__()
-        self.op = op
-        self.type = type
-        self.immediate_args = immediate_args if immediate_args is not None else []
-        self.args = args if args is not None else []
-        self.slotOk = ScratchSlot()
-        self.slotValue = ScratchSlot()
+        types = [type, TealType.uint64]
+        super().__init__(op, types, immediate_args=immediate_args, args=args)
 
     def hasValue(self) -> ScratchLoad:
         """Check if the value exists.
 
         This will return 1 if the value exists, otherwise 0.
         """
-        return self.slotOk.load(TealType.uint64)
+        return self.output_slots[1].load(self.types[1])
 
     def value(self) -> ScratchLoad:
         """Get the value.
 
         If the value exists, it will be returned. Otherwise, the zero value for this type will be
         returned (i.e. either 0 or an empty byte string, depending on the type).
         """
-        return self.slotValue.load(self.type)
-
-    def __str__(self):
-        ret_str = "(({}".format(self.op)
-        for a in self.immediate_args:
-            ret_str += " " + a.__str__()
-
-        for a in self.args:
-            ret_str += " " + a.__str__()
-        ret_str += ") "
-
-        storeOk = self.slotOk.store()
-        storeValue = self.slotValue.store()
-
-        ret_str += storeOk.__str__() + " " + storeValue.__str__() + ")"
+        return self.output_slots[0].load(self.types[0])
 
-        return ret_str
+    @property
+    def slotOk(self) -> ScratchSlot:
+        """Get the scratch slot that stores hasValue.
 
-    def __teal__(self, options: "CompileOptions"):
-        tealOp = TealOp(self, self.op, *self.immediate_args)
-        callStart, callEnd = TealBlock.FromOp(options, tealOp, *self.args)
-
-        storeOk = self.slotOk.store()
-        storeValue = self.slotValue.store()
-
-        storeOkStart, storeOkEnd = storeOk.__teal__(options)
-        storeValueStart, storeValueEnd = storeValue.__teal__(options)
-
-        callEnd.setNextBlock(storeOkStart)
-        storeOkEnd.setNextBlock(storeValueStart)
+        Note: This is mainly added for backwards compatability and normally shouldn't be used
+        directly in pyteal code.
+        """
+        return self.output_slots[1]
 
-        return callStart, storeValueEnd
+    @property
+    def slotValue(self) -> ScratchSlot:
+        """Get the scratch slot that stores the value or the zero value for the type if the value
+        doesn't exist.
 
-    def type_of(self):
-        return TealType.none
+        Note: This is mainly added for backwards compatability and normally shouldn't be used
+        directly in pyteal code.
+        """
+        return self.output_slots[0]
 
 
 MaybeValue.__module__ = "pyteal"

pyteal/ast/multi.py

@@ -0,0 +1,79 @@
+from typing import Callable, List, Union, TYPE_CHECKING
+
+from ..types import TealType
+from ..ir import TealOp, Op, TealBlock
+from .expr import Expr
+from .leafexpr import LeafExpr
+from .scratch import ScratchSlot
+from .seq import Seq
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+
+class MultiValue(LeafExpr):
+    """Represents an operation that returns more than one value"""
+
+    def __init__(
+        self,
+        op: Op,
+        types: List[TealType],
+        *,
+        immediate_args: List[Union[int, str]] = None,
+        args: List[Expr] = None
+    ):
+        """Create a new MultiValue.
+
+        Args:
+            op: The operation that returns values.
+            types: The types of the returned values.
+            immediate_args (optional): Immediate arguments for the op. Defaults to None.
+            args (optional): Stack arguments for the op. Defaults to None.
+        """
+        super().__init__()
+        self.op = op
+        self.types = types
+        self.immediate_args = immediate_args if immediate_args is not None else []
+        self.args = args if args is not None else []
+
+        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))
+
+    def __str__(self):
+        ret_str = "(({}".format(self.op)
+        for a in self.immediate_args:
+            ret_str += " " + a.__str__()
+
+        for a in self.args:
+            ret_str += " " + a.__str__()
+        ret_str += ") "
+
+        ret_str += " ".join([slot.store().__str__() for slot in self.output_slots])
+        ret_str += ")"
+
+        return ret_str
+
+    def __teal__(self, options: "CompileOptions"):
+        tealOp = TealOp(self, self.op, *self.immediate_args)
+        callStart, callEnd = TealBlock.FromOp(options, tealOp, *self.args)
+
+        curEnd = callEnd
+        # the list is reversed in order to preserve the ordering of the opcode's returned
+        # values. ie the output to stack [A, B, C] should correspond to C->output_slots[2]
+        # B->output_slots[1], and A->output_slots[0].
+        for slot in reversed(self.output_slots):
+            store = slot.store()
+            storeStart, storeEnd = store.__teal__(options)
+            curEnd.setNextBlock(storeStart)
+            curEnd = storeEnd
+
+        return callStart, curEnd
+
+    def type_of(self):
+        return TealType.none
+
+
+MultiValue.__module__ = "pyteal"

pyteal/ast/multi_test.py

@@ -0,0 +1,191 @@
+import pytest
+
+from .. import *
+from typing import List
+
+# this is not necessary but mypy complains if it's not included
+from .. import CompileOptions
+
+options = CompileOptions()
+
+
+def __test_single(expr: MultiValue):
+    assert expr.output_slots[0] != expr.output_slots[1]
+
+    with TealComponent.Context.ignoreExprEquality():
+        assert expr.output_slots[0].load().__teal__(options) == ScratchLoad(
+            expr.output_slots[0]
+        ).__teal__(options)
+
+    with TealComponent.Context.ignoreExprEquality():
+        assert expr.output_slots[1].load().__teal__(options) == ScratchLoad(
+            expr.output_slots[1]
+        ).__teal__(options)
+
+    assert expr.type_of() == TealType.none
+
+
+def __test_single_conditional(expr: MultiValue, op, args: List[Expr], iargs, reducer):
+    __test_single(expr)
+
+    expected_call = TealSimpleBlock(
+        [
+            TealOp(expr, op, *iargs),
+            TealOp(expr.output_slots[1].store(), Op.store, expr.output_slots[1]),
+            TealOp(expr.output_slots[0].store(), Op.store, expr.output_slots[0]),
+        ]
+    )
+
+    ifExpr = (
+        If(expr.output_slots[1].load())
+        .Then(expr.output_slots[0].load())
+        .Else(Bytes("None"))
+    )
+    ifBlockStart, _ = ifExpr.__teal__(options)
+
+    expected_call.setNextBlock(ifBlockStart)
+
+    if len(args) == 0:
+        expected: TealBlock = expected_call
+    elif len(args) == 1:
+        expected, after_arg = args[0].__teal__(options)
+        after_arg.setNextBlock(expected_call)
+    elif len(args) == 2:
+        expected, after_arg_1 = args[0].__teal__(options)
+        arg_2, after_arg_2 = args[1].__teal__(options)
+        after_arg_1.setNextBlock(arg_2)
+        after_arg_2.setNextBlock(expected_call)
+
+    expected.addIncoming()
+    expected = TealBlock.NormalizeBlocks(expected)
+
+    actual, _ = expr.outputReducer(reducer).__teal__(options)
+    actual.addIncoming()
+    actual = TealBlock.NormalizeBlocks(actual)
+
+    with TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
+def __test_single_assert(expr: MultiValue, op, args: List[Expr], iargs, reducer):
+    __test_single(expr)
+
+    expected_call = TealSimpleBlock(
+        [
+            TealOp(expr, op, *iargs),
+            TealOp(expr.output_slots[1].store(), Op.store, expr.output_slots[1]),
+            TealOp(expr.output_slots[0].store(), Op.store, expr.output_slots[0]),
+        ]
+    )
+
+    assertExpr = Seq(Assert(expr.output_slots[1].load()), expr.output_slots[0].load())
+    assertBlockStart, _ = assertExpr.__teal__(options)
+
+    expected_call.setNextBlock(assertBlockStart)
+
+    if len(args) == 0:
+        expected: TealBlock = expected_call
+    elif len(args) == 1:
+        expected, after_arg = args[0].__teal__(options)
+        after_arg.setNextBlock(expected_call)
+    elif len(args) == 2:
+        expected, after_arg_1 = args[0].__teal__(options)
+        arg_2, after_arg_2 = args[1].__teal__(options)
+        after_arg_1.setNextBlock(arg_2)
+        after_arg_2.setNextBlock(expected_call)
+
+    expected.addIncoming()
+    expected = TealBlock.NormalizeBlocks(expected)
+
+    actual, _ = expr.outputReducer(reducer).__teal__(options)
+    actual.addIncoming()
+    actual = TealBlock.NormalizeBlocks(actual)
+
+    with TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
+def __test_single_with_vars(
+    expr: MultiValue, op, args: List[Expr], iargs, var1, var2, reducer
+):
+    __test_single(expr)
+
+    expected_call = TealSimpleBlock(
+        [
+            TealOp(expr, op, *iargs),
+            TealOp(expr.output_slots[1].store(), Op.store, expr.output_slots[1]),
+            TealOp(expr.output_slots[0].store(), Op.store, expr.output_slots[0]),
+        ]
+    )
+
+    varExpr = Seq(
+        var1.store(expr.output_slots[1].load()), var2.store(expr.output_slots[0].load())
+    )
+    varBlockStart, _ = varExpr.__teal__(options)
+
+    expected_call.setNextBlock(varBlockStart)
+
+    if len(args) == 0:
+        expected: TealBlock = expected_call
+    elif len(args) == 1:
+        expected, after_arg = args[0].__teal__(options)
+        after_arg.setNextBlock(expected_call)
+    elif len(args) == 2:
+        expected, after_arg_1 = args[0].__teal__(options)
+        arg_2, after_arg_2 = args[1].__teal__(options)
+        after_arg_1.setNextBlock(arg_2)
+        after_arg_2.setNextBlock(expected_call)
+
+    expected.addIncoming()
+    expected = TealBlock.NormalizeBlocks(expected)
+
+    actual, _ = expr.outputReducer(reducer).__teal__(options)
+    actual.addIncoming()
+    actual = TealBlock.NormalizeBlocks(actual)
+
+    with TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
+def test_multi_value():
+    ops = (
+        Op.app_global_get_ex,
+        Op.app_local_get_ex,
+        Op.asset_holding_get,
+        Op.asset_params_get,
+    )
+    types = (TealType.uint64, TealType.bytes, TealType.anytype)
+    immedate_argv = ([], ["AssetFrozen"])
+    argv = ([], [Int(0)], [Int(1), Int(2)])
+
+    for op in ops:
+        for type in types:
+            for iargs in immedate_argv:
+                for args in argv:
+                    reducer = (
+                        lambda value, hasValue: If(hasValue)
+                        .Then(value)
+                        .Else(Bytes("None"))
+                    )
+                    expr = MultiValue(
+                        op, [type, TealType.uint64], immediate_args=iargs, args=args
+                    )
+                    __test_single_conditional(expr, op, args, iargs, reducer)
+
+                    reducer = lambda value, hasValue: Seq(Assert(hasValue), value)
+                    expr = MultiValue(
+                        op, [type, TealType.uint64], immediate_args=iargs, args=args
+                    )
+                    __test_single_assert(expr, op, args, iargs, reducer)
+
+                    hasValueVar = ScratchVar(TealType.uint64)
+                    valueVar = ScratchVar(type)
+                    reducer = lambda value, hasValue: Seq(
+                        hasValueVar.store(hasValue), valueVar.store(value)
+                    )
+                    expr = MultiValue(
+                        op, [type, TealType.uint64], immediate_args=iargs, args=args
+                    )
+                    __test_single_with_vars(
+                        expr, op, args, iargs, hasValueVar, valueVar, reducer
+                    )