Skip to content

Commit

Permalink
feat: allow TEAL optimizer to handle dup/dupn ops when removing stack…
Browse files Browse the repository at this point in the history 
… shuffling of constants, and collapse any repeated elements with a dup/dupn
  • Loading branch information
@achidlow
achidlow committed Mar 15, 2024
1 parent 3bbe933 commit b48db43
Show file tree
Hide file tree
Showing 54 changed files with 135 additions and 123 deletions.
2 changes: 1 addition & 1 deletion examples/amm/out/ConstantProductAMM.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -390,7 +390,7 @@ _create_pool_token:
global CurrentApplicationAddress
// amm/contract.py:279
// reserve=Global.current_application_address,
global CurrentApplicationAddress
dup
// amm/contract.py:280
// fee=0,
int 0
Expand Down
2 changes: 1 addition & 1 deletion examples/amm/out/ConstantProductAMM.arc32.json
View file

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion examples/amm/out_O2/ConstantProductAMM.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -257,7 +257,7 @@ _create_pool_token:
assert // asset exists
concat
global CurrentApplicationAddress
global CurrentApplicationAddress
dup
int 0
itxn_field Fee
itxn_field ConfigAssetReserve
Expand Down
2 changes: 1 addition & 1 deletion examples/calculator/out/MyContract.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ examples.calculator.contract.MyContract.approval_program:
int 0
// calculator/contract.py:24
// action = UInt64(0)
int 0
dup
bury 3
// calculator/contract.py:22
// a = UInt64(0)
Expand Down
2 changes: 1 addition & 1 deletion examples/calculator/out_O2/MyContract.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ examples.calculator.contract.MyContract.approval_program:
int 0
bury 3
int 0
int 0
dup
bury 3
int 0
itob
Expand Down
16 changes: 8 additions & 8 deletions examples/sizes.txt
View file
Original file line number Diff line number Diff line change
@@ -1,17 +1,17 @@
Name O0 size O1 size O1 ⏷ O2 size O2 ⏷
abi_routing/Reference 1179 1019 160 1019 0
amm/ConstantProductAMM 1213 1114 99 1114 0
amm/ConstantProductAMM 1213 1113 100 1113 0
application/Reference 175 168 7 168 0
arc4_numeric_comparisons/UIntNOrdering 1220 908 312 908 0
arc4_types/Arc4Arrays 588 376 212 376 0
arc4_types/Arc4BoolEval 731 20 711 20 0
arc4_types/Arc4BoolEval 731 19 712 19 0
arc4_types/Arc4BoolType 329 57 272 57 0
arc4_types/Arc4DynamicBytes 247 128 119 128 0
arc4_types/Arc4DynamicStringArray 230 112 118 112 0
arc4_types/Arc4MutableParams 362 222 140 220 2
arc4_types/Arc4Mutation 2803 1452 1351 1451 1
arc4_types/Arc4NumericTypes 538 8 530 8 0
arc4_types/Arc4RefTypes 47 43 4 43 0
arc4_types/Arc4RefTypes 47 39 8 39 0
arc4_types/Arc4StringTypes 309 8 301 8 0
arc4_types/Arc4StructsFromAnotherModule 67 12 55 12 0
arc4_types/Arc4StructsType 296 237 59 237 0
Expand All @@ -21,7 +21,7 @@
augmented_assignment/Augmented 157 156 1 156 0
avm_types_in_abi/Test 226 173 53 173 0
biguint_binary_ops/BiguintBinaryOps 280 216 64 216 0
boolean_binary_ops/BooleanBinaryOps 308 302 6 302 0
boolean_binary_ops/BooleanBinaryOps 308 298 10 298 0
bytes_ops/BiguintBinaryOps 143 139 4 139 0
calculator 346 317 29 315 2
callsub 32 32 0 32 0
Expand All @@ -38,15 +38,15 @@
global_state/AppState 305 301 4 301 0
hello_world/HelloWorld 23 22 1 22 0
hello_world_arc4/HelloWorld 110 89 21 89 0
inner_transactions 1847 1174 673 1174 0
inner_transactions 1847 1173 674 1173 0
inner_transactions/Greeter 325 302 23 302 0
inner_transactions/itxn_loop 203 184 19 184 0
intrinsics/ImmediateVariants 164 162 2 162 0
koopman 16 9 7 9 0
less_simple 171 148 23 148 0
local_state/LocalState 305 298 7 286 12
local_state/LocalStateWithOffsets 318 309 9 297 12
log 172 168 4 168 0
log 172 167 5 167 0
match 490 455 35 455 0
merkle/MerkleTree 217 210 7 210 0
nested_loops/Nested 243 201 42 201 0
Expand All @@ -65,9 +65,9 @@
stubs/Bytes 1769 258 1511 258 0
stubs/Uint64 371 8 363 8 0
template_variables/TemplateVariables 168 155 13 155 0
too_many_permutations 108 107 1 107 0
too_many_permutations 108 106 2 106 0
transaction/Transaction 893 849 44 849 0
tuple_support/TupleSupport 442 294 148 294 0
tuple_support/TupleSupport 442 292 150 292 0
typed_abi_call/Greeter 1285 1168 117 1168 0
typed_abi_call/Logger 569 478 91 478 0
unary/Unary 134 96 38 96 0
Expand Down
4 changes: 2 additions & 2 deletions examples/voting/out/VotingRoundApp.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -552,7 +552,7 @@ close:
int 0
// voting/voting.py:129
// for question_index, question_options_arc in uenumerate(self.option_counts):
int 0
dup
byte "option_counts"
app_global_get_ex
assert // check option_counts exists
Expand Down Expand Up @@ -1187,7 +1187,7 @@ vote:
int 0
// voting/voting.py:192
// for question_index in urange(questions_count):
int 0
dup

