diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py index 9712dc76d..fa3697da3 100644 --- a/pyteal/ast/__init__.py +++ b/pyteal/ast/__init__.py @@ -88,8 +88,7 @@ ) # ternary ops -from .ternaryexpr import Ed25519Verify, SetBit, SetByte -from .substring import Substring, Extract, Suffix +from .ternaryexpr import Ed25519Verify, Substring, Extract, SetBit, SetByte # more ops from .naryexpr import NaryExpr, And, Or, Concat @@ -190,7 +189,6 @@ "Ed25519Verify", "Substring", "Extract", - "Suffix", "SetBit", "SetByte", "NaryExpr", diff --git a/pyteal/ast/substring.py b/pyteal/ast/substring.py deleted file mode 100644 index a36ed0722..000000000 --- a/pyteal/ast/substring.py +++ /dev/null @@ -1,309 +0,0 @@ -from enum import Enum -from typing import cast, Tuple, TYPE_CHECKING - -from ..types import TealType, require_type -from ..errors import TealCompileError, verifyTealVersion -from ..ir import TealOp, Op, TealBlock, TealSimpleBlock -from .expr import Expr -from .int import Int -from .ternaryexpr import TernaryExpr - -if TYPE_CHECKING: - from ..compiler import CompileOptions - - -class SubstringExpr(Expr): - """An expression for taking the substring of a byte string given start and end indices""" - - def __init__(self, stringArg: Expr, startArg: Expr, endArg: Expr) -> None: - super().__init__() - - require_type(stringArg.type_of(), TealType.bytes) - require_type(startArg.type_of(), TealType.uint64) - require_type(endArg.type_of(), TealType.uint64) - - self.stringArg = stringArg - self.startArg = startArg - self.endArg = endArg - - # helper method for correctly populating op - def __getOp(self, options: "CompileOptions"): - s, e = cast(Int, self.startArg).value, cast(Int, self.endArg).value - l = e - s - - if l < 0: - raise TealCompileError( - "The end index must be greater than or equal to the start index", - self, - ) - - if l > 0 and options.version >= Op.extract.min_version: - if s < 2 ** 8 and l < 2 ** 8: - return Op.extract - else: - return Op.extract3 - else: - if s < 2 ** 8 and e < 2 ** 8: - return Op.substring - else: - return Op.substring3 - - def __teal__(self, options: "CompileOptions"): - if not isinstance(self.startArg, Int) or not isinstance(self.endArg, Int): - return TernaryExpr( - Op.substring3, - (TealType.bytes, TealType.uint64, TealType.uint64), - TealType.bytes, - self.stringArg, - self.startArg, - self.endArg, - ).__teal__(options) - - op = self.__getOp(options) - - verifyTealVersion( - op.min_version, - options.version, - "TEAL version too low to use op {}".format(op), - ) - - start, end = cast(Int, self.startArg).value, cast(Int, self.endArg).value - if op == Op.extract: - length = end - start - return TealBlock.FromOp( - options, - TealOp(self, op, self.startArg.value, length), - self.stringArg, - ) - elif op == Op.extract3: - length = end - start - return TealBlock.FromOp( - options, - TealOp(self, op), - self.stringArg, - self.startArg, - Int(length), - ) - elif op == Op.substring: - return TealBlock.FromOp( - options, TealOp(self, op, start, end), self.stringArg - ) - elif op == Op.substring3: - return TealBlock.FromOp( - options, - TealOp(self, op), - self.stringArg, - self.startArg, - self.endArg, - ) - - def __str__(self): - return "(Substring {} {} {})".format(self.stringArg, self.startArg, self.endArg) - - def type_of(self): - return TealType.bytes - - def has_return(self): - return False - - -class ExtractExpr(Expr): - """An expression for extracting a section of a byte string given a start index and length""" - - def __init__(self, stringArg: Expr, startArg: Expr, lenArg: Expr) -> None: - super().__init__() - - require_type(stringArg.type_of(), TealType.bytes) - require_type(startArg.type_of(), TealType.uint64) - require_type(lenArg.type_of(), TealType.uint64) - - self.stringArg = stringArg - self.startArg = startArg - self.lenArg = lenArg - - # helper method for correctly populating op - def __getOp(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 - else: - return Op.extract3 - - def __teal__(self, options: "CompileOptions"): - if not isinstance(self.startArg, Int) or not isinstance(self.lenArg, Int): - return TernaryExpr( - Op.extract3, - (TealType.bytes, TealType.uint64, TealType.uint64), - TealType.bytes, - self.stringArg, - self.startArg, - self.lenArg, - ).__teal__(options) - - op = self.__getOp(options) - - verifyTealVersion( - op.min_version, - options.version, - "TEAL version too low to use op {}".format(op), - ) - - s, l = cast(Int, self.startArg).value, cast(Int, self.lenArg).value - if op == Op.extract: - return TealBlock.FromOp(options, TealOp(self, op, s, l), self.stringArg) - elif op == Op.extract3: - return TealBlock.FromOp( - options, - TealOp(self, op), - self.stringArg, - self.startArg, - self.lenArg, - ) - - def __str__(self): - return "(Extract {} {} {})".format(self.stringArg, self.startArg, self.lenArg) - - def type_of(self): - return TealType.bytes - - def has_return(self): - return False - - -class SuffixExpr(Expr): - """An expression for taking the suffix of a byte string given start index""" - - def __init__( - self, - stringArg: Expr, - startArg: Expr, - ) -> None: - super().__init__() - - require_type(stringArg.type_of(), TealType.bytes) - require_type(startArg.type_of(), TealType.uint64) - - self.stringArg = stringArg - self.startArg = startArg - - # helper method for correctly populating op - def __getOp(self, options: "CompileOptions"): - if not isinstance(self.startArg, Int): - return Op.substring3 - - s = cast(Int, self.startArg).value - if s < 2 ** 8: - return Op.extract - else: - return Op.substring3 - - def __teal__(self, options: "CompileOptions"): - op = self.__getOp(options) - - verifyTealVersion( - op.min_version, - options.version, - "TEAL version too low to use op {}".format(op), - ) - - if op == Op.extract: - # if possible, exploit optimization in the extract opcode that takes the suffix - # when the length argument is 0 - return TealBlock.FromOp( - options, - TealOp(self, op, cast(Int, self.startArg).value, 0), - self.stringArg, - ) - elif op == Op.substring3: - strBlockStart, strBlockEnd = self.stringArg.__teal__(options) - nextBlockStart, nextBlockEnd = self.startArg.__teal__(options) - strBlockEnd.setNextBlock(nextBlockStart) - - finalBlock = TealSimpleBlock( - [ - TealOp(self, Op.dig, 1), - TealOp(self, Op.len), - TealOp(self, Op.substring3), - ] - ) - - nextBlockEnd.setNextBlock(finalBlock) - return strBlockStart, finalBlock - - def __str__(self): - return "(Suffix {} {})".format(self.stringArg, self.startArg) - - def type_of(self): - return TealType.bytes - - def has_return(self): - return False - - -def Substring(string: Expr, start: Expr, end: Expr) -> Expr: - """Take a substring of a byte string. - - Produces a new byte string consisting of the bytes starting at :code:`start` up to but not - including :code:`end`. - - This expression is similar to :any:`Extract`, except this expression uses start and end indexes, - while :code:`Extract` uses a start index and length. - - Requires TEAL version 2 or higher. - - Args: - string: The byte string. - start: The starting index for the substring. Must be an integer less than or equal to - :code:`Len(string)`. - end: The ending index for the substring. Must be an integer greater or equal to start, but - less than or equal to Len(string). - """ - return SubstringExpr( - string, - start, - end, - ) - - -def Extract(string: Expr, start: Expr, length: Expr) -> Expr: - """Extract a section of a byte string. - - Produces a new byte string consisting of the bytes starting at :code:`start` up to but not - including :code:`start + length`. - - This expression is similar to :any:`Substring`, except this expression uses a start index and - length, while :code:`Substring` uses start and end indexes. - - Requires TEAL version 5 or higher. - - Args: - string: The byte string. - start: The starting index for the extraction. Must be an integer less than or equal to - :code:`Len(string)`. - length: The number of bytes to extract. Must be an integer such that :code:`start + length <= Len(string)`. - """ - return ExtractExpr( - string, - start, - length, - ) - - -def Suffix(string: Expr, start: Expr) -> Expr: - """Take a suffix of a byte string. - - Produces a new byte string consisting of the suffix of the byte string starting at :code:`start` - - This expression is similar to :any:`Substring` and :any:`Extract`, except this expression only uses a - start index. - - Requires TEAL version 5 or higher. - - Args: - string: The byte string. - start: The starting index for the suffix. Must be an integer less than or equal to :code:`Len(string)`. - """ - return SuffixExpr( - string, - start, - ) diff --git a/pyteal/ast/substring_test.py b/pyteal/ast/substring_test.py deleted file mode 100644 index c532f299f..000000000 --- a/pyteal/ast/substring_test.py +++ /dev/null @@ -1,283 +0,0 @@ -import pytest - -from .. import * - -# this is not necessary but mypy complains if it's not included -from .. import CompileOptions - -teal2Options = CompileOptions(version=2) -teal3Options = CompileOptions(version=3) -teal4Options = CompileOptions(version=4) -teal5Options = CompileOptions(version=5) - - -def test_substring_immediate_v2(): - args = [Bytes("my string"), Int(0), Int(2)] - expr = Substring(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"my string"'), - TealOp(expr, Op.substring, 0, 2), - ] - ) - - actual, _ = expr.__teal__(teal2Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - assert actual == expected - - -def test_substring_immediate_v5(): - args = [Bytes("my string"), Int(1), Int(2)] - expr = Substring(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"my string"'), - TealOp(expr, Op.extract, 1, 1), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - assert actual == expected - - -def test_substring_to_extract(): - my_string = "a" * 257 - args = [Bytes(my_string), Int(255), Int(257)] - expr = Substring(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"{my_string}"'.format(my_string=my_string)), - TealOp(expr, Op.extract, 255, 2), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - with TealComponent.Context.ignoreExprEquality(): - assert actual == expected - - -def test_substring_stack_v2(): - my_string = "a" * 257 - args = [Bytes(my_string), Int(256), Int(257)] - expr = Substring(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"{my_string}"'.format(my_string=my_string)), - TealOp(args[1], Op.int, 256), - TealOp(args[2], Op.int, 257), - TealOp(expr, Op.substring3), - ] - ) - - actual, _ = expr.__teal__(teal2Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - with TealComponent.Context.ignoreExprEquality(): - assert actual == expected - - -def test_substring_stack_v5(): - my_string = "a" * 257 - args = [Bytes(my_string), Int(256), Int(257)] - expr = Substring(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"{my_string}"'.format(my_string=my_string)), - TealOp(args[1], Op.int, 256), - TealOp(Int(1), Op.int, 1), - TealOp(expr, Op.extract3), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - with TealComponent.Context.ignoreExprEquality(): - assert actual == expected - - -def test_zero_length_substring_immediate(): - my_string = "a" * 257 - args = [Bytes(my_string), Int(1), Int(1)] - expr = Substring(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"{my_string}"'.format(my_string=my_string)), - TealOp(expr, Op.substring, 1, 1), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - with TealComponent.Context.ignoreExprEquality(): - assert actual == expected - - -def test_substring_invalid(): - with pytest.raises(TealTypeError): - Substring(Int(0), Int(0), Int(2)) - - with pytest.raises(TealTypeError): - Substring(Bytes("my string"), Txn.sender(), Int(2)) - - with pytest.raises(TealTypeError): - Substring(Bytes("my string"), Int(0), Txn.sender()) - - with pytest.raises(Exception): - Substring(Bytes("my string"), Int(1), Int(0)).__teal__(teal5Options) - - -def test_extract_immediate(): - args = [Bytes("my string"), Int(0), Int(2)] - expr = Extract(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"my string"'), - TealOp(expr, Op.extract, 0, 2), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - assert actual == expected - - with pytest.raises(TealInputError): - expr.__teal__(teal4Options) - - -def test_extract_zero(): - args = [Bytes("my string"), Int(1), Int(0)] - expr = Extract(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"my string"'), - TealOp(args[1], Op.int, 1), - TealOp(args[2], Op.int, 0), - TealOp(expr, Op.extract3), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - assert actual == expected - - with pytest.raises(TealInputError): - expr.__teal__(teal4Options) - - -def test_extract_stack(): - my_string = "*" * 257 - args = [Bytes(my_string), Int(256), Int(257)] - expr = Extract(args[0], args[1], args[2]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"{my_string}"'.format(my_string=my_string)), - TealOp(args[1], Op.int, 256), - TealOp(args[2], Op.int, 257), - TealOp(expr, Op.extract3), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - assert actual == expected - - with pytest.raises(TealInputError): - expr.__teal__(teal4Options) - - -def test_extract_invalid(): - with pytest.raises(TealTypeError): - Extract(Int(0), Int(0), Int(2)) - - with pytest.raises(TealTypeError): - Extract(Bytes("my string"), Txn.sender(), Int(2)) - - with pytest.raises(TealTypeError): - Extract(Bytes("my string"), Int(0), Txn.sender()) - - -def test_suffix_immediate(): - args = [Bytes("my string"), Int(1)] - expr = Suffix(args[0], args[1]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"my string"'), - TealOp(expr, Op.extract, 1, 0), - ] - ) - - actual, _ = expr.__teal__(teal5Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - assert actual == expected - - -def test_suffix_stack(): - my_string = "*" * 1024 - args = [Bytes(my_string), Int(257)] - expr = Suffix(args[0], args[1]) - assert expr.type_of() == TealType.bytes - - expected = TealSimpleBlock( - [ - TealOp(args[0], Op.byte, '"{my_string}"'.format(my_string=my_string)), - TealOp(args[1], Op.int, 257), - TealOp(expr, Op.dig, 1), - TealOp(expr, Op.len), - TealOp(expr, Op.substring3), - ] - ) - - actual, _ = expr.__teal__(teal2Options) - actual.addIncoming() - actual = TealBlock.NormalizeBlocks(actual) - - assert actual == expected - - -def test_suffix_invalid(): - with pytest.raises(TealTypeError): - Suffix(Int(0), Int(0)) - - with pytest.raises(TealTypeError): - Suffix(Bytes("my string"), Txn.sender()) diff --git a/pyteal/ast/ternaryexpr.py b/pyteal/ast/ternaryexpr.py index 333079faf..2e34fb22f 100644 --- a/pyteal/ast/ternaryexpr.py +++ b/pyteal/ast/ternaryexpr.py @@ -4,8 +4,6 @@ from ..errors import verifyTealVersion from ..ir import TealOp, Op, TealBlock from .expr import Expr -from .int import Int -from .unaryexpr import Len if TYPE_CHECKING: from ..compiler import CompileOptions @@ -79,6 +77,59 @@ def Ed25519Verify(data: Expr, sig: Expr, key: Expr) -> TernaryExpr: ) +def Substring(string: Expr, start: Expr, end: Expr) -> TernaryExpr: + """Take a substring of a byte string. + + Produces a new byte string consisting of the bytes starting at :code:`start` up to but not + including :code:`end`. + + This expression is similar to :any:`Extract`, except this expression uses start and end indexes, + while :code:`Extract` uses a start index and length. + + Args: + string: The byte string. + start: The starting index for the substring. Must be an integer less than or equal to + :code:`Len(string)`. + end: The ending index for the substring. Must be an integer greater or equal to start, but + less than or equal to Len(string). + """ + return TernaryExpr( + Op.substring3, + (TealType.bytes, TealType.uint64, TealType.uint64), + TealType.bytes, + string, + start, + end, + ) + + +def Extract(string: Expr, start: Expr, length: Expr) -> TernaryExpr: + """Extract a section of a byte string. + + Produces a new byte string consisting of the bytes starting at :code:`start` up to but not + including :code:`start + length`. + + This expression is similar to :any:`Substring`, except this expression uses a start index and + length, while :code:`Substring` uses start and end indexes. + + Requires TEAL version 5 or higher. + + Args: + string: The byte string. + start: The starting index for the extraction. Must be an integer less than or equal to + :code:`Len(string)`. + length: The number of bytes to extract. Must be an integer such that :code:`start + length <= Len(string)`. + """ + return TernaryExpr( + Op.extract3, + (TealType.bytes, TealType.uint64, TealType.uint64), + TealType.bytes, + string, + start, + length, + ) + + def SetBit(value: Expr, index: Expr, newBitValue: Expr) -> TernaryExpr: """Set the bit value of an expression at a specific index. diff --git a/pyteal/ast/ternaryexpr_test.py b/pyteal/ast/ternaryexpr_test.py index 22b4817bd..adcf081bb 100644 --- a/pyteal/ast/ternaryexpr_test.py +++ b/pyteal/ast/ternaryexpr_test.py @@ -43,6 +43,73 @@ def test_ed25519verify_invalid(): Ed25519Verify(Bytes("data"), Bytes("sig"), Int(0)) +def test_substring(): + args = [Bytes("my string"), Int(0), Int(2)] + expr = Substring(args[0], args[1], args[2]) + assert expr.type_of() == TealType.bytes + + expected = TealSimpleBlock( + [ + TealOp(args[0], Op.byte, '"my string"'), + TealOp(args[1], Op.int, 0), + TealOp(args[2], Op.int, 2), + TealOp(expr, Op.substring3), + ] + ) + + actual, _ = expr.__teal__(teal2Options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + assert actual == expected + + +def test_substring_invalid(): + with pytest.raises(TealTypeError): + Substring(Int(0), Int(0), Int(2)) + + with pytest.raises(TealTypeError): + Substring(Bytes("my string"), Txn.sender(), Int(2)) + + with pytest.raises(TealTypeError): + Substring(Bytes("my string"), Int(0), Txn.sender()) + + +def test_extract(): + args = [Bytes("my string"), Int(0), Int(2)] + expr = Extract(args[0], args[1], args[2]) + assert expr.type_of() == TealType.bytes + + expected = TealSimpleBlock( + [ + TealOp(args[0], Op.byte, '"my string"'), + TealOp(args[1], Op.int, 0), + TealOp(args[2], Op.int, 2), + TealOp(expr, Op.extract3), + ] + ) + + actual, _ = expr.__teal__(teal5Options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + assert actual == expected + + with pytest.raises(TealInputError): + expr.__teal__(teal4Options) + + +def test_extract_invalid(): + with pytest.raises(TealTypeError): + Extract(Int(0), Int(0), Int(2)) + + with pytest.raises(TealTypeError): + Extract(Bytes("my string"), Txn.sender(), Int(2)) + + with pytest.raises(TealTypeError): + Extract(Bytes("my string"), Int(0), Txn.sender()) + + def test_set_bit_int(): args = [Int(0), Int(2), Int(1)] expr = SetBit(args[0], args[1], args[2])