fix[venom]: promote additional memory locations to variables

vyper/venom/ir_node_to_venom.py

@@ -543,7 +543,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)
             _global_symbols[ir.value] = ptr
 
         return _global_symbols.get(ir.value) or symbols.get(ir.value)

vyper/venom/passes/mem2var.py

@@ -13,6 +13,10 @@ class Mem2Var(IRPass):
     It does yet do any memory aliasing analysis, so it is conservative.
     """
 
+    # Number of memory operations over which we promote
+    # a palloca to a stack variable
+    OPS_THREASHOLD = 6
+
     function: IRFunction
     defs: dict[IRVariable, OrderedSet[IRBasicBlock]]
 
@@ -23,9 +27,10 @@ def run_pass(self):
 
         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)
@@ -57,3 +62,36 @@ def _process_alloca_var(self, dfg: DFGAnalysis, var: IRVariable):
                     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/return
+        instructions, it is promoted to a stack variable. Otherwise, it is left as is.
+        """
+        uses = dfg.get_uses(var)
+        work_instructions = [inst.opcode in ["mstore", "mload"] for inst in uses]
+        count = sum(work_instructions)
+        if count > self.OPS_THREASHOLD and count == len(work_instructions):
+            bb = palloca_inst.parent
+            var_name = f"addr{var.name}_{self.var_name_count}"
+            self.var_name_count += 1
+
+            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)]
+                elif inst.opcode == "return":
+                    idx = bb.instructions.index(inst)
+                    bb.insert_instruction(
+                        IRInstruction("mstore", [IRVariable(var_name), inst.operands[1]]), idx
+                    )
+        else:
+            palloca_inst.opcode = "store"