Add JumpTable as Terminator

core/engine/src/optimizer/control_flow_graph/basic_block.rs

@@ -2,7 +2,7 @@ use std::hash::Hash;
 
 use bitflags::bitflags;
 
-use crate::vm::Instruction;
+use crate::vm::{Handler, Instruction};
 
 use super::{BasicBlockKey, Terminator};
 
@@ -13,13 +13,18 @@ bitflags! {
     }
 }
 
+pub(crate) struct BasicBlockHandler {
+    basic_block: BasicBlockKey,
+    handler: Handler,
+}
+
 /// TODO: doc
 #[derive(Default, Clone)]
 pub struct BasicBlock {
     pub(crate) previous: Vec<BasicBlockKey>,
     pub(crate) instructions: Vec<Instruction>,
     pub(crate) terminator: Terminator,
-    pub(crate) handler: Option<BasicBlockKey>,
+    pub(crate) handler: Option<(BasicBlockKey, Handler)>,
 
     pub(crate) flags: BasicBlockFlags,
 }
@@ -61,17 +66,21 @@ impl BasicBlock {
     }
 
     pub(crate) fn next_into(&self, nexts: &mut Vec<BasicBlockKey>) {
-        match self.terminator {
+        match &self.terminator {
             Terminator::None => {}
             Terminator::JumpUnconditional { target, .. } => {
-                nexts.push(target);
+                nexts.push(*target);
             }
             Terminator::JumpConditional { no, yes, .. }
             | Terminator::TemplateLookup { no, yes, .. } => {
-                nexts.push(no);
-                nexts.push(yes);
+                nexts.push(*no);
+                nexts.push(*yes);
+            }
+            Terminator::JumpTable { default, addresses } => {
+                nexts.push(*default);
+                nexts.extend_from_slice(addresses);
             }
-            Terminator::Return { end } => nexts.push(end),
+            Terminator::Return { end } => nexts.push(*end),
         }
     }
 }

core/engine/src/optimizer/control_flow_graph/mod.rs

@@ -9,15 +9,17 @@ use std::fmt::Debug;
 
 use rustc_hash::FxHashMap;
 use slotmap::{new_key_type, SlotMap};
+use thin_vec::ThinVec;
 
 use crate::vm::{CodeBlock, Handler, Instruction, InstructionIterator, Opcode};
 
 pub use self::basic_block::BasicBlock;
 
 new_key_type! { pub(crate) struct BasicBlockKey; }
+new_key_type! { pub(crate) struct HandlerKey; }
 
 /// TODO: doc
-#[derive(Default, Clone, Copy)]
+#[derive(Default, Clone)]
 pub(crate) enum Terminator {
     /// TODO: doc
     #[default]
@@ -53,6 +55,11 @@ pub(crate) enum Terminator {
         site: u64,
     },
 
+    JumpTable {
+        default: BasicBlockKey,
+        addresses: ThinVec<BasicBlockKey>,
+    },
+
     /// TODO: doc
     Return { end: BasicBlockKey },
 }
