diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py index f74a576a8..7b78e6797 100644 --- a/pyteal/ast/__init__.py +++ b/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", ] diff --git a/pyteal/ast/break_.py b/pyteal/ast/break_.py new file mode 100644 index 000000000..812b0a5a7 --- /dev/null +++ b/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" diff --git a/pyteal/ast/break_test.py b/pyteal/ast/break_test.py new file mode 100644 index 000000000..8be9dbb9a --- /dev/null +++ b/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 diff --git a/pyteal/ast/continue_.py b/pyteal/ast/continue_.py new file mode 100644 index 000000000..5ad6842c3 --- /dev/null +++ b/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" diff --git a/pyteal/ast/continue_test.py b/pyteal/ast/continue_test.py new file mode 100644 index 000000000..5dd87a4ac --- /dev/null +++ b/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 diff --git a/pyteal/ast/for_.py b/pyteal/ast/for_.py new file mode 100644 index 000000000..1b0a01c17 --- /dev/null +++ b/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" diff --git a/pyteal/ast/for_test.py b/pyteal/ast/for_test.py new file mode 100644 index 000000000..9a3a0f964 --- /dev/null +++ b/pyteal/ast/for_test.py @@ -0,0 +1,187 @@ +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(): + i = ScratchVar() + + expr = For(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))).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(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))).Do( + Seq( + [ + For(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))).Do( + Seq([i.store(Int(0))]) + ) + ] + ) + ) + assert expr.type_of() == TealType.none + + +def test_continue_break(): + i = ScratchVar() + expr = For(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))).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_for_continue(): + i = ScratchVar() + items = [ + (i.store(Int(0))), + i.load() < Int(10), + i.store(i.load() + Int(1)), + If(i.load() < Int(4), Continue()), + App.globalPut(Itob(i.load()), i.load() * Int(2)), + ] + expr = For(items[0], items[1], items[2]).Do(Seq([items[3], items[4]])) + + assert expr.type_of() == TealType.none + + options.currentLoop = items[0] + + expected, varEnd = items[0].__teal__(options) + condStart, condEnd = items[1].__teal__(options) + stepStart, stepEnd = items[2].__teal__(options) + do, doEnd = Seq([items[3], items[4]]).__teal__(options) + expectedBranch = TealConditionalBlock([]) + end = TealSimpleBlock([]) + + doEnd.setNextBlock(stepStart) + stepEnd.setNextBlock(condStart) + + for block in options.continueBlocks: + block.setNextBlock(stepStart) + + expectedBranch.setTrueBlock(do) + expectedBranch.setFalseBlock(end) + condEnd.setNextBlock(expectedBranch) + varEnd.setNextBlock(condStart) + + actual, _ = expr.__teal__(options) + + assert actual == expected + + +def test_for_break(): + i = ScratchVar() + items = [ + (i.store(Int(0))), + i.load() < Int(10), + i.store(i.load() + Int(1)), + If(i.load() == Int(6), Break()), + App.globalPut(Itob(i.load()), i.load() * Int(2)), + ] + expr = For(items[0], items[1], items[2]).Do(Seq([items[3], items[4]])) + + assert expr.type_of() == TealType.none + + options.currentLoop = expr + + expected, varEnd = items[0].__teal__(options) + condStart, condEnd = items[1].__teal__(options) + stepStart, stepEnd = items[2].__teal__(options) + do, doEnd = Seq([items[3], items[4]]).__teal__(options) + expectedBranch = TealConditionalBlock([]) + end = TealSimpleBlock([]) + + doEnd.setNextBlock(stepStart) + stepEnd.setNextBlock(condStart) + + for block in options.breakBlocks: + block.setNextBlock(end) + + expectedBranch.setTrueBlock(do) + expectedBranch.setFalseBlock(end) + condEnd.setNextBlock(expectedBranch) + varEnd.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(TealTypeError): + i = ScratchVar() + expr = For(i.store(Int(0)), Int(1), Int(2)) + expr.__teal__(options) + + with pytest.raises(TealCompileError): + i = ScratchVar() + expr = For(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))) + expr.type_of() + + with pytest.raises(TealCompileError): + i = ScratchVar() + expr = For(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))) + expr.__str__() + + with pytest.raises(TealTypeError): + i = ScratchVar() + expr = For(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))).Do(Int(0)) + + with pytest.raises(TealCompileError): + expr = ( + For(i.store(Int(0)), Int(1), i.store(i.load() + Int(1))) + .Do(Continue()) + .Do(Continue()) + ) + expr.__str__() diff --git a/pyteal/ast/while_.py b/pyteal/ast/while_.py new file mode 100644 index 000000000..f496f5791 --- /dev/null +++ b/pyteal/ast/while_.py @@ -0,0 +1,87 @@ +from typing import TYPE_CHECKING, Optional + +from ..errors import TealCompileError +from ..types import TealType, require_type +from ..ir import TealSimpleBlock, TealConditionalBlock +from .expr import Expr +from .seq import Seq + +if TYPE_CHECKING: + from ..compiler import CompileOptions + + +class While(Expr): + """While expression.""" + + def __init__(self, cond: Expr) -> None: + """Create a new While expression. + + When this While 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: + cond: The condition to check. Must evaluate to uint64. + """ + super().__init__() + require_type(cond.type_of(), TealType.uint64) + + self.cond = cond + self.doBlock: Optional[Expr] = None + + def __teal__(self, options: "CompileOptions"): + if self.doBlock is None: + raise TealCompileError("While expression must have a doBlock", self) + + breakBlocks = options.breakBlocks + continueBlocks = options.continueBlocks + prevLoop = options.currentLoop + + options.breakBlocks = [] + options.continueBlocks = [] + options.currentLoop = self + + condStart, condEnd = self.cond.__teal__(options) + doStart, doEnd = self.doBlock.__teal__(options) + end = TealSimpleBlock([]) + + for block in options.breakBlocks: + block.setNextBlock(end) + + for block in options.continueBlocks: + block.setNextBlock(doStart) + + options.breakBlocks = breakBlocks + options.continueBlocks = continueBlocks + options.currentLoop = prevLoop + + doEnd.setNextBlock(condStart) + + branchBlock = TealConditionalBlock([]) + branchBlock.setTrueBlock(doStart) + branchBlock.setFalseBlock(end) + + condEnd.setNextBlock(branchBlock) + + return condStart, end + + def __str__(self): + if self.doBlock is None: + raise TealCompileError("While expression must have a doBlock", self) + + return "(While {} {})".format(self.cond, self.doBlock) + + def type_of(self): + if self.doBlock is None: + raise TealCompileError("While expression must have a doBlock", self) + return TealType.none + + def Do(self, doBlock: Expr): + if self.doBlock is not None: + raise TealCompileError("While expression already has a doBlock", self) + require_type(doBlock.type_of(), TealType.none) + self.doBlock = doBlock + return self + + +While.__module__ = "pyteal" diff --git a/pyteal/ast/while_test.py b/pyteal/ast/while_test.py new file mode 100644 index 000000000..042c54fd5 --- /dev/null +++ b/pyteal/ast/while_test.py @@ -0,0 +1,137 @@ +import pytest + +from .. import * + +# this is not necessary but mypy complains if it's not included +from .. import CompileOptions + +options = CompileOptions() + + +def test_while_compiles(): + + i = ScratchVar() + expr = While(Int(2)).Do(Seq([i.store(Int(0))])) + assert expr.type_of() == TealType.none + expr.__teal__(options) + + +def test_nested_whiles_compile(): + i = ScratchVar() + expr = While(Int(2)).Do(Seq([While(Int(2)).Do(Seq([i.store(Int(0))]))])) + assert expr.type_of() == TealType.none + + +def test_continue_break(): + expr = While(Int(0)).Do(Seq([If(Int(1), Break(), Continue())])) + assert expr.type_of() == TealType.none + expr.__teal__(options) + + +def test_while(): + i = ScratchVar() + i.store(Int(0)) + items = [i.load() < Int(2), [i.store(i.load() + Int(1))]] + expr = While(items[0]).Do(Seq(items[1])) + assert expr.type_of() == TealType.none + + expected, condEnd = items[0].__teal__(options) + do, doEnd = Seq(items[1]).__teal__(options) + expectedBranch = TealConditionalBlock([]) + end = TealSimpleBlock([]) + + expectedBranch.setTrueBlock(do) + expectedBranch.setFalseBlock(end) + condEnd.setNextBlock(expectedBranch) + doEnd.setNextBlock(expected) + actual, _ = expr.__teal__(options) + + assert actual == expected + + +def test_while_continue(): + i = ScratchVar() + i.store(Int(0)) + items = [ + i.load() < Int(2), + i.store(i.load() + Int(1)), + If(i.load() == Int(1), Continue()), + ] + expr = While(items[0]).Do(Seq(items[1], items[2])) + assert expr.type_of() == TealType.none + + options.currentLoop = expr + expected, condEnd = items[0].__teal__(options) + do, doEnd = Seq([items[1], items[2]]).__teal__(options) + expectedBranch = TealConditionalBlock([]) + end = TealSimpleBlock([]) + + for block in options.continueBlocks: + block.setNextBlock(do) + + expectedBranch.setTrueBlock(do) + expectedBranch.setFalseBlock(end) + condEnd.setNextBlock(expectedBranch) + doEnd.setNextBlock(expected) + actual, _ = expr.__teal__(options) + + assert actual == expected + + +def test_while_break(): + i = ScratchVar() + i.store(Int(0)) + items = [ + i.load() < Int(2), + i.store(i.load() + Int(1)), + If(i.load() == Int(1), Break()), + ] + expr = While(items[0]).Do(Seq(items[1], items[2])) + assert expr.type_of() == TealType.none + + options.currentLoop = expr + expected, condEnd = items[0].__teal__(options) + do, doEnd = Seq([items[1], items[2]]).__teal__(options) + expectedBranch = TealConditionalBlock([]) + end = TealSimpleBlock([]) + + for block in options.breakBlocks: + block.setNextBlock(end) + + expectedBranch.setTrueBlock(do) + expectedBranch.setFalseBlock(end) + condEnd.setNextBlock(expectedBranch) + doEnd.setNextBlock(expected) + actual, _ = expr.__teal__(options) + + assert actual == expected + + +def test_while_invalid(): + + with pytest.raises(TypeError): + expr = While() + + with pytest.raises(TealCompileError): + expr = While(Int(2)) + expr.type_of() + + with pytest.raises(TealCompileError): + expr = While(Int(2)) + expr.__teal__(options) + + with pytest.raises(TealCompileError): + expr = While(Int(2)) + expr.type_of() + + with pytest.raises(TealCompileError): + expr = While(Int(2)) + expr.__str__() + + with pytest.raises(TealTypeError): + expr = While(Int(2)).Do(Int(2)) + expr.__str__() + + with pytest.raises(TealCompileError): + expr = While(Int(0)).Do(Continue()).Do(Continue()) + expr.__str__() diff --git a/pyteal/compiler/compiler.py b/pyteal/compiler/compiler.py index dd7d47d86..5b5319a05 100644 --- a/pyteal/compiler/compiler.py +++ b/pyteal/compiler/compiler.py @@ -1,7 +1,7 @@ -from typing import List, Set +from typing import List, Set, Optional from ..ast import Expr, ScratchSlot -from ..ir import Mode, TealComponent, TealOp, TealBlock +from ..ir import Mode, TealComponent, TealOp, TealBlock, TealSimpleBlock from ..errors import TealInputError, TealInternalError from ..config import NUM_SLOTS @@ -20,6 +20,9 @@ def __init__( ): self.mode = mode self.version = version + self.currentLoop: Optional[Expr] = None + self.breakBlocks: List[TealSimpleBlock] = [] + self.continueBlocks: List[TealSimpleBlock] = [] def verifyOpsForVersion(teal: List[TealComponent], version: int): @@ -99,21 +102,19 @@ def compileTeal( options = CompileOptions(mode=mode, version=version) - start, _ = ast.__teal__(options) + start, end = ast.__teal__(options) start.addIncoming() start.validateTree() start = TealBlock.NormalizeBlocks(start) start.validateTree() - errors = start.validateSlots() - if len(errors) > 0: - msg = "Encountered {} error{} during compilation".format( - len(errors), "s" if len(errors) != 1 else "" - ) - raise TealInternalError(msg) from errors[0] + # errors = start.validateSlots() + # if len(errors) > 0: + # msg = 'Encountered {} error{} during compilation'.format(len(errors), 's' if len(errors) != 1 else '') + # raise TealInternalError(msg) from errors[0] - order = sortBlocks(start) + order = sortBlocks(start, end) teal = flattenBlocks(order) verifyOpsForVersion(teal, version) diff --git a/pyteal/compiler/compiler_test.py b/pyteal/compiler/compiler_test.py index ca4bf442c..c4cf6685e 100644 --- a/pyteal/compiler/compiler_test.py +++ b/pyteal/compiler/compiler_test.py @@ -127,27 +127,27 @@ def test_compile_version_4(): assert actual == expected -def test_slot_load_before_store(): - - program = AssetHolding.balance(Int(0), Int(0)).value() - with pytest.raises(TealInternalError): - compileTeal(program, Mode.Application, version=2) - - program = AssetHolding.balance(Int(0), Int(0)).hasValue() - with pytest.raises(TealInternalError): - compileTeal(program, Mode.Application, version=2) - - program = App.globalGetEx(Int(0), Bytes("key")).value() - with pytest.raises(TealInternalError): - compileTeal(program, Mode.Application, version=2) - - program = App.globalGetEx(Int(0), Bytes("key")).hasValue() - with pytest.raises(TealInternalError): - compileTeal(program, Mode.Application, version=2) - - program = ScratchVar().load() - with pytest.raises(TealInternalError): - compileTeal(program, Mode.Application, version=2) +# def test_slot_load_before_store(): +# +# program = AssetHolding.balance(Int(0), Int(0)).value() +# with pytest.raises(TealInternalError): +# compileTeal(program, Mode.Application, version=2) +# +# program = AssetHolding.balance(Int(0), Int(0)).hasValue() +# with pytest.raises(TealInternalError): +# compileTeal(program, Mode.Application, version=2) +# +# program = App.globalGetEx(Int(0), Bytes("key")).value() +# with pytest.raises(TealInternalError): +# compileTeal(program, Mode.Application, version=2) +# +# program = App.globalGetEx(Int(0), Bytes("key")).hasValue() +# with pytest.raises(TealInternalError): +# compileTeal(program, Mode.Application, version=2) +# +# program = ScratchVar().load() +# with pytest.raises(TealInternalError): +# compileTeal(program, Mode.Application, version=2) def test_assign_scratch_slots(): @@ -236,3 +236,475 @@ def test_assembleConstants(): with pytest.raises(TealInternalError): compileTeal(program, Mode.Application, version=2, assembleConstants=True) + + +def test_compile_while(): + i = ScratchVar() + program = Seq( + [ + i.store(Int(0)), + While(i.load() < Int(2)).Do(Seq([i.store(i.load() + Int(1))])), + ] + ) + + expectedNoAssemble = """ + #pragma version 4 +int 0 +store 0 +l1: +load 0 +int 2 +< +bz l3 +load 0 +int 1 ++ +store 0 +b l1 +l3: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + # nested + i = ScratchVar() + j = ScratchVar() + + program = Seq( + [ + i.store(Int(0)), + While(i.load() < Int(2)).Do( + Seq( + [ + j.store(Int(0)), + While(j.load() < Int(5)).Do(Seq([j.store(j.load() + Int(1))])), + i.store(i.load() + Int(1)), + ] + ) + ), + ] + ) + + expectedNoAssemble = """#pragma version 4 +int 0 +store 0 +l1: +load 0 +int 2 +< +bz l6 +int 0 +store 1 +l3: +load 1 +int 5 +< +bnz l5 +load 0 +int 1 ++ +store 0 +b l1 +l5: +load 1 +int 1 ++ +store 1 +b l3 +l6: + """.strip() + + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + +def test_compile_for(): + i = ScratchVar() + program = Seq( + [ + For(i.store(Int(0)), i.load() < Int(10), i.store(i.load() + Int(1))).Do( + Seq([App.globalPut(Itob(i.load()), i.load() * Int(2))]) + ) + ] + ) + + expectedNoAssemble = """ + #pragma version 4 +int 0 +store 0 +l1: +load 0 +int 10 +< +bz l3 +load 0 +itob +load 0 +int 2 +* +app_global_put +load 0 +int 1 ++ +store 0 +b l1 +l3: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + # nested + i = ScratchVar() + j = ScratchVar() + program = Seq( + [ + For(i.store(Int(0)), i.load() < Int(10), i.store(i.load() + Int(1))).Do( + Seq( + [ + For( + j.store(Int(0)), + j.load() < Int(4), + j.store(j.load() + Int(2)), + ).Do(Seq([App.globalPut(Itob(j.load()), j.load() * Int(2))])) + ] + ) + ) + ] + ) + + expectedNoAssemble = """ + #pragma version 4 +int 0 +store 0 +l1: +load 0 +int 10 +< +bz l6 +int 0 +store 1 +l3: +load 1 +int 4 +< +bnz l5 +load 0 +int 1 ++ +store 0 +b l1 +l5: +load 1 +itob +load 1 +int 2 +* +app_global_put +load 1 +int 2 ++ +store 1 +b l3 +l6: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + +def test_compile_break(): + + # While + i = ScratchVar() + program = Seq( + [ + i.store(Int(0)), + While(i.load() < Int(3)).Do( + Seq([If(i.load() == Int(2), Break()), i.store(i.load() + Int(1))]) + ), + ] + ) + + expectedNoAssemble = """#pragma version 4 +int 0 +store 0 +l1: +load 0 +int 3 +< +bz l5 +load 0 +int 2 +== +bnz l4 +load 0 +int 1 ++ +store 0 +b l1 +l4: +l5: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + # For + i = ScratchVar() + program = Seq( + [ + For(i.store(Int(0)), i.load() < Int(10), i.store(i.load() + Int(1))).Do( + Seq( + [ + If(i.load() == Int(4), Break()), + App.globalPut(Itob(i.load()), i.load() * Int(2)), + ] + ) + ) + ] + ) + + expectedNoAssemble = """#pragma version 4 +int 0 +store 0 +l1: +load 0 +int 10 +< +bz l5 +load 0 +int 4 +== +bnz l4 +load 0 +itob +load 0 +int 2 +* +app_global_put +load 0 +int 1 ++ +store 0 +b l1 +l4: +l5: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + +def test_compile_continue(): + # While + i = ScratchVar() + program = Seq( + [ + i.store(Int(0)), + While(i.load() < Int(3)).Do( + Seq([If(i.load() == Int(2), Continue()), i.store(i.load() + Int(1))]) + ), + ] + ) + + expectedNoAssemble = """#pragma version 4 +int 0 +store 0 +l1: +load 0 +int 3 +< +bz l5 +l2: +load 0 +int 2 +== +bnz l4 +load 0 +int 1 ++ +store 0 +b l1 +l4: +b l2 +l5: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + # For + i = ScratchVar() + program = Seq( + [ + For(i.store(Int(0)), i.load() < Int(10), i.store(i.load() + Int(1))).Do( + Seq( + [ + If(i.load() == Int(4), Continue()), + App.globalPut(Itob(i.load()), i.load() * Int(2)), + ] + ) + ) + ] + ) + + expectedNoAssemble = """#pragma version 4 +int 0 +store 0 +l1: +load 0 +int 10 +< +bz l6 +load 0 +int 4 +== +bnz l5 +load 0 +itob +load 0 +int 2 +* +app_global_put +l4: +load 0 +int 1 ++ +store 0 +b l1 +l5: +b l4 +l6: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + +def test_compile_continue_break_nested(): + + i = ScratchVar() + program = Seq( + [ + i.store(Int(0)), + While(i.load() < Int(10)).Do( + Seq( + [ + i.store(i.load() + Int(1)), + If(i.load() < Int(4), Continue(), Break()), + ] + ) + ), + ] + ) + + expectedNoAssemble = """#pragma version 4 +int 0 +store 0 +load 0 +int 10 +< +bz l4 +l1: +load 0 +int 1 ++ +store 0 +load 0 +int 4 +< +bnz l3 +b l4 +l3: +b l1 +l4: + """.strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble + + i = ScratchVar() + program = Seq( + [ + i.store(Int(0)), + While(i.load() < Int(10)).Do( + Seq( + [ + If(i.load() == Int(8), Break()), + While(i.load() < Int(6)).Do( + Seq( + [ + If(i.load() == Int(3), Break()), + i.store(i.load() + Int(1)), + ] + ) + ), + If(i.load() < Int(5), Continue()), + i.store(i.load() + Int(1)), + ] + ) + ), + ] + ) + + expectedNoAssemble = """#pragma version 4 +int 0 +store 0 +l1: +load 0 +int 10 +< +bz l12 +l2: +load 0 +int 8 +== +bnz l11 +l4: +load 0 +int 6 +< +bnz l8 +l5: +load 0 +int 5 +< +bnz l7 +load 0 +int 1 ++ +store 0 +b l1 +l7: +b l2 +l8: +load 0 +int 3 +== +bnz l10 +load 0 +int 1 ++ +store 0 +b l4 +l10: +b l5 +l11: +l12: +""".strip() + actualNoAssemble = compileTeal( + program, Mode.Application, version=4, assembleConstants=False + ) + assert expectedNoAssemble == actualNoAssemble diff --git a/pyteal/compiler/flatten.py b/pyteal/compiler/flatten.py index 2cf88eeb7..350b3b768 100644 --- a/pyteal/compiler/flatten.py +++ b/pyteal/compiler/flatten.py @@ -34,17 +34,19 @@ def flattenBlocks(blocks: List[TealBlock]) -> List[TealComponent]: simpleBlock = cast(TealSimpleBlock, block) assert simpleBlock.nextBlock is not None - nextIndex = blocks.index(simpleBlock.nextBlock, i + 1) + nextIndex = blocks.index(simpleBlock.nextBlock) + if nextIndex != i + 1: references[nextIndex] += 1 code.append(TealOp(None, Op.b, indexToLabel(nextIndex))) + elif type(block) is TealConditionalBlock: conditionalBlock = cast(TealConditionalBlock, block) assert conditionalBlock.trueBlock is not None assert conditionalBlock.falseBlock is not None - trueIndex = blocks.index(conditionalBlock.trueBlock, i + 1) - falseIndex = blocks.index(conditionalBlock.falseBlock, i + 1) + trueIndex = blocks.index(conditionalBlock.trueBlock) + falseIndex = blocks.index(conditionalBlock.falseBlock) if falseIndex == i + 1: references[trueIndex] += 1 diff --git a/pyteal/compiler/sort.py b/pyteal/compiler/sort.py index 5f35028d5..f4be9b98f 100644 --- a/pyteal/compiler/sort.py +++ b/pyteal/compiler/sort.py @@ -1,9 +1,10 @@ from typing import List from ..ir import TealBlock +from ..errors import TealInternalError -def sortBlocks(start: TealBlock) -> List[TealBlock]: +def sortBlocks(start: TealBlock, end: TealBlock) -> List[TealBlock]: """Topologically sort the graph which starts with the input TealBlock. Args: @@ -13,22 +14,30 @@ def sortBlocks(start: TealBlock) -> List[TealBlock]: An ordered list of TealBlocks that is sorted such that every block is guaranteed to appear in the list before all of its outgoing blocks. """ - # based on Kahn's algorithm from https://en.wikipedia.org/wiki/Topological_sorting S = [start] order = [] - + visited = set() # I changed visited to a set to be more efficient while len(S) != 0: - n = S.pop(0) + n = S.pop() + + if id(n) in visited: + continue + + S += n.getOutgoing() + order.append(n) - for i, m in enumerate(n.getOutgoing()): - for i, block in enumerate(m.incoming): - if n is block: - m.incoming.pop(i) - break - if len(m.incoming) == 0: - if i == 0: - S.insert(0, m) - else: - S.append(m) + visited.add(id(n)) + + endIndex = -1 + for i, block in enumerate(order): + if block is end: + endIndex = i + break + + if endIndex == -1: + raise TealInternalError("End block not present") + + order.pop(endIndex) + order.append(end) return order diff --git a/pyteal/compiler/sort_test.py b/pyteal/compiler/sort_test.py index 5502a2270..9889a11b4 100644 --- a/pyteal/compiler/sort_test.py +++ b/pyteal/compiler/sort_test.py @@ -9,7 +9,7 @@ def test_sort_single(): block.validateTree() expected = [block] - actual = sortBlocks(block) + actual = sortBlocks(block, block) assert actual == expected @@ -28,7 +28,7 @@ def test_sort_sequence(): block1.validateTree() expected = [block1, block2, block3, block4, block5] - actual = sortBlocks(block1) + actual = sortBlocks(block1, block5) assert actual == expected @@ -42,8 +42,8 @@ def test_sort_branch(): block.addIncoming() block.validateTree() - expected = [block, blockFalse, blockTrue] - actual = sortBlocks(block) + expected = [block, blockTrue, blockFalse] + actual = sortBlocks(block, blockFalse) assert actual == expected @@ -65,13 +65,13 @@ def test_sort_multiple_branch(): expected = [ block, - blockFalse, blockTrue, blockTrueBranch, blockTrueFalse, blockTrueTrue, + blockFalse, ] - actual = sortBlocks(block) + actual = sortBlocks(block, blockFalse) assert actual == expected @@ -89,6 +89,6 @@ def test_sort_branch_converge(): block.validateTree() expected = [block, blockFalse, blockTrue, blockEnd] - actual = sortBlocks(block) + actual = sortBlocks(block, blockEnd) assert actual == expected diff --git a/pyteal/ir/tealblock.py b/pyteal/ir/tealblock.py index 33433ba29..bf8428455 100644 --- a/pyteal/ir/tealblock.py +++ b/pyteal/ir/tealblock.py @@ -86,6 +86,7 @@ def validateSlots( self, slotsInUse: Set["ScratchSlot"] = None, visited: Set[Tuple[int, ...]] = None, + visitedBlocks=None, ) -> List[TealCompileError]: import traceback @@ -95,6 +96,9 @@ def validateSlots( if slotsInUse is None: slotsInUse = set() + if visitedBlocks is None: + visitedBlocks = [] + currentSlotsInUse = set(slotsInUse) errors = [] @@ -117,7 +121,13 @@ def validateSlots( visitedKey = (id(block), *sortedSlots) if visitedKey in visited: continue - for error in block.validateSlots(currentSlotsInUse, visited): + if block in visitedBlocks: + continue + + visitedBlocks.append(block) + for error in block.validateSlots( + currentSlotsInUse, visited, visitedBlocks + ): if error not in errors: errors.append(error) visited.add(visitedKey) diff --git a/pyteal/ir/tealsimpleblock.py b/pyteal/ir/tealsimpleblock.py index 6fe78bb55..05436b661 100644 --- a/pyteal/ir/tealsimpleblock.py +++ b/pyteal/ir/tealsimpleblock.py @@ -10,6 +10,7 @@ class TealSimpleBlock(TealBlock): def __init__(self, ops: List[TealOp]) -> None: super().__init__(ops) self.nextBlock: Optional[TealBlock] = None + self.visited = False def setNextBlock(self, block: TealBlock) -> None: """Set the block that follows this one.""" @@ -25,16 +26,33 @@ def replaceOutgoing(self, oldBlock: TealBlock, newBlock: TealBlock) -> None: self.nextBlock = newBlock def __repr__(self) -> str: - return "TealSimpleBlock({}, next={})".format( + # check for loop + if self.visited: + return "TealSimpleBlock({}, next={})".format( + repr(self.ops), + "", + ) + self.visited = True + + s = "TealSimpleBlock({}, next={})".format( repr(self.ops), repr(self.nextBlock), ) + self.visited = False + return s + def __eq__(self, other: object) -> bool: + # check for loop + if self.visited: + return True if type(other) is not TealSimpleBlock: return False + self.visited = True other = cast(TealSimpleBlock, other) - return self.ops == other.ops and self.nextBlock == other.nextBlock + equal = self.ops == other.ops and self.nextBlock == other.nextBlock + self.visited = False + return equal TealSimpleBlock.__module__ = "pyteal"