feat: allow variable rebinding of mutable parameter values

algorandfoundation · Nov 1, 2024 · 4de5a3e · 4de5a3e
1 parent c0051e1
commit 4de5a3e
Show file tree

Hide file tree

Showing 71 changed files with 11,750 additions and 2,177 deletions.
diff --git a/examples/merkle/out/MerkleTree.ssa.ir b/examples/merkle/out/MerkleTree.ssa.ir
diff --git a/examples/merkle/puya.log b/examples/merkle/puya.log
diff --git a/examples/sizes.txt b/examples/sizes.txt
@@ -12,14 +12,15 @@
  arc4_types/Arc4BoolType                        381      69       -    |        307      46       - 
  arc4_types/Arc4DynamicBytes                    377     185       -    |        213     100       - 
  arc4_types/Arc4DynamicStringArray              283     124       -    |        172      53       - 
- arc4_types/Arc4MutableParams                   471     286       -    |        292     141       - 
+ arc4_types/Arc4MutableParams                  1110     644       -    |        699     343       - 
  arc4_types/Arc4Mutation                       2958    1426       -    |       1977     593       - 
  arc4_types/Arc4NumericTypes                    749     186       -    |        243      26       - 
  arc4_types/Arc4RefTypes                         85      46       -    |         32      27       - 
  arc4_types/Arc4StringTypes                     455      35       -    |        245      13       - 
  arc4_types/Arc4StructsFromAnotherModule         67      12       -    |         49       6       - 
- arc4_types/Arc4StructsType                     391     239       -    |        259     121       - 
+ arc4_types/Arc4StructsType                     389     239       -    |        259     121       - 
  arc4_types/Arc4TuplesType                      795     136       -    |        537      58       - 
+ arc4_types/MutableParams2                      334     202       -    |        193      97       - 
  arc_28/EventEmitter                            186     133       -    |        100      64       - 
  asset/Reference                                268     261       -    |        144     141       - 
  auction/Auction                                592     522       -    |        328     281       - 
@@ -108,7 +109,7 @@
  state_totals                                    65      32       -    |         32      16       - 
  stress_tests/BruteForceRotationSearch          228     163       -    |        152     106       - 
  string_ops                                     156     154       -    |         58      55       - 
- struct_in_box/Example                          243     206       -    |        127      99       - 
+ struct_in_box/Example                          242     206       -    |        127      99       - 
  stubs/BigUInt                                  192     121       -    |        126      73       - 
  stubs/Bytes                                    944     279       -    |        606     153       - 
  stubs/String                                   701     167       -    |        416      54       - 
@@ -128,6 +129,6 @@
  unassigned_expression/Unassigned               158     119       -    |         81      57       - 
  undefined_phi_args/Baddie                      319     282       -    |        174     157       - 
  unssa/UnSSA                                    432     368       -    |        241     204       - 
- voting/VotingRoundApp                         1593    1483       -    |        734     649       - 
+ voting/VotingRoundApp                         1590    1483       -    |        733     649       - 
  with_reentrancy/WithReentrancy                 255     242       -    |        132     122       - 
- Total                                        69283   53532   53473    |      32895   21757   21713 
+ Total                                        70250   54092   54033    |      33494   22056   22012 
diff --git a/examples/struct_in_box/out/ExampleContract.ssa.ir b/examples/struct_in_box/out/ExampleContract.ssa.ir
diff --git a/examples/struct_in_box/out_unoptimized/ExampleContract.approval.teal b/examples/struct_in_box/out_unoptimized/ExampleContract.approval.teal
@@ -165,9 +165,10 @@ write_to_box:
     // @subroutine
     // def write_to_box(self, user: UserStruct) -> None:
     proto 1 1
+    frame_dig -1
     // struct_in_box/contract.py:20
     // box_key = user.id.bytes
-    frame_dig -1
+    dup
     intc_2 // 2
     intc_3 // 8
     extract3 // on error: Index access is out of bounds
@@ -181,7 +182,6 @@ write_to_box:
     // op.Box.put(box_key, user.bytes)
     frame_dig -1
     box_put
-    frame_dig -1
     retsub
 
 

diff --git a/examples/struct_in_box/out_unoptimized/ExampleContract.destructured.ir b/examples/struct_in_box/out_unoptimized/ExampleContract.destructured.ir
diff --git a/examples/struct_in_box/puya.log b/examples/struct_in_box/puya.log
diff --git a/examples/voting/out/VotingRoundApp.ssa.ir b/examples/voting/out/VotingRoundApp.ssa.ir
diff --git a/examples/voting/out_unoptimized/VotingRoundApp.approval.teal b/examples/voting/out_unoptimized/VotingRoundApp.approval.teal
@@ -281,9 +281,10 @@ store_option_counts:
     // @subroutine
     // def store_option_counts(self, option_counts: VoteIndexArray) -> None:
     proto 1 1
+    frame_dig -1
     // voting/voting.py:219
     // assert option_counts.length, "option_counts should be non-empty"
-    frame_dig -1
+    dup
     intc_0 // 0
     extract_uint16
     intc_0 // 0
@@ -310,33 +311,33 @@ store_option_counts:
 store_option_counts_for_header@1:
     // voting/voting.py:223
     // for item in option_counts:
+    frame_dig 3
     frame_dig 2
-    frame_dig 1
     <
     bz store_option_counts_after_for@4
     frame_dig -1
     extract 2 0
-    frame_dig 2
+    frame_dig 3
     intc_1 // 1
     *
     intc_1 // 1
     extract3 // on error: Index access is out of bounds
     // voting/voting.py:224
     // total_options += item.native
     btoi
