loop (algorand#95)

*Added support for while and for 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,36 @@
+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
+    start, _ = items[1].__teal__(options)
+
+    assert len(options.breakBlocks) == 1
+    assert start in options.breakBlocks

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,35 @@
+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
+    start, _ = items[1].__teal__(options)
+
+    assert len(options.continueBlocks) == 1
+    assert start in options.continueBlocks

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)
+        require_type(start.type_of(), TealType.none)
+        require_type(step.type_of(), TealType.none)
+
+        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)
+
+        options.breakBlocks = breakBlocks
+        options.continueBlocks = continueBlocks
+        options.currentLoop = prevLoop
+
+        branchBlock = TealConditionalBlock([])
+        branchBlock.setTrueBlock(doStart)
+        branchBlock.setFalseBlock(end)
+
+        condEnd.setNextBlock(branchBlock)
+
+        startEnd.setNextBlock(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 doBlock", 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 doBlock", 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"