update normalization check

charles-cooper · charles-cooper · commit 28f0ba2dd439 · 2024-09-19T09:35:38.000-04:00
diff --git a/vyper/venom/function.py b/vyper/venom/function.py
@@ -148,13 +148,17 @@ def _compute_reachability_from(self, bb: IRBasicBlock) -> None:
     def normalized(self) -> bool:
         """
         Check if function is normalized. A function is normalized if in the
-        CFG, no basic block simultaneously has multiple inputs and outputs.
+        CFG, no basic block simultaneously has multiple inputs and outputs,
+        and no basic block has a single input/output.
         That is, a basic block can be jumped to *from* multiple blocks, or it
         can jump *to* multiple blocks, but it cannot simultaneously do both.
         Having a normalized CFG makes calculation of stack layout easier when
         emitting assembly.
         """
         for bb in self.get_basic_blocks():
+            if len(bb.cfg_in) == 1 and len(bb.cfg_out) == 1:
+                return False
+
             # Ignore if there are no multiple predecessors
             if len(bb.cfg_in) <= 1:
                 continue
diff --git a/vyper/venom/venom_to_assembly.py b/vyper/venom/venom_to_assembly.py
@@ -394,10 +394,13 @@ def _generate_evm_for_instruction(
 
         # Step 3: Reorder stack before join points
         if opcode == "jmp":
-            # after cfg normalization, join points are represented by jmp instructions
             # prepare stack for jump into another basic block
-            assert inst.parent and isinstance(inst.parent.cfg_out, OrderedSet)
-            b = next(iter(inst.parent.cfg_out))
+            # we only need to reorder stack before join points, which after
+            # cfg simplification and normalization, join points are represented
+            # by jmp instructions.
+            assert isinstance(inst.parent.cfg_out, OrderedSet)
+            assert len(inst.parent.cfg_out) == 1
+            b = inst.parent.cfg_out.first()
             target_stack = self.liveness_analysis.input_vars_from(inst.parent, b)
             # TODO optimize stack reordering at entry and exit from basic blocks
             # NOTE: stack in general can contain multiple copies of the same variable,
