feat[venom]: skip pops of halting blocks

tests/unit/compiler/venom/test_duplicate_operands.py

@@ -15,7 +15,7 @@ def test_duplicate_operands():
     %3 = mul %1, %2
     stop
 
-    Should compile to: [PUSH1, 10, DUP1, DUP2, ADD, MUL, POP, STOP]
+    Should compile to: [PUSH1, 10, DUP1, DUP2, ADD, MUL, STOP]
     """
     ctx = IRContext()
     fn = ctx.create_function("test")
@@ -30,4 +30,4 @@ def test_duplicate_operands():
 
     optimize = OptimizationLevel.GAS
     asm = generate_assembly_experimental(ctx, optimize=optimize)
-    assert asm == ["PUSH1", 10, "DUP1", "DUP2", "ADD", "MUL", "POP", "STOP"]
+    assert asm == ["PUSH1", 10, "DUP1", "DUP2", "ADD", "MUL", "STOP"]

tests/unit/compiler/venom/test_venom_to_assembly.py

@@ -39,34 +39,6 @@ def test_optimistic_swap_params():
     assert asm == ["SWAP2", "PUSH1", 117, "POP", "MSTORE", "MSTORE", "JUMP"]
 
 
-def test_invoke_middle_output_unused():
-    code = """
-    function main {
-    main:
-        %a, %b, %c = invoke @callee
-        return %a, %c
-    }
-
-    function callee {
-    callee:
-        %retpc = param
-        %x = 1
-        %y = 2
-        %z = 3
-        ret %x, %y, %z, %retpc
-    }
-    """
-    ctx = parse_venom(code)
-    asm = VenomCompiler(ctx).generate_evm_assembly()
-
-    assert "POP" in asm, f"expected POP to remove dead output, got {asm}"
-    pop_idx = asm.index("POP")
-    assert pop_idx > 0 and asm[pop_idx - 1] == "SWAP1", asm
-    assert "RETURN" in asm, asm
-    return_idx = asm.index("RETURN")
-    assert return_idx > pop_idx and asm[return_idx - 1] == "SWAP1", asm
-
-
 def test_popmany_bulk_removal_of_suffix():
     compiler = VenomCompiler(IRContext())
     stack = StackModel()

vyper/venom/venom_to_assembly.py

@@ -111,6 +111,9 @@
 
 _REVERT_POSTAMBLE = [Label("revert"), *PUSH(0), "DUP1", "REVERT"]
 
+# Terminators that halt execution - no need to clean up stack before these
+_HALTING_TERMINATORS = frozenset(["return", "revert", "stop"])
+
 
 def apply_line_numbers(inst: IRInstruction, asm) -> list[str]:
     ret = []
@@ -355,13 +358,19 @@ def _generate_evm_for_basicblock_r(
 
         all_insts = [inst for inst in basicblock.instructions if inst.opcode != "param"]
 
+        # Check if this block ends with a halting terminator (return, revert, stop)
+        # If so, we don't need to pop dead variables since execution halts anyway
+        is_halting_block = all_insts[-1].opcode in _HALTING_TERMINATORS
+
         for i, inst in enumerate(all_insts):
             if i + 1 < len(all_insts):
                 next_liveness = self.liveness.live_vars_at(all_insts[i + 1])
             else:
                 next_liveness = self.liveness.out_vars(basicblock)
 
-            asm.extend(self._generate_evm_for_instruction(inst, stack, next_liveness))
+            asm.extend(
+                self._generate_evm_for_instruction(inst, stack, next_liveness, is_halting_block)
+            )
 
         if DEBUG_SHOW_COST:
             print(" ".join(map(str, asm)), file=sys.stderr)
@@ -395,7 +404,11 @@ def clean_stack_from_cfg_in(
         self.popmany(asm, to_pop, stack)
 
     def _generate_evm_for_instruction(
-        self, inst: IRInstruction, stack: StackModel, next_liveness: OrderedSet
+        self,
+        inst: IRInstruction,
+        stack: StackModel,
+        next_liveness: OrderedSet,
+        skip_pops: bool = False,
     ) -> list[str]:
         assembly: list[AssemblyInstruction] = []
         opcode = inst.opcode
@@ -605,8 +618,10 @@ def _generate_evm_for_instruction(
         if len(outputs) == 0:
             return apply_line_numbers(inst, assembly)
 
-        dead_outputs = [out for out in outputs if out not in next_liveness]
-        self.popmany(assembly, dead_outputs, stack)
+        # Skip popping dead outputs if we're in a halting block (return/revert/stop)
+        if not skip_pops:
+            dead_outputs = [out for out in outputs if out not in next_liveness]
+            self.popmany(assembly, dead_outputs, stack)
 
         live_outputs = [out for out in outputs if out in next_liveness]
         if len(live_outputs) == 0: