fix[venom]: check operand position in mem2var promotion

tests/unit/compiler/venom/test_mem2var.py

@@ -0,0 +1,65 @@
+from tests.venom_utils import parse_from_basic_block
+from vyper.venom.analysis import IRAnalysesCache
+from vyper.venom.passes import MakeSSA
+from vyper.venom.passes.mem2var import Mem2Var
+
+
+def _find_insts(fn, opcode):
+    return [
+        inst for bb in fn.get_basic_blocks() for inst in bb.instructions if inst.opcode == opcode
+    ]
+
+
+def _run_mem2var(pre: str):
+    ctx = parse_from_basic_block(pre)
+    fn = next(iter(ctx.functions.values()))
+    ac = IRAnalysesCache(fn)
+    MakeSSA(ac, fn).run_pass()
+    Mem2Var(ac, fn).run_pass()
+    return fn
+
+
+def test_mem2var_promotes_simple_alloca():
+    """
+    A 32-byte alloca whose only uses are mload/mstore address operands
+    is promoted to a stack variable.
+    """
+    pre = """
+    main:
+        %ptr = alloca 32
+        mstore %ptr, 42
+        %v = mload %ptr
+        sink %v
+    """
+    fn = _run_mem2var(pre)
+
+    assert len(_find_insts(fn, "mload")) == 0
+    assert len(_find_insts(fn, "mstore")) == 0
+
+
+def test_mem2var_skips_alloca_used_as_stored_value():
+    """
+    Mem2Var must NOT promote an alloca whose pointer is used as the
+    VALUE operand of an mstore (`mstore %b, %a` stores the pointer %a
+    into the slot %b) -- the pointer escapes through memory.
+    See issue #5070.
+    """
+    pre = """
+    main:
+        %a = alloca 32
+        %b = source
+        mstore %b, %a
+        %x = mload %a
+        return %b, 32
+    """
+    fn = _run_mem2var(pre)
+
+    # %a must not be promoted: the store of the pointer must survive
+    allocas = _find_insts(fn, "alloca")
+    assert len(allocas) == 1, f"alloca should be preserved, got {allocas}"
+    mstores = _find_insts(fn, "mstore")
+    assert len(mstores) == 1, f"mstore of the pointer must survive, got {mstores}"
+    val, _ptr = mstores[0].operands
+    assert val == allocas[0].output, f"stored value should be the pointer, got {val}"
+    # the load through %a must survive as well
+    assert len(_find_insts(fn, "mload")) == 1

vyper/venom/passes/mem2var.py

@@ -35,11 +35,26 @@ def _mk_varname(self, varname: str):
         self.var_name_count += 1
         return varname
 
+    @staticmethod
+    def _is_pointer_use(inst: IRInstruction, ptr: IRVariable) -> bool:
+        """
+        Check that `inst` uses `ptr` only as a memory address, and not
+        e.g. as the value operand of an mstore.
+        """
+        if inst.opcode == "mload":
+            # the only operand is the address
+            return True
+        if inst.opcode in ("mstore", "return"):
+            # mstore [val, ptr]; return [size, ptr] -- the pointer must
+            # not appear in the value/size position.
+            return inst.operands[0] != ptr
+        return False
+
     def _process_alloca_var(self, dfg: DFGAnalysis, alloca_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.
+        Process alloca allocated variable. If it is only used as the
+        memory address of mstore/mload/return instructions, it is promoted
+        to a stack variable. Otherwise, it is left as is.
         """
 
         assert len(alloca_inst.operands) == 1, (alloca_inst, alloca_inst.parent)
@@ -51,7 +66,7 @@ def _process_alloca_var(self, dfg: DFGAnalysis, alloca_inst: IRInstruction, var:
 
         uses = dfg.get_uses(alloca_inst.output)
 
-        if not all2(inst.opcode in ["mstore", "mload", "return"] for inst in uses):
+        if not all2(self._is_pointer_use(inst, alloca_inst.output) for inst in uses):
             return
 
         size = size_lit.value