fix[venom]: promote additional memory locations to variables (vyperla…

…ng#4039) promote additional memory locations (specifically, internal function params allocated with `palloca`) to stack variables. --------- Co-authored-by: Charles Cooper <[email protected]>
charles-cooper · Oct 10, 2024 · fdc05d6 · fdc05d6
1 parent d8f4032
commit fdc05d6
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 13 deletions.
diff --git a/vyper/venom/ir_node_to_venom.py b/vyper/venom/ir_node_to_venom.py
@@ -548,7 +548,7 @@ def emit_body_blocks():
             _global_symbols[ir.value] = ptr
         elif ir.value.startswith("$palloca") and ir.value not in _global_symbols:
             alloca = ir.passthrough_metadata["alloca"]
-            ptr = fn.get_basic_block().append_instruction("store", alloca.offset)
+            ptr = fn.get_basic_block().append_instruction("palloca", alloca.offset, alloca.size)
             _global_symbols[ir.value] = ptr
 
         return _global_symbols.get(ir.value) or symbols.get(ir.value)

diff --git a/vyper/venom/passes/mem2var.py b/vyper/venom/passes/mem2var.py
@@ -1,8 +1,7 @@
-from vyper.utils import OrderedSet
 from vyper.venom.analysis.cfg import CFGAnalysis
 from vyper.venom.analysis.dfg import DFGAnalysis
 from vyper.venom.analysis.liveness import LivenessAnalysis
-from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IRVariable
+from vyper.venom.basicblock import IRInstruction, IRVariable
 from vyper.venom.function import IRFunction
 from vyper.venom.passes.base_pass import IRPass
 
@@ -14,21 +13,27 @@ class Mem2Var(IRPass):
     """
 
     function: IRFunction
-    defs: dict[IRVariable, OrderedSet[IRBasicBlock]]
 
     def run_pass(self):
         self.analyses_cache.request_analysis(CFGAnalysis)
         dfg = self.analyses_cache.request_analysis(DFGAnalysis)
 
         self.var_name_count = 0
         for var, inst in dfg.outputs.items():
-            if inst.opcode != "alloca":
-                continue
-            self._process_alloca_var(dfg, var)
+            if inst.opcode == "alloca":
+                self._process_alloca_var(dfg, var)
+            elif inst.opcode == "palloca":
+                self._process_palloca_var(dfg, inst, var)
 
         self.analyses_cache.invalidate_analysis(DFGAnalysis)
         self.analyses_cache.invalidate_analysis(LivenessAnalysis)
 
+    def _mk_varname(self, varname: str):
+        varname = varname.removeprefix("%")
+        varname = f"var{varname}_{self.var_name_count}"
+        self.var_name_count += 1
+        return varname
+
     def _process_alloca_var(self, dfg: DFGAnalysis, var: IRVariable):
         """
         Process alloca allocated variable. If it is only used by mstore/mload/return
@@ -40,8 +45,7 @@ def _process_alloca_var(self, dfg: DFGAnalysis, var: IRVariable):
         elif all([inst.opcode == "mstore" for inst in uses]):
             return
         elif all([inst.opcode in ["mstore", "mload", "return"] for inst in uses]):
-            var_name = f"addr{var.name}_{self.var_name_count}"
-            self.var_name_count += 1
+            var_name = self._mk_varname(var.name)
             for inst in uses:
                 if inst.opcode == "mstore":
                     inst.opcode = "store"
@@ -52,7 +56,32 @@ def _process_alloca_var(self, dfg: DFGAnalysis, var: IRVariable):
                     inst.operands = [IRVariable(var_name)]
                 elif inst.opcode == "return":
                     bb = inst.parent
-                    idx = bb.instructions.index(inst)
+                    idx = len(bb.instructions) - 1
+                    assert inst == bb.instructions[idx]  # sanity
                     bb.insert_instruction(
                         IRInstruction("mstore", [IRVariable(var_name), inst.operands[1]]), idx
                     )
+
+    def _process_palloca_var(self, dfg: DFGAnalysis, palloca_inst: IRInstruction, var: IRVariable):
+        """
+        Process alloca allocated variable. If it is only used by mstore/mload
+        instructions, it is promoted to a stack variable. Otherwise, it is left as is.
+        """
+        uses = dfg.get_uses(var)
+        if not all(inst.opcode in ["mstore", "mload"] for inst in uses):
+            return
+
+        var_name = self._mk_varname(var.name)
+
+        palloca_inst.opcode = "mload"
+        palloca_inst.operands = [palloca_inst.operands[0]]
+        palloca_inst.output = IRVariable(var_name)
+
+        for inst in uses:
+            if inst.opcode == "mstore":
+                inst.opcode = "store"
+                inst.output = IRVariable(var_name)
+                inst.operands = [inst.operands[0]]
+            elif inst.opcode == "mload":
+                inst.opcode = "store"
+                inst.operands = [IRVariable(var_name)]
diff --git a/vyper/venom/passes/sccp/sccp.py b/vyper/venom/passes/sccp/sccp.py
@@ -179,7 +179,7 @@ def _visit_phi(self, inst: IRInstruction):
 
     def _visit_expr(self, inst: IRInstruction):
         opcode = inst.opcode
-        if opcode in ["store", "alloca"]:
+        if opcode in ["store", "alloca", "palloca"]:
             assert inst.output is not None, "Got store/alloca without output"
             out = self._eval_from_lattice(inst.operands[0])
             self._set_lattice(inst.output, out)

diff --git a/vyper/venom/venom_to_assembly.py b/vyper/venom/venom_to_assembly.py
@@ -363,7 +363,7 @@ def _generate_evm_for_instruction(
 
         if opcode in ["jmp", "djmp", "jnz", "invoke"]:
             operands = list(inst.get_non_label_operands())
-        elif opcode == "alloca":
+        elif opcode in ("alloca", "palloca"):
             offset, _size = inst.operands
             operands = [offset]
 
@@ -463,7 +463,7 @@ def _generate_evm_for_instruction(
         # Step 5: Emit the EVM instruction(s)
         if opcode in _ONE_TO_ONE_INSTRUCTIONS:
             assembly.append(opcode.upper())
-        elif opcode == "alloca":
+        elif opcode in ("alloca", "palloca"):
             pass
         elif opcode == "param":
             pass