loop

pyteal/ast/__init__.py

@@ -93,6 +93,11 @@
 from .cond import Cond
 from .seq import Seq
 from .assert_ import Assert
+from .while_ import While
+from .for_ import For
+from .break_ import Break
+from .continue_ import Continue
+
 
 # misc
 from .scratch import ScratchSlot, ScratchLoad, ScratchStore, ScratchStackStore
@@ -200,4 +205,8 @@
     "BytesGe",
     "BytesNot",
     "BytesZero",
+    "While",
+    "For",
+    "Break",
+    "Continue",
 ]

pyteal/ast/break_.py

@@ -0,0 +1,40 @@
+from typing import TYPE_CHECKING
+
+from ..types import TealType
+from ..errors import TealCompileError
+from .expr import Expr
+from ..ir import TealSimpleBlock
+
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+
+class Break(Expr):
+    """A break expression"""
+
+    def __init__(self) -> None:
+        """Create a new break expression.
+
+        This operation is only permitted in a loop.
+
+        """
+        super().__init__()
+
+    def __str__(self) -> str:
+        return "break"
+
+    def __teal__(self, options: "CompileOptions"):
+        if options.currentLoop is None:
+            raise TealCompileError("break is only allowed in a loop", self)
+
+        start = TealSimpleBlock([])
+        options.breakBlocks.append(start)
+
+        return start, start
+
+    def type_of(self):
+        return TealType.none
+
+
+Break.__module__ = "pyteal"

pyteal/ast/break_test.py

@@ -0,0 +1,42 @@
+import pytest
+
+from .. import *
+
+# this is not necessary but mypy complains if it's not included
+from .. import CompileOptions
+
+options = CompileOptions()
+
+
+def test_break_fail():
+
+    with pytest.raises(TealCompileError):
+        Break().__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        If(Int(1), Break()).__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        Seq([Break()]).__teal__(options)
+
+    with pytest.raises(TypeError):
+        Break(Int(1))
+
+
+def test_break():
+
+    items = [Int(1), Seq([Break()])]
+    expr = While(items[0]).Do(items[1])
+    actual, _ = expr.__teal__(options)
+
+    options.currentLoop = expr
+    expected, condEnd = items[0].__teal__(options)
+    do, doEnd = items[1].__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    end = TealSimpleBlock([])
+    expectedBranch.setTrueBlock(do)
+    expectedBranch.setFalseBlock(end)
+    condEnd.setNextBlock(expectedBranch)
+    doEnd.setNextBlock(expected)
+
+    actual, _ = expr.__teal__(options)

pyteal/ast/continue_.py

@@ -0,0 +1,40 @@
+from typing import TYPE_CHECKING
+
+from ..types import TealType
+from ..errors import TealCompileError
+from .expr import Expr
+from ..ir import TealSimpleBlock
+
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+
+class Continue(Expr):
+    """A continue expression"""
+
+    def __init__(self) -> None:
+        """Create a new continue expression.
+
+        This operation is only permitted in a loop.
+
+        """
+        super().__init__()
+
+    def __str__(self) -> str:
+        return "continue"
+
+    def __teal__(self, options: "CompileOptions"):
+        if options.currentLoop is None:
+            raise TealCompileError("continue is only allowed in a loop", self)
+
+        start = TealSimpleBlock([])
+        options.continueBlocks.append(start)
+
+        return start, start
+
+    def type_of(self):
+        return TealType.none
+
+
+Continue.__module__ = "pyteal"

pyteal/ast/continue_test.py

@@ -0,0 +1,41 @@
+import pytest
+
+from .. import *
+
+# this is not necessary but mypy complains if it's not included
+from .. import CompileOptions
+
+options = CompileOptions()
+
+
+def test_continue_fail():
+    with pytest.raises(TealCompileError):
+        Continue().__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        If(Int(1), Continue()).__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        Seq([Continue()]).__teal__(options)
+
+    with pytest.raises(TypeError):
+        Continue(Int(1))
+
+
+def test_continue():
+
+    items = [Int(1), Seq([Continue()])]
+    expr = While(items[0]).Do(items[1])
+    actual, _ = expr.__teal__(options)
+
+    options.currentLoop = expr
+    expected, condEnd = items[0].__teal__(options)
+    do, doEnd = items[1].__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    end = TealSimpleBlock([])
+    expectedBranch.setTrueBlock(do)
+    expectedBranch.setFalseBlock(end)
+    condEnd.setNextBlock(expectedBranch)
+    doEnd.setNextBlock(expected)
+
+    actual, _ = expr.__teal__(options)

pyteal/ast/for_.py