vote_for_header@1:
// voting/voting.py:192
Expand Down
2 changes: 1 addition & 1 deletion examples/voting/out/VotingRoundApp.arc32.json
View file

Large diffs are not rendered by default.

4 changes: 2 additions & 2 deletions examples/voting/out_O2/VotingRoundApp.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -301,7 +301,7 @@ close:
byte ",\"tallies\":["
concat
int 0
int 0
dup
byte "option_counts"
app_global_get_ex
assert // check option_counts exists
Expand Down Expand Up @@ -686,7 +686,7 @@ vote:
==
assert // Payment must be the exact min balance
int 0
int 0
dup

vote_for_header@1:
frame_dig 3
Expand Down
41 changes: 26 additions & 15 deletions src/puya/teal/optimize/constant_stack_shuffling.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import functools
import operator
import attrs

from puya.teal import models
from puya.teal.optimize._data import LOAD_OP_CODES
Expand All @@ -10,10 +9,13 @@ def perform_constant_stack_shuffling(block: models.TealBlock) -> bool:
loads = list[models.TealOp]()
modified = False
for op in block.ops:
if op.op_code in LOAD_OP_CODES: # noqa: SIM114
loads.append(op)
elif isinstance(op, models.FrameDig) and op.n < 0:
if _is_constant_load(op):
loads.append(op)
elif loads and op.op_code.startswith("dup"):
(n,) = op.immediates or (1,)
assert isinstance(n, int)
loads.extend([attrs.evolve(loads[-1], source_location=op.source_location)] * n)
modified = True
elif loads:
match op:
case models.Uncover(n=n) if n < len(loads):
Expand Down Expand Up @@ -48,9 +50,7 @@ def constant_dupn_insertion(block: models.TealBlock) -> bool:
modified = _collapse_loads(loads) or modified
result.extend(loads)
loads = []
if op.op_code in LOAD_OP_CODES or (
op.op_code == "frame_dig" and int(op.immediates[0]) < 0
):
if _is_constant_load(op):
loads.append(op)
else:
result.append(op)
Expand All @@ -63,10 +63,21 @@ def constant_dupn_insertion(block: models.TealBlock) -> bool:

def _collapse_loads(loads: list[models.TealOp]) -> bool:
n = len(loads) - 1
if n >= 2:
dupn_source_location = functools.reduce(
operator.add, (op.source_location for op in loads[1:])
)
loads[1:] = [models.DupN(n=n, source_location=dupn_source_location)]
return True
return False
if n < 1:
return False

if n == 1:
dup_op: models.TealOp = models.Dup(source_location=loads[1].source_location)
else:
dupn_source_location = None
for op in loads[1:]:
if op.source_location is not None:
# TODO: it'd be better to only merge these if they're adjacent
dupn_source_location = op.source_location + dupn_source_location
dup_op = models.DupN(n=n, source_location=dupn_source_location)
loads[1:] = [dup_op]
return True


def _is_constant_load(op: models.TealOp) -> bool:
return op.op_code in LOAD_OP_CODES or (isinstance(op, models.FrameDig) and op.n < 0)
6 changes: 5 additions & 1 deletion src/puya/teal/optimize/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,14 @@ def optimize_block(block: models.TealBlock, *, level: int) -> None:
modified = simplify_repeated_rotation_ops(block) or modified
modified = peephole(block) or modified

# don't do this one in a loop, only after everything else
# we don't do dup/dupn collapse in the above loop, but after it.
# it's easier to deal with expanded dup/dupn instructions above when looking at
# stack shuffling etc, but once it's done we save ops / program size by collapsing them
constant_dupn_insertion(block)

if level >= 2:
# this is a brute-force search which can be slow at times,
# so it's only done once and only at higher optimisation levels
block.ops = repeated_rotation_ops_search(block.ops)


Expand Down
2 changes: 1 addition & 1 deletion test_cases/arc4_types/out/Arc4ArraysContract.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ test_cases.arc4_types.array.Arc4ArraysContract.approval_program:
// arc4_types/array.py:26
// total = UInt64(0)
int 0
int 0
dup