@@ -91,6 +98,7 @@ pub struct ControlFlowGraph {
     basic_block_start: BasicBlockKey,
     basic_block_end: BasicBlockKey,
     basic_blocks: SlotMap<BasicBlockKey, BasicBlock>,
+    // handlers: SlotMap<HandlerKey, BasicBlockHandler>,
 }
 
 impl Debug for ControlFlowGraph {
@@ -148,7 +156,7 @@ impl Debug for ControlFlowGraph {
             }
 
             if let Some(handler) = &basic_block.handler {
-                let index = index_from_basic_block(*handler);
+                let index = index_from_basic_block(handler.0);
                 write!(f, " -- handler B{index}")?;
             }
 
@@ -179,6 +187,21 @@ impl Debug for ControlFlowGraph {
                         let target = index_from_basic_block(*yes);
                         write!(f, "TemplateLookup B{target}")?;
                     }
+                    Terminator::JumpTable { default, addresses } => {
+                        let default = index_from_basic_block(*default);
+                        write!(f, "JumpTable default: B{default}")?;
+
+                        if !addresses.is_empty() {
+                            write!(f, " ")?;
+                            for (i, address) in addresses.iter().enumerate() {
+                                let address = index_from_basic_block(*address);
+                                write!(f, "{}: B{address}", i + 1)?;
+                                if i + 1 != addresses.len() {
+                                    write!(f, ", ")?;
+                                }
+                            }
+                        }
+                    }
                     Terminator::Return { end } => {
                         let target = index_from_basic_block(*end);
                         write!(f, "Return B{target}")?;
@@ -234,6 +257,10 @@ impl ControlFlowGraph {
                     leaders.push(iter.pc() as u32);
                     leaders.push(exit);
                 }
+                Instruction::JumpTable { default, addresses } => {
+                    leaders.push(default);
+                    leaders.extend_from_slice(&addresses);
+                }
                 instruction => {
                     if let Some(target) = is_jump_kind_instruction(&instruction) {
                         leaders.push(iter.pc() as u32);
@@ -297,9 +324,9 @@ impl ControlFlowGraph {
                 .rev()
                 .find(|handler| handler.contains(iter.pc() as u32));
             if let Some(handler) = handler {
-                let handler = basic_block_from_bytecode_position(handler.handler());
+                let basic_block_handler = basic_block_from_bytecode_position(handler.handler());
 
-                basic_blocks[key].handler = Some(handler);
+                basic_blocks[key].handler = Some((basic_block_handler, *handler));
             }
 
             let mut bytecode = Vec::new();
@@ -333,6 +360,23 @@ impl ControlFlowGraph {
 
                         terminator = Terminator::TemplateLookup { no, yes, site };
                     }
+                    Instruction::JumpTable { default, addresses } => {
+                        let default = basic_block_from_bytecode_position(default);
+
+                        basic_blocks[default].previous.push(key);
+
+                        let mut basic_block_addresses = ThinVec::with_capacity(addresses.len());
+                        for address in addresses {
+                            let address = basic_block_from_bytecode_position(address);
+                            basic_blocks[address].previous.push(key);
+                            basic_block_addresses.push(address);
+                        }
+
+                        terminator = Terminator::JumpTable {
+                            default,
+                            addresses: basic_block_addresses,
+                        }
+                    }
                     instruction => {
                         if let Some(address) = is_jump_kind_instruction(&instruction) {
                             let yes = basic_block_from_bytecode_position(address);
@@ -427,6 +471,24 @@ impl ControlFlowGraph {
                     let target = index_from_basic_block(*yes);
                     labels.push((start as u32, target));
                 }
+                Terminator::JumpTable { default, addresses } => {
+                    results.extend_from_slice(&[Opcode::JumpTable as u8]);
+                    let start = results.len();
+                    results.extend_from_slice(&[0, 0, 0, 0]);
+
+                    let default = index_from_basic_block(*default);
+                    labels.push((start as u32, default));
+
+                    results.extend_from_slice(&(addresses.len() as u32).to_ne_bytes());
+
+                    for address in addresses {
+                        let start = results.len();
+                        results.extend_from_slice(&[0, 0, 0, 0]);
+
+                        let target = index_from_basic_block(*address);
+                        labels.push((start as u32, target));
+                    }
+                }
                 Terminator::Return { .. } => {
                     results.push(Opcode::Return as u8);
                 }
@@ -561,4 +623,22 @@ mod test {
 
         assert_eq!(bytecode, actual.as_slice());
     }
+
+    #[test]
+    fn preserve_jump_table() {
+        let bytecode = &[
+            193, 68, 0, 18, 67, 1, 72, 1, 140, 1, 76, 153, 2, 18, 155, 6, 17, 118, 38, 0, 0, 0,
+            155, 5, 17, 118, 38, 0, 0, 0, 126, 16, 118, 38, 0, 0, 0, 16, 151, 153, 3, 18, 71, 1,
+            143, 0, 152, 155, 152, 120, 55, 0, 0, 0, 124, 123, 71, 0, 0, 0, 1, 0, 0, 0, 68, 0, 0,
+            0, 152, 147, 148, 18, 152, 147, 148,
+        ];
+
+        let graph = ControlFlowGraph::generate(bytecode, &[]);
+
+        println!("{graph:?}");
+
+        let actual = graph.finalize();
+
+        assert_eq!(bytecode, actual.as_slice());
+    }
 }