@@ -0,0 +1,107 @@
+from typing import TYPE_CHECKING, Optional
+
+from ..types import TealType, require_type
+from ..ir import TealSimpleBlock, TealConditionalBlock
+from ..errors import TealCompileError
+from .expr import Expr
+from .seq import Seq
+from .int import Int
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+
+class For(Expr):
+    """For expression."""
+
+    def __init__(self, start: Expr, cond: Expr, step: Expr) -> None:
+        """Create a new For expression.
+
+        When this For expression is executed, the condition will be evaluated, and if it produces a
+        true value, doBlock will be executed and return to the start of the expression execution.
+        Otherwise, no branch will be executed.
+
+        Args:
+            start: Expression setting the variable's initial value
+            cond: The condition to check. Must evaluate to uint64.
+            step: Expression to update the variable's value.
+        """
+        super().__init__()
+        require_type(cond.type_of(), TealType.uint64)
+
+        self.start = start
+        self.cond = cond
+        self.step = step
+        self.doBlock: Optional[Expr] = None
+
+    def __teal__(self, options: "CompileOptions"):
+        if self.doBlock is None:
+            raise TealCompileError("For expression must have a doBlock", self)
+
+        breakBlocks = options.breakBlocks
+        continueBlocks = options.continueBlocks
+        prevLoop = options.currentLoop
+
+        options.breakBlocks = []
+        options.continueBlocks = []
+        options.currentLoop = self
+
+        end = TealSimpleBlock([])
+        start, startEnd = self.start.__teal__(options)
+        condStart, condEnd = self.cond.__teal__(options)
+        doStart, doEnd = self.doBlock.__teal__(options)
+
+        stepStart, stepEnd = self.step.__teal__(options)
+        stepEnd.setNextBlock(condStart)
+        doEnd.setNextBlock(stepStart)
+
+        for block in options.breakBlocks:
+            block.setNextBlock(end)
+
+        for block in options.continueBlocks:
+            block.setNextBlock(stepStart)
+            doEnd.setNextBlock(block)
+
+        options.breakBlocks = breakBlocks
+        options.continueBlocks = continueBlocks
+        options.currentLoop = prevLoop
+
+        branchBlock = TealConditionalBlock([])
+        branchBlock.setTrueBlock(doStart)
+        branchBlock.setFalseBlock(end)
+
+        condEnd.setNextBlock(branchBlock)
+        condStart.addIncoming(doStart)
+
+        startEnd.nextBlock = condStart
+
+        return start, end
+
+    def __str__(self):
+        if self.start is None:
+            raise TealCompileError("For expression must have a start", self)
+        if self.cond is None:
+            raise TealCompileError("For expression must have a condition", self)
+        if self.step is None:
+            raise TealCompileError("For expression must have a end", self)
+        if self.doBlock is None:
+            raise TealCompileError("For expression must have a thenBranch", self)
+
+        return "(For {} {} {} {})".format(
+            self.start, self.cond, self.step, self.doBlock
+        )
+
+    def type_of(self):
+        if self.doBlock is None:
+            raise TealCompileError("For expression must have a thenBranch", self)
+        return TealType.none
+
+    def Do(self, doBlock: Expr):
+        if self.doBlock is not None:
+            raise TealCompileError("For expression already has a doBlock", self)
+        require_type(doBlock.type_of(), TealType.none)
+        self.doBlock = doBlock
+        return self
+
+
+For.__module__ = "pyteal"

pyteal/ast/for_test.py

@@ -0,0 +1,94 @@
+import pytest
+
+from .. import *
+
+# this is not necessary but mypy complains if it's not included
+from .. import CompileOptions
+
+options = CompileOptions()
+
+
+def test_for_compiles():
+    expr = For(Int(0), Int(1), Int(2)).Do(App.globalPut(Itob(Int(0)), Itob(Int(2))))
+    assert expr.type_of() == TealType.none
+    expr.__teal__(options)
+
+    expr = For(Int(0), Int(1), Int(2)).Do(App.globalPut(Itob(Int(0)), Itob(Int(2))))
+    assert expr.type_of() == TealType.none
+    expr.__teal__(options)
+
+
+def test_nested_for_compiles():
+    i = ScratchVar()
+    expr = For(Int(0), Int(1), Int(2)).Do(
+        Seq([For(Int(0), Int(1), Int(2)).Do(Seq([i.store(Int(0))]))])
+    )
+    assert expr.type_of() == TealType.none
+
+
+def test_continue_break():
+    expr = For(Int(0), Int(1), Int(2)).Do(Seq([If(Int(1), Break(), Continue())]))
+    assert expr.type_of() == TealType.none
+    expr.__teal__(options)
+
+
+def test_for():
+    i = ScratchVar()
+    items = [
+        (i.store(Int(0))),
+        i.load() < Int(10),
+        i.store(i.load() + Int(1)),
+        App.globalPut(Itob(i.load()), i.load() * Int(2)),
+    ]
+    expr = For(items[0], items[1], items[2]).Do(Seq([items[3]]))
+
+    assert expr.type_of() == TealType.none
+
+    expected, varEnd = items[0].__teal__(options)
+    condStart, condEnd = items[1].__teal__(options)
+    stepStart, stepEnd = items[2].__teal__(options)
+    do, doEnd = Seq([items[3]]).__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    end = TealSimpleBlock([])
+
+    varEnd.setNextBlock(condStart)
+    doEnd.setNextBlock(stepStart)
+
+    expectedBranch.setTrueBlock(do)
+    expectedBranch.setFalseBlock(end)
+    condEnd.setNextBlock(expectedBranch)
+    stepEnd.setNextBlock(condStart)
+
+    actual, _ = expr.__teal__(options)
+
+    assert actual == expected
+
+
+def test_invalid_for():
+    with pytest.raises(TypeError):
+        expr = For()
+
+    with pytest.raises(TypeError):
+        expr = For(Int(2))
+
+    with pytest.raises(TypeError):
+        expr = For(Int(1), Int(2))
+
+    with pytest.raises(TealCompileError):
+        expr = For(Int(0), Int(1), Int(2))
+        expr.__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        expr = For(Int(0), Int(1), Int(2))
+        expr.type_of()
+
+    with pytest.raises(TealCompileError):
+        expr = For(Int(0), Int(1), Int(2))
+        expr.__str__()
+
+    with pytest.raises(TealTypeError):
+        expr = For(Int(0), Int(1), Int(2)).Do(Int(0))
+
+    with pytest.raises(TealCompileError):
+        expr = For(Int(0), Int(1), Int(2)).Do(Continue()).Do(Continue())
+        expr.__str__()