perf(codegen): plan switch lowering

crates/codegen/src/backend/evm/assembler/mod.rs

@@ -97,6 +97,8 @@ pub(crate) struct Assembler {
     label_relocations: Vec<(ir::BlockId, usize, Label)>,
     /// Unresolved deferred constants emitted as push operands.
     deferred_relocations: Vec<(ir::BlockId, usize, DeferredConst)>,
+    /// Indexed jumps whose possible targets are assembler labels.
+    indexed_jump_relocations: Vec<(ir::BlockId, Vec<Label>)>,
     /// Interned push immediates too large for inline storage.
     push_values: LocalInterner<U256, PushValueId>,
     /// Next label ID.
@@ -126,6 +128,7 @@ impl Assembler {
             cold_labels: GrowableBitSet::new_empty(),
             label_relocations: Vec::new(),
             deferred_relocations: Vec::new(),
+            indexed_jump_relocations: Vec::new(),
             push_values: LocalInterner::new(),
             next_label: IdCounter::new(),
             next_deferred: IdCounter::new(),
@@ -142,6 +145,7 @@ impl Assembler {
         self.cold_labels.clear();
         self.label_relocations.clear();
         self.deferred_relocations.clear();
+        self.indexed_jump_relocations.clear();
         self.push_values.clear();
         self.next_label.clear();
         self.next_deferred.clear();
@@ -175,6 +179,13 @@ impl Assembler {
         self.label_relocations.push((block, instruction, label));
     }
 
+    /// Terminates the current block with an indexed jump to one of `targets`.
+    pub(crate) fn emit_indexed_jump(&mut self, targets: Vec<Label>) {
+        assert!(!targets.is_empty(), "indexed jump must have at least one target");
+        let block = self.current_block();
+        self.indexed_jump_relocations.push((block, targets));
+    }
+
     /// Emits a push instruction for a deferred constant.
     pub(crate) fn emit_push_deferred(&mut self, id: DeferredConst) {
         let (block, instruction) =
@@ -265,10 +276,23 @@ impl Assembler {
         for (block, instruction, id) in self.deferred_relocations.drain(..) {
             module.blocks[block].instructions[instruction] = ir::Instruction::push_deferred(id);
         }
+        Self::finalize_evm_ir(&mut module);
+        for (block, targets) in self.indexed_jump_relocations.drain(..) {
+            let targets = targets
+                .into_iter()
+                .map(|label| {
+                    self.label_blocks
+                        .get(&label)
+                        .copied()
+                        .unwrap_or_else(|| panic!("label {label:?} was never defined"))
+                })
+                .collect::<Vec<_>>()
+                .into_boxed_slice();
+            module.blocks[block].terminator =
+                Some(ir::Terminator::new(ir::TerminatorKind::IndexedJump(targets)));
+        }
         self.label_blocks.clear();
         self.cold_labels.clear();
-
-        Self::finalize_evm_ir(&mut module);
         Some((module, std::mem::take(&mut self.block_labels)))
     }
 
@@ -402,10 +426,10 @@ impl Assembler {
 
         // Label-free constructor and deployment snippets need neither offset
         // discovery nor push-width relaxation.
-        if !program
-            .instructions
-            .iter()
-            .any(|inst| matches!(inst.kind(), AsmInstKind::Label(_) | AsmInstKind::PushLabel(_)))
+        if program.indexed_jump_tables.is_empty()
+            && !program.instructions.iter().any(|inst| {
+                matches!(inst.kind(), AsmInstKind::Label(_) | AsmInstKind::PushLabel(_))
+            })
         {
             let mut result =
                 self.emit_bytecode(&program, FxHashMap::default(), &FxHashMap::default());
@@ -480,6 +504,10 @@ impl Assembler {
                 }
             }
         }
+        for &(label, ref targets) in &program.indexed_jump_tables {
+            label_offsets.insert(label, offset);
+            offset += targets.len() * assembly::INDEXED_JUMP_STUB_LEN;
+        }
 
         // Compute new widths based on resolved offsets
         for (idx, inst) in program.instructions.iter().enumerate() {
@@ -533,6 +561,26 @@ impl Assembler {
                 }
             }
         }
+        for (_, targets) in &program.indexed_jump_tables {
+            for &target in targets {
+                out.emit_op(op::JUMPDEST);
+                let target_offset = label_offsets
+                    .get(&target)
+                    .copied()
+                    .unwrap_or_else(|| panic!("label {target:?} was never defined"));
+                assert!(
+                    out.push_width(U256::from(target_offset))
+                        <= assembly::INDEXED_JUMP_TARGET_WIDTH,
+                    "label {target:?} does not fit PUSH{}",
+                    assembly::INDEXED_JUMP_TARGET_WIDTH
+                );
+                out.emit_push_fixed_width(
+                    U256::from(target_offset),
+                    assembly::INDEXED_JUMP_TARGET_WIDTH,
+                );
+                out.emit_op(op::JUMP);
+            }
+        }
 
         out.finish()
     }
@@ -723,6 +771,48 @@ PUSH4 0x80000000
             str![[r#"
 PUSH1 0x02
 
+"#]]
+        );
+    }
+
+    #[test]
+    fn indexed_jump_uses_out_of_line_fixed_stride_table() {
+        let mut asm = Assembler::with_config(AssemblerConfig {
+            evm_version: EvmVersion::Shanghai,
+            optimization: OptimizationMode::None,
+            ..AssemblerConfig::default()
+        });
+        let left = asm.new_label();
+        let right = asm.new_label();
+
+        asm.emit_push(U256::from(1));
+        asm.emit_indexed_jump(vec![left, right]);
+        asm.define_label(left);
+        asm.emit_op(op::STOP);
+        asm.define_label(right);
+        asm.emit_op(op::INVALID);
+
+        assert_data_eq!(
+            disassemble(&asm.assemble().bytecode),
+            str![[r#"
+PUSH1 0x01
+PUSH1 0x06
+MUL
+PUSH1 0x0e
+ADD
+JUMP
+JUMPDEST
+STOP
+JUMPDEST
+INVALID
+STOP
+JUMPDEST
+PUSH3 0x000009
+JUMP
+JUMPDEST
+PUSH3 0x00000b
+JUMP
+
 "#]]
         );
     }