Optimize DupN code generation & fix testing bug (#623)

pyteal/ast/frame.py

@@ -44,14 +44,7 @@ def __init__(self, local_type: TealType, count: int):
                 )
 
     def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBlock]:
-        if self.count == 1:
-            inst_srt, inst_end = self.auto_instance.__teal__(options)
-            return inst_srt, inst_end
-        else:
-            dupn_srt, dupn_end = DupN(self.auto_instance, self.count - 1).__teal__(
-                options
-            )
-            return dupn_srt, dupn_end
+        return DupN(self.auto_instance, self.count - 1).__teal__(options)
 
     def __str__(self) -> str:
         return f"(LocalTypeSegment: (type: {self.local_type}) (count: {self.count}))"
@@ -330,6 +323,14 @@ class DupN(Expr):
     """
 
     def __init__(self, value: Expr, repetition: int):
+        """Create a DupN expression.
+
+        Args:
+            value: The value to be duplicated.
+            repetition: How many additional times the value should be added to the stack. At the end
+                of this operation, `repetition+1` elements will be added to the stack. Zero can be
+                specified here to indicate no duplication.
+        """
         super().__init__()
         require_type(value, TealType.anytype)
         if repetition < 0:
@@ -338,6 +339,15 @@ def __init__(self, value: Expr, repetition: int):
         self.repetition = repetition
 
     def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBlock]:
+        if self.repetition == 0:
+            # no duplication required
+            return self.value.__teal__(options)
+
+        if self.repetition == 1:
+            # use normal dup op for just 1 duplication
+            op = TealOp(self, Op.dup)
+            return TealBlock.FromOp(options, op, self.value)
+
         verifyProgramVersion(
             Op.dupn.min_version,
             options.version,

pyteal/ast/frame_test.py

@@ -92,7 +92,44 @@ def test_frame_bury_invalid():
         FrameBury(pt.Int(1), 1).__teal__(avm7Options)
 
 
-def test_dupn():
+def test_dupn_zero():
+    byte_expr = pt.Bytes("Astartes")
+    expr = DupN(byte_expr, 0)
+    assert not expr.has_return()
+    assert expr.type_of() == byte_expr.type_of()
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(byte_expr, pt.Op.byte, '"Astartes"'),
+        ]
+    )
+    actual, _ = expr.__teal__(avm8Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+
+def test_dupn_single():
+    byte_expr = pt.Bytes("Astartes")
+    expr = DupN(byte_expr, 1)
+    assert not expr.has_return()
+    assert expr.type_of() == byte_expr.type_of()
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(byte_expr, pt.Op.byte, '"Astartes"'),
+            pt.TealOp(expr, pt.Op.dup),
+        ]
+    )
+    actual, _ = expr.__teal__(avm8Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+
+def test_dupn_multiple():
     byte_expr = pt.Bytes("Astartes")
     expr = DupN(byte_expr, 4)
     assert not expr.has_return()
@@ -113,13 +150,13 @@ def test_dupn():
 
 def test_dupn_invalid():
     with pytest.raises(pt.TealTypeError):
-        DupN(pt.Seq(), 1)
+        DupN(pt.Seq(), 10)
 
     with pytest.raises(pt.TealInputError):
-        DupN(pt.Int(1), -1)
+        DupN(pt.Int(1), -10)
 
     with pytest.raises(pt.TealInputError):
-        DupN(pt.Int(1), 1).__teal__(avm7Options)
+        DupN(pt.Int(1), 10).__teal__(avm7Options)
 
 
 def test_local_type_segment_invalid():

tests/compile_asserts.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from difflib import unified_diff
 
 from pyteal.compiler import compileTeal
 from pyteal.ir import Mode
@@ -8,7 +9,7 @@
 GENERATED = PATH / "generated"
 
 
-def compile_and_save(approval, version, test_name):
+def compile_and_save(approval, version: int, test_name: str) -> tuple[Path, str, str]:
     compiled = compileTeal(approval(), mode=Mode.Application, version=version)
     name = approval.__name__
     tealdir = GENERATED / test_name
@@ -24,43 +25,29 @@ def compile_and_save(approval, version, test_name):
     return tealdir, name, compiled
 
 
-def mismatch_ligature(expected, actual):
-    la, le = len(actual), len(expected)
-    mm_idx = -1
-    for i in range(min(la, le)):
-        if expected[i] != actual[i]:
-            mm_idx = i
-            break
-    if mm_idx < 0:
-        return ""
-    return " " * (mm_idx) + "X" + "-" * (max(la, le) - mm_idx - 1)
+def assert_teal_as_expected(path2actual: Path, path2expected: Path):
+    with open(path2actual, "r") as f:
+        actual_lines = f.readlines()
 
+    with open(path2expected, "r") as f:
+        expected_lines = f.readlines()
 
-def assert_teal_as_expected(path2actual, path2expected):
-    with open(path2actual, "r") as fa, open(path2expected, "r") as fe:
-        alines = fa.read().split("\n")
-        elines = fe.read().split("\n")
-
-        assert len(elines) == len(
-            alines
-        ), f"""EXPECTED {len(elines)} lines for {path2expected} but ACTUALLY got {len(alines)} lines in {path2actual}"""
+    diff = list(
+        unified_diff(
+            expected_lines,
+            actual_lines,
+            fromfile=str(path2expected),
+            tofile=str(path2actual),
+            n=3,
+        )
+    )
 
-        for i, actual in enumerate(alines):
-            expected = elines[i]
-            assert expected.startswith(
-                actual
-            ), f"""ACTUAL line in {path2actual}
-LINE{i+1}:
-{actual}
-{mismatch_ligature(expected, actual)}
-DOES NOT prefix the EXPECTED (which should have been actual + some commentary) in {path2expected}:
-LINE{i+1}:
-{expected}  
-{mismatch_ligature(expected, actual)}
-"""
+    assert (
+        len(diff) == 0
+    ), f"Difference between expected and actual TEAL code:\n\n{''.join(diff)}"
 
 
-def assert_new_v_old(approve_func, version, test_name):
+def assert_new_v_old(approve_func, version: int, test_name: str):
     tealdir, name, compiled = compile_and_save(approve_func, version, test_name)
 
     print(

tests/integration/teal/roundtrip/app_roundtrip_()_v8.teal

@@ -16,9 +16,9 @@ tuplecomplement_0:
 proto 1 1
 byte ""
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 byte ""
 frame_bury 0
 retsub
@@ -29,9 +29,9 @@ proto 1 1
 byte ""
 dupn 2
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 frame_dig -1
 callsub tuplecomplement_0
 frame_bury 1

tests/integration/teal/roundtrip/app_roundtrip_(bool)[10]_v8.teal

@@ -18,7 +18,7 @@ byte ""
 int 0
 dupn 3
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 0
 getbit
@@ -39,9 +39,9 @@ proto 1 1
 byte ""
 dupn 10
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 1
 int 0
@@ -170,9 +170,9 @@ proto 1 1
 byte ""
 dupn 2
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 frame_dig -1
 callsub arraycomplement_1
 frame_bury 1

tests/integration/teal/roundtrip/app_roundtrip_(bool)_v8.teal

@@ -18,7 +18,7 @@ byte ""
 int 0
 dupn 3
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 0
 getbit
@@ -39,9 +39,9 @@ proto 1 1
 byte ""
 dupn 2
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 frame_dig -1
 callsub tuplecomplement_0
 frame_bury 1

tests/integration/teal/roundtrip/app_roundtrip_(bool,address,(uint64,bool),byte[10],bool[4],uint64)_v8.teal

@@ -18,7 +18,7 @@ byte ""
 int 0
 dupn 6
 byte ""
-dupn 1
+dup
 int 0
 dupn 97
 byte ""
@@ -91,9 +91,9 @@ proto 1 1
 byte ""
 dupn 4
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 frame_dig -1
 callsub tuplecomplement_0
 frame_bury 2
@@ -140,7 +140,7 @@ byte ""
 int 0
 dupn 97
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 1
 int 0
@@ -598,7 +598,7 @@ byte ""
 int 0
 dupn 5
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 0
 extract_uint64
@@ -644,7 +644,7 @@ byte ""
 int 0
 dupn 11
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 1
 int 0
@@ -805,7 +805,7 @@ byte ""
 int 0
 dupn 5
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 0
 getbit

tests/integration/teal/roundtrip/app_roundtrip_(bool,byte)_v8.teal

@@ -18,7 +18,7 @@ byte ""
 int 0
 dupn 5
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 0
 getbit
@@ -51,9 +51,9 @@ proto 1 1
 byte ""
 dupn 2
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 frame_dig -1
 callsub tuplecomplement_0
 frame_bury 1

tests/integration/teal/roundtrip/app_roundtrip_(bool,byte,address,string)_v8.teal

@@ -18,19 +18,19 @@ byte ""
 int 0
 dupn 4
 byte ""
-dupn 1
+dup
 int 0
 dupn 97
 byte ""
 dupn 3
 int 0
 dupn 7
 byte ""
-dupn 1
+dup
 int 0
 dupn 2
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 0
 getbit
@@ -94,9 +94,9 @@ proto 1 1
 byte ""
 dupn 2
 int 0
-dupn 1
+dup
 byte ""
-dupn 1
+dup
 frame_dig -1
 callsub tuplecomplement_0
 frame_bury 1
@@ -205,7 +205,7 @@ byte ""
 int 0
 dupn 97
 byte ""
-dupn 1
+dup
 frame_dig -1
 int 1
 int 0
@@ -663,7 +663,7 @@ byte ""
 int 0
 dupn 7
 byte ""
-dupn 1
+dup
 int 0
 frame_dig -1
 int 1