Bug fix for #331 continue skips while condition check (#332)

* wire continue-next to do block * update compiler test * add one semantic test
algorand · May 10, 2022 · 60a04f8 · 60a04f8
1 parent 22deba5
commit 60a04f8
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 10 deletions.
diff --git a/pyteal/ast/while_.py b/pyteal/ast/while_.py
@@ -52,7 +52,7 @@ def __teal__(self, options: "CompileOptions"):
             block.setNextBlock(end)
 
         for block in continueBlocks:
-            block.setNextBlock(doStart)
+            block.setNextBlock(condStart)
 
         return condStart, end
 

diff --git a/pyteal/compiler/compiler_test.py b/pyteal/compiler/compiler_test.py
@@ -652,11 +652,10 @@ def test_compile_continue():
 int 3
 <
 bz main_l4
-main_l2:
 load 0
 int 2
 ==
-bnz main_l2
+bnz main_l1
 load 0
 int 1
 +
@@ -665,7 +664,41 @@ def test_compile_continue():
 main_l4:
 int 1
 return
-                """.strip()
+    """.strip()
+    actual = pt.compileTeal(
+        program, pt.Mode.Application, version=4, assembleConstants=False
+    )
+    assert expected == actual
+
+    # pt.While
+    program = pt.Seq(
+        i.store(pt.Int(0)),
+        pt.While(i.load() < pt.Int(30)).Do(
+            pt.Seq(
+                i.store(i.load() + pt.Int(1)),
+                pt.Continue(),
+            )
+        ),
+        pt.Return(pt.Int(1)),
+    )
+
+    expected = """#pragma version 4
+int 0
+store 0
+main_l1:
+load 0
+int 30
+<
+bz main_l3
+load 0
+int 1
++
+store 0
+b main_l1
+main_l3:
+int 1
+return
+    """.strip()
     actual = pt.compileTeal(
         program, pt.Mode.Application, version=4, assembleConstants=False
     )
@@ -744,11 +777,11 @@ def test_compile_continue_break_nested():
     expected = """#pragma version 4
 int 0
 store 0
+main_l1:
 load 0
 int 10
 <
-bz main_l2
-main_l1:
+bz main_l3
 load 0
 int 1
 +
@@ -757,7 +790,7 @@ def test_compile_continue_break_nested():
 int 4
 <
 bnz main_l1
-main_l2:
+main_l3:
 int 1
 return
     """.strip()
@@ -799,7 +832,6 @@ def test_compile_continue_break_nested():
 int 10
 <
 bz main_l8
-main_l2:
 load 0
 int 8
 ==
@@ -813,7 +845,7 @@ def test_compile_continue_break_nested():
 load 0
 int 5
 <
-bnz main_l2
+bnz main_l1
 load 0
 int 1
 +

diff --git a/tests/integration/graviton_test.py b/tests/integration/graviton_test.py
@@ -13,6 +13,9 @@
     Mode,
     ScratchVar,
     Seq,
+    While,
+    Continue,
+    Return,
     Subroutine,
     SubroutineFnWrapper,
     TealType,
@@ -770,9 +773,47 @@ def euclid(x, y):
         Invariant(predicate).validates(property, inputs, inspectors)
 
 
+def blackbox_pyteal_while_continue_test():
+    @Blackbox(input_types=[TealType.uint64])
+    @Subroutine(TealType.uint64)
+    def while_continue_accumulation(n):
+        i = ScratchVar(TealType.uint64)
+        return Seq(
+            i.store(Int(0)),
+            While(i.load() < n).Do(
+                Seq(
+                    i.store(i.load() + Int(1)),
+                    Continue(),
+                )
+            ),
+            Return(i.load()),
+        )
+
+    approval_lsig = blackbox_pyteal(while_continue_accumulation, Mode.Signature)
+    lsig_teal = compileTeal(approval_lsig(), Mode.Signature, version=6)
+    algod = algod_with_assertion()
+
+    for x in range(30):
+        args = [x]
+        lsig_result = DryRunExecutor.dryrun_logicsig(algod, lsig_teal, args)
+        if x == 0:
+            assert not lsig_result.passed()
+        else:
+            assert lsig_result.passed()
+
+        assert lsig_result.stack_top() == x, lsig_result.report(
+            args, "stack_top() gave unexpected results for lsig"
+        )
+
+
 @pytest.mark.parametrize(
     "example",
-    [blackbox_pyteal_example1, blackbox_pyteal_example2, blackbox_pyteal_example3],
+    [
+        blackbox_pyteal_example1,
+        blackbox_pyteal_example2,
+        blackbox_pyteal_example3,
+        blackbox_pyteal_while_continue_test,
+    ],
 )
 def test_blackbox_pyteal_examples(example):
     example()