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/maybe_test.py

@@ -26,6 +26,7 @@ def test_maybe_value():
                     expr = MaybeValue(op, type, immediate_args=iargs, args=args)
 
                     assert expr.slotOk != expr.slotValue
+                    assert expr.output_slots == [expr.slotValue, expr.slotOk]
 
                     assert expr.hasValue().type_of() == TealType.uint64
                     with TealComponent.Context.ignoreExprEquality():

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"