main_for_header@1:
// arc4_types/array.py:27
Expand Down
2 changes: 1 addition & 1 deletion test_cases/arc4_types/out/Arc4BoolEvalContract.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ test_cases.arc4_types.bool_eval.Arc4BoolEvalContract.approval_program:
// arc4_types/bool_eval.py:19
// assert not arc4.Address(Global.zero_address)
global ZeroAddress
global ZeroAddress
dup
==
assert
// arc4_types/bool_eval.py:20
Expand Down
2 changes: 1 addition & 1 deletion test_cases/arc4_types/out/Arc4DynamicBytesContract.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ test_cases.arc4_types.dynamic_bytes.Arc4DynamicBytesContract.approval_program:
// arc4_types/dynamic_bytes.py:11
// total = UInt64(0)
int 0
int 0
dup

main_for_header@1:
// arc4_types/dynamic_bytes.py:16
Expand Down
16 changes: 10 additions & 6 deletions test_cases/arc4_types/out/Arc4RefTypesContract.approval.teal
View file
Original file line number Diff line number Diff line change
@@ -1,19 +1,23 @@
#pragma version 10

test_cases.arc4_types.reference_types.Arc4RefTypesContract.approval_program:
// arc4_types/reference_types.py:6-8
// arc4_types/reference_types.py:9-11
// # When creating an address from bytes, we check the length is 32 as we don't know the
// # source of the bytes
// checked_address = arc4.Address(op.Txn.sender.bytes)
txn Sender
// arc4_types/reference_types.py:6-11
// # When creating an address from an account no need to check the length as we assume the
// # Account is valid
// sender_address = arc4.Address(op.Txn.sender)
txn Sender
// # When creating an address from bytes, we check the length is 32 as we don't know the
// # source of the bytes
// checked_address = arc4.Address(op.Txn.sender.bytes)
dupn 3
// arc4_types/reference_types.py:9-11
// # When creating an address from bytes, we check the length is 32 as we don't know the
// # source of the bytes
// checked_address = arc4.Address(op.Txn.sender.bytes)
txn Sender
dup
cover 2
dup
len
int 32
==
Expand Down
2 changes: 1 addition & 1 deletion test_cases/arc4_types/out/Arc4StructsTypeContract.approval.teal
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ test_cases.arc4_types.structs.Arc4StructsTypeContract.approval_program:
byte 0x000000083cfbf217000000230384b842
// arc4_types/structs.py:28
// coord_2 = Vector(x=Decimal("35.382882839"), y=Decimal("150.382884930"))
byte 0x000000083cfbf217000000230384b842
dup
// arc4_types/structs.py:29
// coord_3 = add(coord_1.copy(), coord_2.copy())
callsub add
Expand Down
2 changes: 1 addition & 1 deletion test_cases/arc4_types/out/array.O1.log
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ PC Teal Stack
49 byte "" 0, 0x
50 dupn 2 0, 0x, 0x, 0x
52 int 0 0, 0x, 0x, 0x, 0
53 int 0 0, 0x, 0x, 0x, 0, 0
53 dup 0, 0x, 0x, 0x, 0, 0
54 dup 0, 0x, 0x, 0x, 0, 0, 0
55 int 2 0, 0x, 0x, 0x, 0, 0, 0, 2
56 < 0, 0x, 0x, 0x, 0, 0, 1
Expand Down
2 changes: 1 addition & 1 deletion test_cases/arc4_types/out/array.O2.log
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ PC Teal Stack
49 byte "" 0, 0x
50 dupn 2 0, 0x, 0x, 0x
52 int 0 0, 0x, 0x, 0x, 0
53 int 0 0, 0x, 0x, 0x, 0, 0
53 dup 0, 0x, 0x, 0x, 0, 0
54 dup 0, 0x, 0x, 0x, 0, 0, 0
55 int 2 0, 0x, 0x, 0x, 0, 0, 0, 2
56 < 0, 0x, 0x, 0x, 0, 0, 1
Expand Down
6 changes: 3 additions & 3 deletions test_cases/arc4_types/out/structs.O1.log
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
PC Teal Stack
1 <intcblock>
6 <bytecblock>
29 byte "" 0x
31 byte 0x000000083cfbf217000000230384b842 0x, 0x000000083CFBF217000000230384B842
32 byte 0x000000083cfbf217000000230384b842 0x, 0x000000083CFBF217000000230384B842, 0x000000083CFBF217000000230384B842
12 byte "" 0x
14 byte 0x000000083cfbf217000000230384b842 0x, 0x000000083CFBF217000000230384B842
32 dup 0x, 0x000000083CFBF217000000230384B842, 0x000000083CFBF217000000230384B842
33 callsub add 0x, 0x000000083CFBF217000000230384B842, 0x000000083CFBF217000000230384B842
100 proto 2 3 0x, 0x000000083CFBF217000000230384B842, 0x000000083CFBF217000000230384B842
103 frame_dig -2 0x, 0x000000083CFBF217000000230384B842, 0x000000083CFBF217000000230384B842, 0x000000083CFBF217000000230384B842
Expand Down
Loading

0 comments on commit b48db43

Please sign in to comment.