feat[venom]: add small heuristic for cleaning input stack (#4251)

charles-cooper · web-flow · commit d8f403235577 · 2024-10-10T17:43:41.000Z
this commit adds a small heuristic for cleaning the input stack from
`cfg_in`, which is to pop the shallowest items first.

it also cleans up the code a little bit and adds stronger preconditions
to the code.

a further optimization would be to try to align the output stack as best
as possible at the `jnz`, but that belongs in the DFT pass.
diff --git a/vyper/venom/venom_to_assembly.py b/vyper/venom/venom_to_assembly.py
@@ -295,7 +295,8 @@ def _generate_evm_for_basicblock_r(
         asm.append(f"_sym_{basicblock.label}")
         asm.append("JUMPDEST")
 
-        self.clean_stack_from_cfg_in(asm, basicblock, stack)
+        if len(basicblock.cfg_in) == 1:
+            self.clean_stack_from_cfg_in(asm, basicblock, stack)
 
         all_insts = sorted(basicblock.instructions, key=lambda x: x.opcode != "param")
 
@@ -321,26 +322,28 @@ def _generate_evm_for_basicblock_r(
     def clean_stack_from_cfg_in(
         self, asm: list, basicblock: IRBasicBlock, stack: StackModel
     ) -> None:
-        if len(basicblock.cfg_in) == 0:
-            return
-
-        to_pop = OrderedSet[IRVariable]()
-        for in_bb in basicblock.cfg_in:
-            # inputs is the input variables we need from in_bb
-            inputs = self.liveness_analysis.input_vars_from(in_bb, basicblock)
-
-            # layout is the output stack layout for in_bb (which works
-            # for all possible cfg_outs from the in_bb).
-            layout = in_bb.out_vars
-
-            # pop all the stack items which in_bb produced which we don't need.
-            to_pop |= layout.difference(inputs)
-
+        # the input block is a splitter block, like jnz or djmp
+        assert len(basicblock.cfg_in) == 1
+        in_bb = basicblock.cfg_in.first()
+        assert len(in_bb.cfg_out) > 1
+
+        # inputs is the input variables we need from in_bb
+        inputs = self.liveness_analysis.input_vars_from(in_bb, basicblock)
+
+        # layout is the output stack layout for in_bb (which works
+        # for all possible cfg_outs from the in_bb, in_bb is responsible
+        # for making sure its output stack layout works no matter which
+        # bb it jumps into).
+        layout = in_bb.out_vars
+        to_pop = list(layout.difference(inputs))
+
+        # small heuristic: pop from shallowest first.
+        to_pop.sort(key=lambda var: -stack.get_depth(var))
+
+        # NOTE: we could get more fancy and try to optimize the swap
+        # operations here, there is probably some more room for optimization.
         for var in to_pop:
             depth = stack.get_depth(var)
-            # don't pop phantom phi inputs
-            if depth is StackModel.NOT_IN_STACK:
-                continue
 
             if depth != 0:
                 self.swap(asm, stack, depth)
