Repository: algorand/pyteal

Issue:
Support importing/exporting values between contracts
# Summary
TEAL v4 introduces the `gload(s)` ops, which can load values from the scratch space of prior app call transactions in the same group. PyTeal should support this feature as well.

# Scope
edit: commenting out the initial idea in favor of a simpler approach inspired by @StylishTriangles's suggestion.

This issue requires the following changes:

1. Allow `ScratchSlot` to have an additional optional parameter `requestedSlotId`, which would be an integer from 0-255 which specifies the slot ID that compiler **must** assign to this slot. This would also require modifications to `compileTeal` to actually assign the correct ID to any `ScratchSlot`s with this parameter, and if multiple unique `ScratchSlot` request the same slot ID, compilation must fail.
2. Also add an optional `requestedSlotId` parameter to `ScratchVar` which would be passed directly to its underlying `ScratchSlot`.
3. Expose the `gload` op in a PyTeal expression (maybe as `ImportScratchValue(txnIndex: Expr, slotId: int)`)

~Ideally we can have a system where one contract can export many values (perhaps with helpful names that don't make it onto the chain) with something like `Export("bidAmount", computedBidAmount)`. Then, another contract can import this exported value with something like `computedBidAmount = ImportFrom(txnIndex, contract1.exports("bidAmount"))`.~

~This would require programs to be compiled in a specific order to make dependencies work, or all at once and the compiler can sort out dependencies.~

~Additionally, the compile method would need to return more information, perhaps a Python object that contains a dict of all exports from the contract and their assigned slots -- it's important to make this information accessible by the user so that they can properly debug, and it also lets other non-PyTeal contracts import their exports. This object then needs to be obtained by the importing contract in order to know which slots to actually load. Again, here it will be important to allow imports from non-PyTeal contracts (meaning essentially direct access to gload(s)).~

Comments:

Patch:
diff --git a/pyteal/__init__.py b/pyteal/__init__.py
index 95dd111..879ecea 100644
--- a/pyteal/__init__.py
+++ b/pyteal/__init__.py
@@ -5,7 +5,7 @@ from .ir import __all__ as ir_all
 from .compiler import MAX_TEAL_VERSION, MIN_TEAL_VERSION, DEFAULT_TEAL_VERSION, CompileOptions, compileTeal
 from .types import TealType
 from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
-from .config import MAX_GROUP_SIZE
+from .config import MAX_GROUP_SIZE, NUM_SLOTS
 
 __all__ = ast_all + ir_all + [
     "MAX_TEAL_VERSION",
@@ -19,4 +19,5 @@ __all__ = ast_all + ir_all + [
     "TealInputError",
     "TealCompileError",
     "MAX_GROUP_SIZE",
+    "NUM_SLOTS",
 ]
diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py
index 08a29e8..ed7e0b5 100644
--- a/pyteal/ast/__init__.py
+++ b/pyteal/ast/__init__.py
@@ -13,6 +13,7 @@ from .arg import Arg
 from .txn import TxnType, TxnField, TxnExpr, TxnaExpr, TxnArray, TxnObject, Txn
 from .gtxn import GtxnExpr, GtxnaExpr, TxnGroup, Gtxn
 from .gaid import GeneratedID
+from .gload import ImportScratchValue
 from .global_ import Global, GlobalField
 from .app import App, AppField, OnComplete
 from .asset import AssetHolding, AssetParam
@@ -66,6 +67,7 @@ __all__ = [
     "TxnGroup",
     "Gtxn",
     "GeneratedID",
+    "ImportScratchValue",
     "Global",
     "GlobalField",
     "App",
diff --git a/pyteal/ast/gload.py b/pyteal/ast/gload.py
new file mode 100644
index 0000000..f8774ab
--- /dev/null
+++ b/pyteal/ast/gload.py
@@ -0,0 +1,58 @@
+from typing import cast, Union, TYPE_CHECKING
+
+from ..types import TealType, require_type
+from ..ir import TealOp, Op, TealBlock
+from ..errors import TealInputError, verifyTealVersion
+from ..config import MAX_GROUP_SIZE, NUM_SLOTS
+from .expr import Expr
+from .int import Int
+from .leafexpr import LeafExpr
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+class ImportScratchValue(LeafExpr):
+    """An expression to load a scratch value created by another transaction in the current group"""
+
+    def __init__(self, txnIndex: Union[int, Expr], slotId: int) -> None:
+        """ Create an expression to load a scratch space slot from a transaction in the current group.
+
+        Requires TEAL version 4 or higher. This operation is only permitted in application mode.
+
+        Args:
+            txnIndex: The index of the transaction from which the created ID should be obtained.
+                This index may be a Python int, or it may be a PyTeal expression that evaluates at
+                runtime. If it's an expression, it must evaluate to a uint64. In all cases, the index
+                must be less than the index of the current transaction.
+            slotId: The index of the scratch slot that should be loaded. The index must be a Python int
+                in the range [0-256). 
+        """
+        super().__init__()
+        if type(txnIndex) == int:
+            if txnIndex < 0 or txnIndex >= MAX_GROUP_SIZE:
+                raise TealInputError("Invalid transaction index {}, shoud be in [0, {})".format(txnIndex, MAX_GROUP_SIZE))
+        else:
+            require_type(cast(Expr, txnIndex).type_of(), TealType.uint64)
+        if slotId < 0 or slotId >= NUM_SLOTS:
+                raise TealInputError("Invalid slot ID {}, shoud be in [0, {})".format(slotId, NUM_SLOTS))
+        
+        self.txnIndex = txnIndex
+        self.slotId = slotId
+
+    def __str__(self) -> str:
+        return "(Gload {} {})".format(self.txnIndex, self.slotId)
+
+    def __teal__(self, options: 'CompileOptions'):
+        verifyTealVersion(Op.gload.min_version, options.version, "TEAL version too low to use Gload expression")
+
+        if type(self.txnIndex) == int:
+            op = TealOp(self, Op.gload, cast(int, self.txnIndex), cast(int, self.slotId))
+            return TealBlock.FromOp(options, op)
+        
+        op = TealOp(self, Op.gloads, self.slotId)
+        return TealBlock.FromOp(options, op, cast(Expr, self.txnIndex))
+
+    def type_of(self):
+        return TealType.anytype
+
+ImportScratchValue.__module__ = "pyteal"
diff --git a/pyteal/ast/scratch.py b/pyteal/ast/scratch.py
index 38de0f5..cd342fc 100644
--- a/pyteal/ast/scratch.py
+++ b/pyteal/ast/scratch.py
@@ -1,6 +1,8 @@
 from typing import TYPE_CHECKING
 
 from ..types import TealType
+from ..config import NUM_SLOTS
+from ..errors import TealInputError
 from .expr import Expr
 
 if TYPE_CHECKING:
@@ -9,12 +11,28 @@ if TYPE_CHECKING:
 class ScratchSlot:
     """Represents the allocation of a scratch space slot."""
 
-    slotId = 0
+    # Unique identifier for the compiler to automatically assign slots
+    # The id field is used by the compiler to map to an actual slot in the source code
+    # Slot ids under 256 are manually reserved slots
+    nextSlotId = NUM_SLOTS 
 
-    def __init__(self):
-        self.id = ScratchSlot.slotId
-        ScratchSlot.slotId += 1
+    def __init__(self, requestedSlotId: int = None):
+        """Initializes a scratch slot with a particular id
 
+        Args:
+            requestedSlotId (optional): A scratch slot id that the compiler must store the value. 
+            This id may be a Python int in the range [0-256).       
+        """
+        if requestedSlotId is None:
+            self.id = ScratchSlot.nextSlotId
+            ScratchSlot.nextSlotId += 1
+            self.isReservedSlot = False
+        else:
+            if requestedSlotId < 0 or requestedSlotId >= NUM_SLOTS:
+                raise TealInputError("Invalid slot ID {}, shoud be in [0, {})".format(requestedSlotId, NUM_SLOTS))
+            self.id = requestedSlotId
+            self.isReservedSlot = True
+        
     def store(self, value: Expr = None) -> Expr:
         """Get an expression to store a value in this slot.
         
diff --git a/pyteal/ast/scratchvar.py b/pyteal/ast/scratchvar.py
index 888bded..53398e1 100644
--- a/pyteal/ast/scratchvar.py
+++ b/pyteal/ast/scratchvar.py
@@ -16,15 +16,17 @@ class ScratchVar:
             ])
     """
 
-    def __init__(self, type: TealType = TealType.anytype):
+    def __init__(self, type: TealType = TealType.anytype, slotId: int = None):
         """Create a new ScratchVar with an optional type.
 
         Args:
             type (optional): The type that this variable can hold. An error will be thrown if an
                 expression with an incompatiable type is stored in this variable. Defaults to
                 TealType.anytype.
+            slotId (optional): A scratch slot id that the compiler must store the value. 
+                This id may be a Python int in the range [0-256).  
         """
-        self.slot = ScratchSlot()
+        self.slot = ScratchSlot(requestedSlotId=slotId)
         self.type = type
     
     def storage_type(self) -> TealType:
diff --git a/pyteal/compiler/compiler.py b/pyteal/compiler/compiler.py
index 37635ee..d57f169 100644
--- a/pyteal/compiler/compiler.py
+++ b/pyteal/compiler/compiler.py
@@ -1,6 +1,6 @@
-from typing import List
+from typing import List, Set
 
-from ..ast import Expr
+from ..ast import Expr, ScratchSlot
 from ..ir import Mode, TealComponent, TealOp, TealBlock
 from ..errors import TealInputError, TealInternalError
 from ..config import NUM_SLOTS
@@ -96,18 +96,32 @@ def compileTeal(ast: Expr, mode: Mode, *, version: int = DEFAULT_TEAL_VERSION, a
     verifyOpsForVersion(teal, version)
     verifyOpsForMode(teal, mode)
 
-    slots = set()
+    slots: Set[ScratchSlot] = set()
+    slotIds: Set[int] = set()
+    nextSlotIndex = 0
     for stmt in teal:
         for slot in stmt.getSlots():
+            # If there are two unique slots with same IDs, raise an error
+            if slot.id in slotIds and id(slot) not in [id(s) for s in slots]:
+                raise TealInternalError("Slot ID {} has been assigned multiple times".format(slot.id))
+            slotIds.add(slot.id)
             slots.add(slot)
     
     if len(slots) > NUM_SLOTS:
         # TODO: identify which slots can be reused
         raise TealInternalError("Too many slots in use: {}, maximum is {}".format(len(slots), NUM_SLOTS))
     
-    for index, slot in enumerate(sorted(slots, key=lambda slot: slot.id)):
+    for slot in sorted(slots, key=lambda slot: slot.id):
+        # Find next vacant slot that compiler can assign to
+        while nextSlotIndex in slotIds:
+            nextSlotIndex += 1
         for stmt in teal:
-            stmt.assignSlot(slot, index)
+            if slot.isReservedSlot:
+                # Slot ids under 256 are manually reserved slots
+                stmt.assignSlot(slot, slot.id)
+            else:
+                stmt.assignSlot(slot, nextSlotIndex)
+                slotIds.add(nextSlotIndex)
     
     if assembleConstants:
         if version < 3:


Test patch:
diff --git a/pyteal/ast/gload_test.py b/pyteal/ast/gload_test.py
new file mode 100644
index 0000000..7d58037
--- /dev/null
+++ b/pyteal/ast/gload_test.py
@@ -0,0 +1,57 @@
+import pytest
+
+from .. import *
+# this is not necessary but mypy complains if it's not included
+from .. import MAX_GROUP_SIZE, NUM_SLOTS, CompileOptions
+
+teal3Options = CompileOptions(version=3)
+teal4Options = CompileOptions(version=4)
+
+def test_gload_teal_3():
+    with pytest.raises(TealInputError):
+        ImportScratchValue(0, 1).__teal__(teal3Options)
+
+    with pytest.raises(TealInputError):
+        ImportScratchValue(Int(0), 1).__teal__(teal3Options)
+
+def test_gload():
+    expr = ImportScratchValue(0, 1)
+    assert expr.type_of() == TealType.anytype
+
+    expected = TealSimpleBlock([
+        TealOp(expr, Op.gload, 0, 1)
+    ])
+    
+    actual, _ = expr.__teal__(teal4Options)
+
+    assert actual == expected
+
+def test_gload_dynamic():
+    arg = Int(1)
+    expr = ImportScratchValue(arg, 0)
+    assert expr.type_of() == TealType.anytype
+
+    expected = TealSimpleBlock([
+        TealOp(arg, Op.int, 1),
+        TealOp(expr, Op.gloads, 0)
+    ])
+    
+    actual, _ = expr.__teal__(teal4Options)
+    actual.addIncoming()
+    actual = TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+def test_gload_invalid():
+    with pytest.raises(TealInputError):
+        ImportScratchValue(-1, 0)
+
+    with pytest.raises(TealInputError):
+        ImportScratchValue(MAX_GROUP_SIZE, 0)
+
+    with pytest.raises(TealInputError):
+        ImportScratchValue(0, -1)
+
+    with pytest.raises(TealInputError):
+        ImportScratchValue(0, NUM_SLOTS)
+    
\ No newline at end of file
diff --git a/pyteal/ast/scratch_test.py b/pyteal/ast/scratch_test.py
index 9236182..dd78053 100644
--- a/pyteal/ast/scratch_test.py
+++ b/pyteal/ast/scratch_test.py
@@ -75,3 +75,23 @@ def test_scratch_stack_store():
     actual, _ = expr.__teal__(options)
 
     assert actual == expected
+
+def test_scratch_assign_id():
+    slot = ScratchSlot(255)
+    expr = ScratchStackStore(slot)
+    assert expr.type_of() == TealType.none
+    
+    expected = TealSimpleBlock([
+        TealOp(expr, Op.store, slot)
+    ])
+
+    actual, _ = expr.__teal__(options)
+
+    assert actual == expected
+
+def test_scratch_assign_id_invalid():
+    with pytest.raises(TealInputError):
+        slot = ScratchSlot(-1)
+
+    with pytest.raises(TealInputError):
+        slot = ScratchSlot(NUM_SLOTS)
diff --git a/pyteal/ast/scratchvar_test.py b/pyteal/ast/scratchvar_test.py
index 519202c..0f9b0c0 100644
--- a/pyteal/ast/scratchvar_test.py
+++ b/pyteal/ast/scratchvar_test.py
@@ -66,3 +66,35 @@ def test_scratchvar_load():
     actual = TealBlock.NormalizeBlocks(actual)
 
     assert actual == expected
+
+def test_scratchvar_assign_store():
+    slotId = 2
+    myvar = ScratchVar(TealType.uint64, slotId)
+    arg = Int(10)
+    expr = myvar.store(arg)
+
+    expected = TealSimpleBlock([
+        TealOp(arg, Op.int, 10),
+        TealOp(expr, Op.store, myvar.slot),
+    ])
+
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+def test_scratchvar_assign_load():
+    slotId = 5
+    myvar = ScratchVar(slotId=slotId)
+    expr = myvar.load()
+
+    expected = TealSimpleBlock([
+        TealOp(expr, Op.load, myvar.slot)
+    ])
+
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
diff --git a/pyteal/compiler/compiler_test.py b/pyteal/compiler/compiler_test.py
index b68eb67..f10e8c9 100644
--- a/pyteal/compiler/compiler_test.py
+++ b/pyteal/compiler/compiler_test.py
@@ -139,6 +139,42 @@ def test_slot_load_before_store():
     with pytest.raises(TealInternalError):
         compileTeal(program, Mode.Application, version=2)
 
+def test_assign_scratch_slots():
+    myScratch       = ScratchVar(TealType.uint64)
+    otherScratch    = ScratchVar(TealType.uint64, 1)
+    anotherScratch  = ScratchVar(TealType.uint64, 0)
+    lastScratch     = ScratchVar(TealType.uint64)
+    prog            = Seq([
+                        myScratch.store(Int(5)),      # Slot 2
+                        otherScratch.store(Int(0)),   # Slot 1
+                        anotherScratch.store(Int(7)), # Slot 0
+                        lastScratch.store(Int(9)),    # Slot 3
+                      ])
+
+    expected = """
+#pragma version 4
+int 5
+store 2
+int 0
+store 1
+int 7
+store 0
+int 9
+store 3
+""".strip()
+    actual = compileTeal(prog, mode=Mode.Signature, version=4) 
+    assert actual == expected
+
+def test_scratchvar_double_assign_invalid():
+    myvar    = ScratchVar(TealType.uint64, 10)
+    otherVar = ScratchVar(TealType.uint64, 10)
+    prog  = Seq([
+                myvar.store(Int(5)),
+                otherVar.store(Int(0))
+            ])
+    with pytest.raises(TealInternalError):
+        compileTeal(prog, mode=Mode.Signature, version=4) 
+
 def test_assembleConstants():
     program = Itob(Int(1) + Int(1) + Tmpl.Int("TMPL_VAR")) == Concat(Bytes("test"), Bytes("test"), Bytes("test2"))