Fix bug where String.set duplicates the argument Expr in the AST (#…

…508)

* Fix `String.set` when called with `Expr`

* Add changelog entry

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Unreleased
 
+## Fixed
+* Fix AST duplication bug in `String.set` when called with an `Expr` argument ([#508](https://github.com/algorand/pyteal/pull/508))
+
 # 0.16.0
 
 ## Added

pyteal/ast/abi/string.py

@@ -1,6 +1,8 @@
 from typing import Union, Sequence, cast
 from collections.abc import Sequence as CollectionSequence
 
+from algosdk.abi import ABIType
+
 from pyteal.ast.abi.uint import Byte
 from pyteal.ast.abi.type import ComputedValue, BaseType
 from pyteal.ast.abi.array_dynamic import DynamicArray, DynamicArrayTypeSpec
@@ -9,15 +11,27 @@
 from pyteal.ast.int import Int
 from pyteal.ast.expr import Expr
 from pyteal.ast.bytes import Bytes
+from pyteal.ast.seq import Seq
+from pyteal.ast.scratchvar import ScratchVar
 from pyteal.ast.unaryexpr import Itob, Len
 from pyteal.ast.substring import Suffix
 from pyteal.ast.naryexpr import Concat
 
 from pyteal.errors import TealInputError
 
 
-def _encoded_string(s: Expr):
-    return Concat(Suffix(Itob(Len(s)), Int(6)), s)
+def _encoded_byte_string(s: bytes | bytearray) -> Expr:
+    prefix = ABIType.from_string("uint16").encode(len(s))
+    return Bytes(prefix + s)
+
+
+def _store_encoded_expr_byte_string_into_var(value: Expr, location: ScratchVar) -> Expr:
+    return Seq(
+        location.store(value),
+        location.store(
+            Concat(Suffix(Itob(Len(location.load())), Int(6)), location.load())
+        ),
+    )
 
 
 class StringTypeSpec(DynamicArrayTypeSpec):
@@ -101,10 +115,14 @@ def set(
                 raise TealInputError(
                     f"Got {value} with type spec {value.type_spec()}, expected {StringTypeSpec}"
                 )
-            case str() | bytes():
-                return self.stored_value.store(_encoded_string(Bytes(value)))
+            case bytes() | bytearray():
+                return self.stored_value.store(_encoded_byte_string(value))
+            case str():
+                return self.stored_value.store(_encoded_byte_string(value.encode()))
             case Expr():
-                return self.stored_value.store(_encoded_string(value))
+                return _store_encoded_expr_byte_string_into_var(
+                    value, self.stored_value
+                )
             case CollectionSequence():
                 return super().set(cast(Sequence[Byte], value))
 

pyteal/ast/abi/string_test.py

@@ -4,7 +4,6 @@
 from pyteal import abi
 from pyteal.ast.abi.util import substring_for_decoding
 from pyteal.ast.abi.type_test import ContainerType
-from pyteal.util import escapeStr
 
 options = pt.CompileOptions(version=5)
 
@@ -110,62 +109,41 @@ def test_String_get():
         assert actual == expected
 
 
-def test_String_set_static():
+STATIC_SET_TESTCASES: list[tuple[str | bytes | bytearray, bytes]] = [
+    ("stringy", b"\x00\x07stringy"),
+    ("😀", b"\x00\x04\xf0\x9f\x98\x80"),
+    ("0xDEADBEEF", b"\x00\x0a0xDEADBEEF"),
+    (bytes(32), b"\x00\x20" + bytes(32)),
+    (b"alphabet_soup", b"\x00\x0dalphabet_soup"),
+    (bytearray(b"another one"), b"\x00\x0banother one"),
+]
 
-    for value_to_set in ("stringy", "😀", "0xDEADBEEF"):
-        value = abi.String()
-        expr = value.set(value_to_set)
-        assert expr.type_of() == pt.TealType.none
-        assert not expr.has_return()
-
-        expected = pt.TealSimpleBlock(
-            [
-                pt.TealOp(None, pt.Op.byte, escapeStr(value_to_set)),
-                pt.TealOp(None, pt.Op.len),
-                pt.TealOp(None, pt.Op.itob),
-                pt.TealOp(None, pt.Op.extract, 6, 0),
-                pt.TealOp(None, pt.Op.byte, escapeStr(value_to_set)),
-                pt.TealOp(None, pt.Op.concat),
-                pt.TealOp(None, pt.Op.store, value.stored_value.slot),
-            ]
-        )
 
-        actual, _ = expr.__teal__(options)
-        actual.addIncoming()
-        actual = pt.TealBlock.NormalizeBlocks(actual)
-
-        with pt.TealComponent.Context.ignoreExprEquality():
-            assert actual == expected
-
-    for value_to_set in (bytes(32), b"alphabet_soup"):
-        value = abi.String()
-        expr = value.set(value_to_set)
-        assert expr.type_of() == pt.TealType.none
-        assert not expr.has_return()
+@pytest.mark.parametrize("value_to_set, value_encoded", STATIC_SET_TESTCASES)
+def test_String_set_static(value_to_set, value_encoded):
+    value = abi.String()
+    expr = value.set(value_to_set)
+    assert expr.type_of() == pt.TealType.none
+    assert not expr.has_return()
 
-        teal_val = f"0x{value_to_set.hex()}"
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(None, pt.Op.byte, "0x" + value_encoded.hex()),
+            pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+        ]
+    )
 
-        expected = pt.TealSimpleBlock(
-            [
-                pt.TealOp(None, pt.Op.byte, teal_val),
-                pt.TealOp(None, pt.Op.len),
-                pt.TealOp(None, pt.Op.itob),
-                pt.TealOp(None, pt.Op.extract, 6, 0),
-                pt.TealOp(None, pt.Op.byte, teal_val),
-                pt.TealOp(None, pt.Op.concat),
-                pt.TealOp(None, pt.Op.store, value.stored_value.slot),
-            ]
-        )
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
 
-        actual, _ = expr.__teal__(options)
-        actual.addIncoming()
-        actual = pt.TealBlock.NormalizeBlocks(actual)
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
 
-        with pt.TealComponent.Context.ignoreExprEquality():
-            assert actual == expected
 
+def test_String_set_static_invalid():
     with pytest.raises(pt.TealInputError):
-        value.set(42)
+        abi.String().set(42)
 
 
 def test_String_set_expr():
@@ -175,18 +153,23 @@ def test_String_set_expr():
         assert expr.type_of() == pt.TealType.none
         assert not expr.has_return()
 
-        vts, _ = value_to_set.__teal__(options)
-        expected = pt.TealSimpleBlock(
+        value_start, value_end = value_to_set.__teal__(options)
+        expected_body = pt.TealSimpleBlock(
             [
-                vts.ops[0],
+                pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+                pt.TealOp(None, pt.Op.load, value.stored_value.slot),
                 pt.TealOp(None, pt.Op.len),
                 pt.TealOp(None, pt.Op.itob),
                 pt.TealOp(None, pt.Op.extract, 6, 0),
-                vts.ops[0],
+                pt.TealOp(None, pt.Op.load, value.stored_value.slot),
                 pt.TealOp(None, pt.Op.concat),
                 pt.TealOp(None, pt.Op.store, value.stored_value.slot),
             ]
         )
+        value_end.setNextBlock(expected_body)
+        expected = value_start
+        expected.addIncoming()
+        expected = pt.TealBlock.NormalizeBlocks(expected)
 
         actual, _ = expr.__teal__(options)
         actual.addIncoming()

pyteal/compiler/compiler_test.py

@@ -2813,12 +2813,7 @@ def approve_if_odd(condition_encoding: pt.abi.Uint32) -> pt.Expr:
 
 // log_creation
 logcreation_8:
-byte "logging creation"
-len
-itob
-extract 6 0
-byte "logging creation"
-concat
+byte 0x00106c6f6767696e67206372656174696f6e
 store 68
 load 68
 retsub""".strip()
@@ -3341,12 +3336,7 @@ def approve_if_odd(condition_encoding: pt.abi.Uint32) -> pt.Expr:
 
 // log_creation
 logcreation_8:
-byte "logging creation"
-len
-itob
-extract 6 0
-byte "logging creation"
-concat
+byte 0x00106c6f6767696e67206372656174696f6e
 store 68
 load 68
 retsub""".strip()