Add Replace

pyteal/__init__.pyi

@@ -138,6 +138,7 @@ __all__ = [
     "Or",
     "Pop",
     "Reject",
+    "Replace",
     "Return",
     "ScratchIndex",
     "ScratchLoad",

pyteal/ast/__init__.py

@@ -102,6 +102,7 @@
 # ternary ops
 from pyteal.ast.ternaryexpr import Divw, Ed25519Verify, SetBit, SetByte
 from pyteal.ast.substring import Substring, Extract, Suffix
+from pyteal.ast.replace import Replace
 
 # more ops
 from pyteal.ast.naryexpr import NaryExpr, And, Or, Concat
@@ -272,6 +273,7 @@
     "ExtractUint16",
     "ExtractUint32",
     "ExtractUint64",
+    "Replace",
     "Log",
     "While",
     "For",

pyteal/ast/replace.py

@@ -0,0 +1,86 @@
+from typing import cast, TYPE_CHECKING
+
+from pyteal.types import TealType, require_type
+from pyteal.errors import verifyTealVersion
+from pyteal.ir import TealOp, Op, TealBlock
+from pyteal.ast.expr import Expr
+from pyteal.ast.int import Int
+from pyteal.ast.ternaryexpr import TernaryExpr
+
+if TYPE_CHECKING:
+    from pyteal.compiler import CompileOptions
+
+
+class ReplaceExpr(Expr):
+    """An expression for replacing a section of a byte string at a given start index"""
+
+    def __init__(self, original: Expr, start: Expr, replacement: Expr) -> None:
+        super().__init__()
+
+        require_type(original, TealType.bytes)
+        require_type(start, TealType.uint64)
+        require_type(replacement, TealType.bytes)
+
+        self.original = original
+        self.start = start
+        self.replacement = replacement
+
+    # helper method for correctly populating op
+    def __get_op(self, options: "CompileOptions"):
+        s = cast(Int, self.start).value
+        if s < 2**8:
+            return Op.replace2
+        else:
+            return Op.replace3
+
+    def __teal__(self, options: "CompileOptions"):
+        if not isinstance(self.start, Int):
+            return TernaryExpr(
+                Op.replace3,
+                (TealType.bytes, TealType.uint64, TealType.bytes),
+                TealType.bytes,
+                self.original,
+                self.start,
+                self.replacement,
+            ).__teal__(options)
+
+        op = self.__get_op(options)
+
+        verifyTealVersion(
+            op.min_version,
+            options.version,
+            "TEAL version too low to use op {}".format(op),
+        )
+
+        s = cast(Int, self.start).value
+        if op == Op.replace2:
+            return TealBlock.FromOp(
+                options, TealOp(self, op, s), self.original, self.replacement
+            )
+        elif op == Op.replace3:
+            return TealBlock.FromOp(
+                options, TealOp(self, op), self.original, self.start, self.replacement
+            )
+
+    def __str__(self):
+        return "(Replace {} {} {})".format(self.original, self.start, self.replacement)
+
+    def type_of(self):
+        return TealType.bytes
+
+    def has_return(self):
+        return False
+
+
+def Replace(original: Expr, start: Expr, replacement: Expr) -> Expr:
+    """
+    Replace a portion of original bytes with new bytes at a given starting point.
+
+    Requires TEAL version 7 or higher.
+
+    Args:
+        original: The value containing the original bytes. Must evaluate to bytes.
+        start: The index of the byte where replacement starts. Must evaluate to an integer less than Len(original).
+        replacement: The value containing the replacement bytes. Must evaluate to bytes with length at most Len(original) - start.
+    """
+    return ReplaceExpr(original, start, replacement)

pyteal/ast/replace_test.py

@@ -0,0 +1,97 @@
+import pytest
+
+import pyteal as pt
+
+teal6Options = pt.CompileOptions(version=6)
+teal7Options = pt.CompileOptions(version=7)
+
+
+def test_replace_immediate():
+    args = [pt.Bytes("my string"), pt.Int(0), pt.Bytes("abcdefghi")]
+    expr = pt.Replace(args[0], args[1], args[2])
+    assert expr.type_of() == pt.TealType.bytes
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"my string"'),
+            pt.TealOp(args[2], pt.Op.byte, '"abcdefghi"'),
+            pt.TealOp(expr, pt.Op.replace2, 0),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal7Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        expr.__teal__(teal6Options)
+
+
+def test_replace_stack_int():
+    my_string = "*" * 257
+    args = [pt.Bytes(my_string), pt.Int(256), pt.Bytes("ab")]
+    expr = pt.Replace(args[0], args[1], args[2])
+    assert expr.type_of() == pt.TealType.bytes
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
+            pt.TealOp(args[1], pt.Op.int, 256),
+            pt.TealOp(args[2], pt.Op.byte, '"ab"'),
+            pt.TealOp(expr, pt.Op.replace3),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal7Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        expr.__teal__(teal6Options)
+
+
+# Mirrors `test_replace_stack_int`, but attempts replacement with start != pt.Int.
+def test_replace_stack_not_int():
+    my_string = "*" * 257
+    add = pt.Add(pt.Int(254), pt.Int(2))
+    args = [pt.Bytes(my_string), add, pt.Bytes("ab")]
+    expr = pt.Replace(args[0], args[1], args[2])
+    assert expr.type_of() == pt.TealType.bytes
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
+            pt.TealOp(pt.Int(254), pt.Op.int, 254),
+            pt.TealOp(pt.Int(2), pt.Op.int, 2),
+            pt.TealOp(add, pt.Op.add),
+            pt.TealOp(args[2], pt.Op.byte, '"ab"'),
+            pt.TealOp(expr, pt.Op.replace3),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal7Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        expr.__teal__(teal6Options)
+
+
+def test_replace_invalid():
+    with pytest.raises(pt.TealTypeError):
+        pt.Replace(pt.Bytes("my string"), pt.Int(0), pt.Int(1))
+
+    with pytest.raises(pt.TealTypeError):
+        pt.Replace(
+            pt.Bytes("my string"), pt.Bytes("should be int"), pt.Bytes("abcdefghi")
+        )
+
+    with pytest.raises(pt.TealTypeError):
+        pt.Replace(pt.Bytes("my string"), pt.Txn.sender(), pt.Bytes("abcdefghi"))

pyteal/ast/substring.py

@@ -26,7 +26,7 @@ def __init__(self, stringArg: Expr, startArg: Expr, endArg: Expr) -> None:
         self.endArg = endArg
 
     # helper method for correctly populating op
-    def __getOp(self, options: "CompileOptions"):
+    def __get_op(self, options: "CompileOptions"):
         s, e = cast(Int, self.startArg).value, cast(Int, self.endArg).value
         l = e - s
 
@@ -58,7 +58,7 @@ def __teal__(self, options: "CompileOptions"):
                 self.endArg,
             ).__teal__(options)
 
