fix[venom]: promote additional memory locations to variables (#4039)

promote additional memory locations (specifically, internal function
params allocated with `palloca`) to stack variables.

---------

Co-authored-by: Charles Cooper <[email protected]>

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)

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)]

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)

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