fix[venom]: invalid jump error (#4214)

fix an issue where `ret` would generate a jump to an invalid location
because the stack is not clean. the solution is to always pop
instruction outputs which are not used.

note this leads to a slight performance regression (roughly 0.07%
bytecode size), since `POP`s are always generated, even in cases when the
stack does not need to be cleaned (e.g. before `STOP`).

---------

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

tests/unit/compiler/venom/test_duplicate_operands.py

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

tests/unit/compiler/venom/test_stack_cleanup.py

@@ -0,0 +1,16 @@
+from vyper.compiler.settings import OptimizationLevel
+from vyper.venom import generate_assembly_experimental
+from vyper.venom.context import IRContext
+
+
+def test_cleanup_stack():
+    ctx = IRContext()
+    fn = ctx.create_function("test")
+    bb = fn.get_basic_block()
+    ret_val = bb.append_instruction("param")
+    op = bb.append_instruction("store", 10)
+    bb.append_instruction("add", op, op)
+    bb.append_instruction("ret", ret_val)
+
+    asm = generate_assembly_experimental(ctx, optimize=OptimizationLevel.GAS)
+    assert asm == ["PUSH1", 10, "DUP1", "ADD", "POP", "JUMP"]

vyper/venom/venom_to_assembly.py

@@ -285,34 +285,16 @@ def _generate_evm_for_basicblock_r(
 
         self.clean_stack_from_cfg_in(asm, basicblock, stack)
 
-        param_insts = [inst for inst in basicblock.instructions if inst.opcode == "param"]
-        main_insts = [inst for inst in basicblock.instructions if inst.opcode != "param"]
+        all_insts = sorted(basicblock.instructions, key=lambda x: x.opcode != "param")
 
-        for inst in param_insts:
-            asm.extend(self._generate_evm_for_instruction(inst, stack))
-
-        self._clean_unused_params(asm, basicblock, stack)
-
-        for i, inst in enumerate(main_insts):
-            next_liveness = main_insts[i + 1].liveness if i + 1 < len(main_insts) else OrderedSet()
+        for i, inst in enumerate(all_insts):
+            next_liveness = all_insts[i + 1].liveness if i + 1 < len(all_insts) else OrderedSet()
 
             asm.extend(self._generate_evm_for_instruction(inst, stack, next_liveness))
 
         for bb in basicblock.reachable:
             self._generate_evm_for_basicblock_r(asm, bb, stack.copy())
 
-    def _clean_unused_params(self, asm: list, bb: IRBasicBlock, stack: StackModel) -> None:
-        for i, inst in enumerate(bb.instructions):
-            if inst.opcode != "param":
-                break
-            if inst.is_volatile and i + 1 < len(bb.instructions):
-                liveness = bb.instructions[i + 1].liveness
-                if inst.output is not None and inst.output not in liveness:
-                    depth = stack.get_depth(inst.output)
-                    if depth != 0:
-                        self.swap(asm, stack, depth)
-                    self.pop(asm, stack)
-
     # pop values from stack at entry to bb
     # note this produces the same result(!) no matter which basic block
     # we enter from in the CFG.
@@ -543,11 +525,12 @@ def _generate_evm_for_instruction(
 
         # Step 6: Emit instructions output operands (if any)
         if inst.output is not None:
-            if "call" in inst.opcode and inst.output not in next_liveness:
+            if inst.output not in next_liveness:
                 self.pop(assembly, stack)
-            elif inst.output in next_liveness:
+            else:
                 # peek at next_liveness to find the next scheduled item,
                 # and optimistically swap with it
+                # TODO: implement OrderedSet.last()
                 next_scheduled = list(next_liveness)[-1]
                 self.swap_op(assembly, stack, next_scheduled)