-    frame_dig 0
+    frame_dig 1
     +
-    frame_bury 0
-    frame_dig 2
+    frame_bury 1
+    frame_dig 3
     intc_1 // 1
     +
-    frame_bury 2
+    frame_bury 3
     b store_option_counts_for_header@1
 
 store_option_counts_after_for@4:
     // voting/voting.py:225
     // assert total_options <= 128, "Can't have more than 128 vote options"
-    frame_dig 0
+    frame_dig 1
     dup
     pushint 128 // 128
     <=
@@ -352,8 +353,6 @@ store_option_counts_after_for@4:
     bytec 14 // "total_options"
     swap
     app_global_put
-    frame_dig -1
-    frame_bury 0
     retsub
 
 

diff --git a/examples/voting/out_unoptimized/VotingRoundApp.destructured.ir b/examples/voting/out_unoptimized/VotingRoundApp.destructured.ir
diff --git a/examples/voting/puya.log b/examples/voting/puya.log
diff --git a/src/puya/awst/validation/arc4_copy.py b/src/puya/awst/validation/arc4_copy.py
@@ -53,18 +53,18 @@ def visit_assignment_expression(self, expr: awst_nodes.AssignmentExpression) ->
         expr.value.accept(self)
 
     def visit_subroutine_call_expression(self, expr: awst_nodes.SubroutineCallExpression) -> None:
-        super().visit_subroutine_call_expression(expr)
-        for arg in expr.args:
-            match arg.value:
-                case awst_nodes.VarExpression():
-                    # Var expressions don't need copy as we implicitly return the latest value and
-                    # update the var
-                    continue
-                case awst_nodes.AppStateExpression() | awst_nodes.AppAccountStateExpression():
-                    message = "being passed to a subroutine from state"
-                case _:
-                    message = "being passed to a subroutine"
-            _check_for_arc4_copy(arg.value, message)
+        for arg_ in expr.args:
+            for arg in _expand_tuple_items(arg_.value):
+                match arg:
+                    case awst_nodes.VarExpression():
+                        # Var expressions don't need copy as we implicitly return the latest value
+                        # and update the var
+                        continue
+                    case awst_nodes.AppStateExpression() | awst_nodes.AppAccountStateExpression():
+                        message = "being passed to a subroutine from state"
+                    case _:
+                        message = "being passed to a subroutine"
+                _check_for_arc4_copy(arg, message)
 
     def visit_new_array(self, expr: awst_nodes.NewArray) -> None:
         super().visit_new_array(expr)
@@ -131,7 +131,8 @@ def _check_for_arc4_copy(expr: awst_nodes.Expression, context_desc: str) -> None
 def _expand_tuple_items(expr: awst_nodes.Expression) -> Iterator[awst_nodes.Expression]:
     match expr:
         case awst_nodes.TupleExpression(items=items):
-            yield from items
+            for item in items:
+                yield from _expand_tuple_items(item)
         case _:
             yield expr
 

diff --git a/src/puya/ir/builder/_utils.py b/src/puya/ir/builder/_utils.py
@@ -3,10 +3,13 @@
 
 import attrs
 
-from puya.awst import nodes as awst_nodes
+from puya.awst import (
+    nodes as awst_nodes,
+    wtypes,
+)
 from puya.errors import InternalError
 from puya.ir.avm_ops import AVMOp
-from puya.ir.context import IRFunctionBuildContext
+from puya.ir.context import TMP_VAR_INDICATOR, IRFunctionBuildContext
 from puya.ir.models import (
     Assignment,
     BytesConstant,
@@ -76,6 +79,55 @@ def assign_targets(
     context.block_builder.add(
         Assignment(targets=targets, source=source, source_location=assignment_location)
     )
+    # also update any implicitly returned variables
+    implicit_params = {p.name for p in context.subroutine.parameters if p.implicit_return}
+    for target in targets:
+        if target.name in implicit_params:
+            _update_implicit_out_var(context, target.name, target.ir_type)
+
+
+def _update_implicit_out_var(context: IRFunctionBuildContext, var: str, ir_type: IRType) -> None:
+    # emit conditional assignment equivalent to
+    # if var%is_original:
+    #   var%out = var
+    loc = SourceLocation(file=None, line=1)
+    wtype = wtypes.bytes_wtype if ir_type == IRType.bytes else wtypes.uint64_wtype
+    node = awst_nodes.IfElse(
+        condition=awst_nodes.VarExpression(
+            name=get_implicit_return_is_original(var),
+            wtype=wtypes.bool_wtype,
+            source_location=loc,
+        ),
+        if_branch=awst_nodes.Block(
+            body=[
+                awst_nodes.AssignmentStatement(
+                    target=awst_nodes.VarExpression(
+                        name=get_implicit_return_out(var),
+                        wtype=wtype,
+                        source_location=loc,
+                    ),
+                    value=awst_nodes.VarExpression(
+                        name=var,
+                        wtype=wtype,
+                        source_location=loc,
+                    ),
+                    source_location=loc,
+                )
+            ],
+            source_location=loc,
+        ),
+        else_branch=None,
+        source_location=loc,
+    )
+    node.accept(context.visitor)
+
+
+def get_implicit_return_is_original(var_name: str) -> str:
+    return f"{var_name}{TMP_VAR_INDICATOR}is_original"
+
+
+def get_implicit_return_out(var_name: str) -> str:
+    return f"{var_name}{TMP_VAR_INDICATOR}out"
 
 
 def mktemp(