-        op = self.__getOp(options)
+        op = self.__get_op(options)
 
         verifyTealVersion(
             op.min_version,
@@ -121,7 +121,7 @@ def __init__(self, stringArg: Expr, startArg: Expr, lenArg: Expr) -> None:
         self.lenArg = lenArg
 
     # helper method for correctly populating op
-    def __getOp(self, options: "CompileOptions"):
+    def __get_op(self, options: "CompileOptions"):
         s, l = cast(Int, self.startArg).value, cast(Int, self.lenArg).value
         if s < 2**8 and l > 0 and l < 2**8:
             return Op.extract
@@ -139,7 +139,7 @@ def __teal__(self, options: "CompileOptions"):
                 self.lenArg,
             ).__teal__(options)
 
-        op = self.__getOp(options)
+        op = self.__get_op(options)
 
         verifyTealVersion(
             op.min_version,
@@ -186,7 +186,7 @@ def __init__(
         self.startArg = startArg
 
     # helper method for correctly populating op
-    def __getOp(self, options: "CompileOptions"):
+    def __get_op(self, options: "CompileOptions"):
         if not isinstance(self.startArg, Int):
             return Op.substring3
 
@@ -197,7 +197,7 @@ def __getOp(self, options: "CompileOptions"):
             return Op.substring3
 
     def __teal__(self, options: "CompileOptions"):
-        op = self.__getOp(options)
+        op = self.__get_op(options)
 
         verifyTealVersion(
             op.min_version,

pyteal/ast/substring_test.py

@@ -272,6 +272,88 @@ def test_suffix_stack():
     assert actual == expected
 
 
+@pytest.mark.parametrize("op", [pt.Op.extract3, pt.Op.substring3])
+def test_startArg_not_int(op: pt.Op):
+    my_string = "*" * 257
+    add = pt.Add(pt.Int(254), pt.Int(2))
+    args = [pt.Bytes(my_string), add, pt.Int(257)]
+
+    def generate_expr() -> pt.Expr:
+        match op:
+            case pt.Op.extract3:
+                return pt.Extract(args[0], args[1], args[2])
+            case pt.Op.substring3:
+                return pt.Substring(args[0], args[1], args[2])
+            case _:
+                raise Exception(f"Unsupported {op=}")
+
+    expr = generate_expr()
+    assert expr.type_of() == pt.TealType.bytes
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
+            pt.TealOp(pt.Int(254), pt.Op.int, 254),
+            pt.TealOp(pt.Int(2), pt.Op.int, 2),
+            pt.TealOp(add, pt.Op.add),
+            pt.TealOp(args[2], pt.Op.int, 257),
+            pt.TealOp(None, op),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal5Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+    if op == pt.Op.extract3:
+        with pytest.raises(pt.TealInputError):
+            expr.__teal__(teal4Options)
+
+
+@pytest.mark.parametrize("op", [pt.Op.extract3, pt.Op.substring3])
+def test_endArg_not_int(op: pt.Op):
+    my_string = "*" * 257
+    add = pt.Add(pt.Int(254), pt.Int(3))
+    args = [pt.Bytes(my_string), pt.Int(256), add]
+
+    def generate_expr() -> pt.Expr:
+        match op:
+            case pt.Op.extract3:
+                return pt.Extract(args[0], args[1], args[2])
+            case pt.Op.substring3:
+                return pt.Substring(args[0], args[1], args[2])
+            case _:
+                raise Exception(f"Unsupported {op=}")
+
+    expr = generate_expr()
+    assert expr.type_of() == pt.TealType.bytes
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(args[0], pt.Op.byte, '"{my_string}"'.format(my_string=my_string)),
+            pt.TealOp(args[1], pt.Op.int, 256),
+            pt.TealOp(pt.Int(254), pt.Op.int, 254),
+            pt.TealOp(pt.Int(3), pt.Op.int, 3),
+            pt.TealOp(add, pt.Op.add),
+            pt.TealOp(None, op),
+        ]
+    )
+
+    actual, _ = expr.__teal__(teal5Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+    if op == pt.Op.extract3:
+        with pytest.raises(pt.TealInputError):
+            expr.__teal__(teal4Options)
+
+
 def test_suffix_invalid():
     with pytest.raises(pt.TealTypeError):
         pt.Suffix(pt.Int(0), pt.Int(0))

pyteal/ir/ops.py

@@ -179,6 +179,8 @@ def min_version(self) -> int:
     gitxnas             = OpType("gitxnas",             Mode.Application,                  6)
     gloadss             = OpType("gloadss",             Mode.Application,                  6)
     acct_params_get     = OpType("acct_params_get",     Mode.Application,                  6)
+    replace2            = OpType("replace2",            Mode.Signature | Mode.Application, 7)
+    replace3            = OpType("replace3",            Mode.Signature | Mode.Application, 7)
     # fmt: on