Skip to content

Commit 155cf7f

Browse files
jasonpaulosjasonpaulos
committed
Revert "Revert "String optimization and addition of Suffix() (#126)""
This reverts commit 564e602.
1 parent eb8d8b1 commit 155cf7f

File tree

5 files changed

+597
-121
lines changed

5 files changed

+597
-121
lines changed

‎pyteal/ast/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@
8888
)
8989

9090
# ternary ops
91-
from .ternaryexpr import Ed25519Verify, Substring, Extract, SetBit, SetByte
91+
from .ternaryexpr import Ed25519Verify, SetBit, SetByte
92+
from .substring import Substring, Extract, Suffix
9293

9394
# more ops
9495
from .naryexpr import NaryExpr, And, Or, Concat
@@ -189,6 +190,7 @@
189190
"Ed25519Verify",
190191
"Substring",
191192
"Extract",
193+
"Suffix",
192194
"SetBit",
193195
"SetByte",
194196
"NaryExpr",

‎pyteal/ast/substring.py‎

Lines changed: 309 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,309 @@
1+
from enum import Enum
2+
from typing import cast, Tuple, TYPE_CHECKING
3+
4+
from ..types import TealType, require_type
5+
from ..errors import TealCompileError, verifyTealVersion
6+
from ..ir import TealOp, Op, TealBlock, TealSimpleBlock
7+
from .expr import Expr
8+
from .int import Int
9+
from .ternaryexpr import TernaryExpr
10+
11+
if TYPE_CHECKING:
12+
from ..compiler import CompileOptions
13+
14+
15+
class SubstringExpr(Expr):
16+
"""An expression for taking the substring of a byte string given start and end indices"""
17+
18+
def __init__(self, stringArg: Expr, startArg: Expr, endArg: Expr) -> None:
19+
super().__init__()
20+
21+
require_type(stringArg.type_of(), TealType.bytes)
22+
require_type(startArg.type_of(), TealType.uint64)
23+
require_type(endArg.type_of(), TealType.uint64)
24+
25+
self.stringArg = stringArg
26+
self.startArg = startArg
27+
self.endArg = endArg
28+
29+
# helper method for correctly populating op
30+
def __getOp(self, options: "CompileOptions"):
31+
s, e = cast(Int, self.startArg).value, cast(Int, self.endArg).value
32+
l = e - s
33+
34+
if l < 0:
35+
raise TealCompileError(
36+
"The end index must be greater than or equal to the start index",
37+
self,
38+
)
39+
40+
if l > 0 and options.version >= Op.extract.min_version:
41+
if s < 2 ** 8 and l < 2 ** 8:
42+
return Op.extract
43+
else:
44+
return Op.extract3
45+
else:
46+
if s < 2 ** 8 and e < 2 ** 8:
47+
return Op.substring
48+
else:
49+
return Op.substring3
50+
51+
def __teal__(self, options: "CompileOptions"):
52+
if not isinstance(self.startArg, Int) or not isinstance(self.endArg, Int):
53+
return TernaryExpr(
54+
Op.substring3,
55+
(TealType.bytes, TealType.uint64, TealType.uint64),
56+
TealType.bytes,
57+
self.stringArg,
58+
self.startArg,
59+
self.endArg,
60+
).__teal__(options)
61+
62+
op = self.__getOp(options)
63+
64+
verifyTealVersion(
65+
op.min_version,
66+
options.version,
67+
"TEAL version too low to use op {}".format(op),
68+
)
69+
70+
start, end = cast(Int, self.startArg).value, cast(Int, self.endArg).value
71+
if op == Op.extract:
72+
length = end - start
73+
return TealBlock.FromOp(
74+
options,
75+
TealOp(self, op, self.startArg.value, length),
76+
self.stringArg,
77+
)
78+
elif op == Op.extract3:
79+
length = end - start
80+
return TealBlock.FromOp(
81+
options,
82+
TealOp(self, op),
83+
self.stringArg,
84+
self.startArg,
85+
Int(length),
86+
)
87+
elif op == Op.substring:
88+
return TealBlock.FromOp(
89+
options, TealOp(self, op, start, end), self.stringArg
90+
)
91+
elif op == Op.substring3:
92+
return TealBlock.FromOp(
93+
options,
94+
TealOp(self, op),
95+
self.stringArg,
96+
self.startArg,
97+
self.endArg,
98+
)
99+
100+
def __str__(self):
101+
return "(Substring {} {} {})".format(self.stringArg, self.startArg, self.endArg)
102+
103+
def type_of(self):
104+
return TealType.bytes
105+
106+
def has_return(self):
107+
return False
108+
109+
110+
class ExtractExpr(Expr):
111+
"""An expression for extracting a section of a byte string given a start index and length"""
112+
113+
def __init__(self, stringArg: Expr, startArg: Expr, lenArg: Expr) -> None:
114+
super().__init__()
115+
116+
require_type(stringArg.type_of(), TealType.bytes)
117+
require_type(startArg.type_of(), TealType.uint64)
118+
require_type(lenArg.type_of(), TealType.uint64)
119+
120+
self.stringArg = stringArg
121+
self.startArg = startArg
122+
self.lenArg = lenArg
123+
124+
# helper method for correctly populating op
125+
def __getOp(self, options: "CompileOptions"):
126+
s, l = cast(Int, self.startArg).value, cast(Int, self.lenArg).value
127+
if s < 2 ** 8 and l > 0 and l < 2 ** 8:
128+
return Op.extract
129+
else:
130+
return Op.extract3
131+
132+
def __teal__(self, options: "CompileOptions"):
133+
if not isinstance(self.startArg, Int) or not isinstance(self.lenArg, Int):
134+
return TernaryExpr(
135+
Op.extract3,
136+
(TealType.bytes, TealType.uint64, TealType.uint64),
137+
TealType.bytes,
138+
self.stringArg,
139+
self.startArg,
140+
self.lenArg,
141+
).__teal__(options)
142+
143+
op = self.__getOp(options)
144+
145+
verifyTealVersion(
146+
op.min_version,
147+
options.version,
148+
"TEAL version too low to use op {}".format(op),
149+
)
150+
151+
s, l = cast(Int, self.startArg).value, cast(Int, self.lenArg).value
152+
if op == Op.extract:
153+
return TealBlock.FromOp(options, TealOp(self, op, s, l), self.stringArg)
154+
elif op == Op.extract3:
155+
return TealBlock.FromOp(
156+
options,
157+
TealOp(self, op),
158+
self.stringArg,
159+
self.startArg,
160+
self.lenArg,
161+
)
162+
163+
def __str__(self):
164+
return "(Extract {} {} {})".format(self.stringArg, self.startArg, self.lenArg)
165+
166+
def type_of(self):
167+
return TealType.bytes
168+
169+
def has_return(self):
170+
return False
171+
172+
173+
class SuffixExpr(Expr):
174+
"""An expression for taking the suffix of a byte string given start index"""
175+
176+
def __init__(
177+
self,
178+
stringArg: Expr,
179+
startArg: Expr,
180+
) -> None:
181+
super().__init__()
182+
183+
require_type(stringArg.type_of(), TealType.bytes)
184+
require_type(startArg.type_of(), TealType.uint64)
185+
186+
self.stringArg = stringArg
187+
self.startArg = startArg
188+
189+
# helper method for correctly populating op
190+
def __getOp(self, options: "CompileOptions"):
191+
if not isinstance(self.startArg, Int):
192+
return Op.substring3
193+
194+
s = cast(Int, self.startArg).value
195+
if s < 2 ** 8:
196+
return Op.extract
197+
else:
198+
return Op.substring3
199+
200+
def __teal__(self, options: "CompileOptions"):
201+
op = self.__getOp(options)
202+
203+
verifyTealVersion(
204+
op.min_version,
205+
options.version,
206+
"TEAL version too low to use op {}".format(op),
207+
)
208+
209+
if op == Op.extract:
210+
# if possible, exploit optimization in the extract opcode that takes the suffix
211+
# when the length argument is 0
212+
return TealBlock.FromOp(
213+
options,
214+
TealOp(self, op, cast(Int, self.startArg).value, 0),
215+
self.stringArg,
216+
)
217+
elif op == Op.substring3:
218+
strBlockStart, strBlockEnd = self.stringArg.__teal__(options)
219+
nextBlockStart, nextBlockEnd = self.startArg.__teal__(options)
220+
strBlockEnd.setNextBlock(nextBlockStart)
221+
222+
finalBlock = TealSimpleBlock(
223+
[
224+
TealOp(self, Op.dig, 1),
225+
TealOp(self, Op.len),
226+
TealOp(self, Op.substring3),
227+
]
228+
)
229+
230+
nextBlockEnd.setNextBlock(finalBlock)
231+
return strBlockStart, finalBlock
232+
233+
def __str__(self):
234+
return "(Suffix {} {})".format(self.stringArg, self.startArg)
235+
236+
def type_of(self):
237+
return TealType.bytes
238+
239+
def has_return(self):
240+
return False
241+
242+
243+
def Substring(string: Expr, start: Expr, end: Expr) -> Expr:
244+
"""Take a substring of a byte string.
245+
246+
Produces a new byte string consisting of the bytes starting at :code:`start` up to but not
247+
including :code:`end`.
248+
249+
This expression is similar to :any:`Extract`, except this expression uses start and end indexes,
250+
while :code:`Extract` uses a start index and length.
251+
252+
Requires TEAL version 2 or higher.
253+
254+
Args:
255+
string: The byte string.
256+
start: The starting index for the substring. Must be an integer less than or equal to
257+
:code:`Len(string)`.
258+
end: The ending index for the substring. Must be an integer greater or equal to start, but
259+
less than or equal to Len(string).
260+
"""
261+
return SubstringExpr(
262+
string,
263+
start,
264+
end,
265+
)
266+
267+
268+
def Extract(string: Expr, start: Expr, length: Expr) -> Expr:
269+
"""Extract a section of a byte string.
270+
271+
Produces a new byte string consisting of the bytes starting at :code:`start` up to but not
272+
including :code:`start + length`.
273+
274+
This expression is similar to :any:`Substring`, except this expression uses a start index and
275+
length, while :code:`Substring` uses start and end indexes.
276+
277+
Requires TEAL version 5 or higher.
278+
279+
Args:
280+
string: The byte string.
281+
start: The starting index for the extraction. Must be an integer less than or equal to
282+
:code:`Len(string)`.
283+
length: The number of bytes to extract. Must be an integer such that :code:`start + length <= Len(string)`.
284+
"""
285+
return ExtractExpr(
286+
string,
287+
start,
288+
length,
289+
)
290+
291+
292+
def Suffix(string: Expr, start: Expr) -> Expr:
293+
"""Take a suffix of a byte string.
294+
295+
Produces a new byte string consisting of the suffix of the byte string starting at :code:`start`
296+
297+
This expression is similar to :any:`Substring` and :any:`Extract`, except this expression only uses a
298+
start index.
299+
300+
Requires TEAL version 5 or higher.
301+
302+
Args:
303+
string: The byte string.
304+
start: The starting index for the suffix. Must be an integer less than or equal to :code:`Len(string)`.
305+
"""
306+
return SuffixExpr(
307+
string,
308+
start,
309+
)

0 commit comments

Comments
 (0)