Implement re-designed Wasmi br_table instructions (#1158)

* initial commit

* redo Wasmi `br_table` instruction design

This new design accounts for Wasmi `br_table`s where each `br_table` target may need to copy values to different registers before taking a branch which was an oversight in the previous design.

* remove no longer needed constructor

* remove unnecessary provider2reg wrapper function

* properly encode non-overlapping branch table targets

* add missing docs

* revert changes to translate_copy_branch_params method

* fix oversight/bug in constructors

* add some more br_table tests

* fix clippy warnings in tests

* add more br_table translation tests with immediate inputs

crates/wasmi/src/engine/bytecode/construct.rs

@@ -202,14 +202,67 @@ impl Instruction {
         Self::branch_i64_ne_imm(condition, 0_i16, offset)
     }
 
-    /// Creates a new [`Instruction::BranchTable`] for the given `index` and `len_targets`.
-    pub fn branch_table(index: Register, len_targets: impl Into<Const32<u32>>) -> Self {
-        Self::BranchTable {
-            index,
-            len_targets: len_targets.into(),
+    /// Creates a new [`Instruction::BranchTable0`] for the given `index` and `len_targets`.
+    pub fn branch_table_0(index: impl Into<Register>, len_targets: u32) -> Self {
+        Self::BranchTable0 {
+            index: index.into(),
+            len_targets,
+        }
+    }
+
+    /// Creates a new [`Instruction::BranchTable1`] for the given `index` and `len_targets`.
+    pub fn branch_table_1(index: impl Into<Register>, len_targets: u32) -> Self {
+        Self::BranchTable1 {
+            index: index.into(),
+            len_targets,
+        }
+    }
+
+    /// Creates a new [`Instruction::BranchTable2`] for the given `index` and `len_targets`.
+    pub fn branch_table_2(index: impl Into<Register>, len_targets: u32) -> Self {
+        Self::BranchTable2 {
+            index: index.into(),
+            len_targets,
+        }
+    }
+
+    /// Creates a new [`Instruction::BranchTable3`] for the given `index` and `len_targets`.
+    pub fn branch_table_3(index: impl Into<Register>, len_targets: u32) -> Self {
+        Self::BranchTable3 {
+            index: index.into(),
+            len_targets,
+        }
+    }
+
+    /// Creates a new [`Instruction::BranchTableSpan`] for the given `index` and `len_targets`.
+    pub fn branch_table_span(index: impl Into<Register>, len_targets: u32) -> Self {
+        Self::BranchTableSpan {
+            index: index.into(),
+            len_targets,
+        }
+    }
+
+    /// Creates a new [`Instruction::BranchTableMany`] for the given `index` and `len_targets`.
+    pub fn branch_table_many(index: impl Into<Register>, len_targets: u32) -> Self {
+        Self::BranchTableMany {
+            index: index.into(),
+            len_targets,
         }
     }
 
+    /// Creates a new [`Instruction::BranchTableTarget`] for the given `index` and `len_targets`.
+    pub fn branch_table_target(results: RegisterSpan, offset: BranchOffset) -> Self {
+        Self::BranchTableTarget { results, offset }
+    }
+
+    /// Creates a new [`Instruction::BranchTableTargetNonOverlapping`] for the given `index` and `len_targets`.
+    pub fn branch_table_target_non_overlapping(
+        results: RegisterSpan,
+        offset: BranchOffset,
+    ) -> Self {
+        Self::BranchTableTargetNonOverlapping { results, offset }
+    }
+
     /// Creates a new [`Instruction::Copy`].
     pub fn copy(result: impl Into<Register>, value: impl Into<Register>) -> Self {
         Self::Copy {
@@ -1001,6 +1054,11 @@ impl Instruction {
         Self::RegisterList([reg0.into(), reg1.into(), reg2.into()])
     }
 
+    /// Creates a new [`Instruction::RegisterSpan`].
+    pub fn register_span(span: RegisterSpanIter) -> Self {
+        Self::RegisterSpan(span)
+    }
+
     /// Creates a new [`Instruction::CallIndirectParams`] for the given `index` and `table`.
     pub fn call_indirect_params(index: Register, table: impl Into<TableIdx>) -> Self {
         Self::CallIndirectParams(CallIndirectParams {

crates/wasmi/src/engine/bytecode/mod.rs

@@ -552,25 +552,113 @@ pub enum Instruction {
     /// A fused [`Instruction::F64Ge`] and Wasm branch instruction.
     BranchF64Ge(BranchBinOpInstr),
 
-    /// A Wasm `br_table` instruction.
+    /// A Wasm `br_table` equivalent Wasmi instruction.
     ///
     /// # Encoding
     ///
-    /// 1. May be followed by one of the copy instructions.
-    /// 1. Must be followed `len_targets` times by any of:
+    /// Followed `len_target` times by
     ///
     /// - [`Instruction::Branch`]
     /// - [`Instruction::Return`]
+    BranchTable0 {
+        /// The register holding the index of the instruction.
+        index: Register,
+        /// The number of branch table targets including the default target.
+        len_targets: u32,
+    },
+    /// A Wasm `br_table` equivalent Wasmi instruction.
+    ///
+    /// # Encoding
+    ///
+    /// 1. Followed by one of
+    ///
+    /// - [`Instruction::Register`]
+    /// - [`Instruction::Const32`]
+    /// - [`Instruction::I64Const32`]
+    /// - [`Instruction::F64Const32`]
+    ///
+    /// 2. Followed `len_target` times by
+    ///
+    /// - [`Instruction::BranchTableTarget`]
     /// - [`Instruction::ReturnReg`]
     /// - [`Instruction::ReturnImm32`]
     /// - [`Instruction::ReturnI64Imm32`]
     /// - [`Instruction::ReturnF64Imm32`]
+    BranchTable1 {
+        /// The register holding the index of the instruction.
+        index: Register,
+        /// The number of branch table targets including the default target.
+        len_targets: u32,
+    },
+    /// A Wasm `br_table` equivalent Wasmi instruction.
+    ///
+    /// # Encoding
+    ///
+    /// 1. Followed by [`Instruction::Register2`].
+    /// 2. Followed `len_target` times by
+    ///
+    /// - [`Instruction::BranchTableTarget`]
+    /// - [`Instruction::ReturnReg2`]
+    BranchTable2 {
+        /// The register holding the index of the instruction.
+        index: Register,
+        /// The number of branch table targets including the default target.
+        len_targets: u32,
+    },
+    /// A Wasm `br_table` equivalent Wasmi instruction.
+    ///
+    /// # Encoding
+    ///
+    /// 1. Followed by [`Instruction::Register3`].
+    /// 2. Followed `len_target` times by
+    ///
+    /// - [`Instruction::BranchTableTarget`]
+    /// - [`Instruction::ReturnReg3`]
+    BranchTable3 {
+        /// The register holding the index of the instruction.
+        index: Register,
+        /// The number of branch table targets including the default target.
+        len_targets: u32,
+    },
+    /// A Wasm `br_table` equivalent Wasmi instruction.
+    ///
+    /// # Note
+    ///
+    /// All branch table targets must share the same destination registers.
+    ///
+    /// # Encoding
+    ///
+    /// 1. Followed by one of [`Instruction::RegisterSpan`].
+    /// 2. Followed `len_target` times by
+    ///
+    /// - [`Instruction::BranchTableTarget`]
+    /// - [`Instruction::BranchTableTargetNonOverlapping`]
     /// - [`Instruction::ReturnSpan`]
-    BranchTable {
+    BranchTableSpan {
+        /// The register holding the index of the instruction.
+        index: Register,
+        /// The number of branch table targets including the default target.
+        len_targets: u32,
+    },
+    /// A Wasm `br_table` equivalent Wasmi instruction.
+    ///
+    /// # Note
+    ///
+    /// All branch table targets must share the same destination registers.
+    ///
+    /// # Encoding
+    ///
+    /// 1. Followed by [`Instruction::RegisterList`] encoding.
+    /// 2. Followed `len_target` times by
+    ///
+    /// - [`Instruction::BranchTableTarget`]
+    /// - [`Instruction::BranchTableTargetNonOverlapping`]
+    /// - [`Instruction::Return`]
+    BranchTableMany {
         /// The register holding the index of the instruction.
         index: Register,
         /// The number of branch table targets including the default target.
-        len_targets: Const32<u32>,
+        len_targets: u32,
     },
 
     /// Copies `value` to `result`.
@@ -3177,6 +3265,36 @@ pub enum Instruction {
     /// This [`Instruction`] only acts as a parameter to another
     /// one and will never be executed itself directly.
     F64Const32(Const32<f64>),
+    /// A Wasm `br_table` branching target which copies values before branching.
+    ///
+    /// # Encoding
+    ///
+    /// This always follows
+    ///
+    /// - [`Instruction::BranchTable1`]
+    /// - [`Instruction::BranchTable2`]
+    /// - [`Instruction::BranchTableSpan`]
+    /// - [`Instruction::BranchTableMany`]
+    BranchTableTarget {
+        /// The registers where the values are going to be copied.
+        results: RegisterSpan,
+        /// The branching offset of the branch table target.
+        offset: BranchOffset,
+    },
+    /// A Wasm `br_table` branching target which copies overlapping values before branching.
+    ///
+    /// # Encoding
+    ///
+    /// This always follows
+    ///
+    /// - [`Instruction::BranchTableSpan`]
+    /// - [`Instruction::BranchTableMany`]
+    BranchTableTargetNonOverlapping {
+        /// The registers where the values are going to be copied.
+        results: RegisterSpan,
+        /// The branching offset of the branch table target.
+        offset: BranchOffset,
+    },
     /// An instruction parameter with a [`Register`] and a 32-bit immediate value.
     RegisterAndImm32 {
         /// The [`Register`] parameter value.
@@ -3188,6 +3306,8 @@ pub enum Instruction {
         /// The 32-bit immediate value.
         imm: AnyConst32,
     },
+    /// A [`RegisterSpanIter`] instruction parameter.
+    RegisterSpan(RegisterSpanIter),
     /// A [`Register`] instruction parameter.
     ///
     /// # Note

crates/wasmi/src/engine/bytecode/utils.rs

@@ -229,6 +229,11 @@ impl Iterator for RegisterSpanIter {
         self.next = self.next.next();
         Some(reg)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.len_as_u16() as usize;
+        (remaining, Some(remaining))
+    }
 }
 
 impl DoubleEndedIterator for RegisterSpanIter {

crates/wasmi/src/engine/executor/instrs.rs

@@ -213,8 +213,23 @@ impl<'engine> Executor<'engine> {
                     ))
                 }
                 Instr::Branch { offset } => self.execute_branch(offset),
-                Instr::BranchTable { index, len_targets } => {
-                    self.execute_branch_table(index, len_targets)
+                Instr::BranchTable0 { index, len_targets } => {
+                    self.execute_branch_table_0(index, len_targets)
+                }
+                Instr::BranchTable1 { index, len_targets } => {
+                    self.execute_branch_table_1(index, len_targets)
+                }
+                Instr::BranchTable2 { index, len_targets } => {
+                    self.execute_branch_table_2(index, len_targets)
+                }
+                Instr::BranchTable3 { index, len_targets } => {
+                    self.execute_branch_table_3(index, len_targets)
+                }
+                Instr::BranchTableSpan { index, len_targets } => {
+                    self.execute_branch_table_span(index, len_targets)
+                }
+                Instr::BranchTableMany { index, len_targets } => {
+                    self.execute_branch_table_many(index, len_targets)
                 }
                 Instr::BranchCmpFallback { lhs, rhs, params } => {
                     self.execute_branch_cmp_fallback(lhs, rhs, params)
@@ -856,10 +871,13 @@ impl<'engine> Executor<'engine> {
                 | Instr::Const32(_)
                 | Instr::I64Const32(_)
                 | Instr::F64Const32(_)
+                | Instr::BranchTableTarget { .. }
+                | Instr::BranchTableTargetNonOverlapping { .. }
                 | Instr::Register(_)
                 | Instr::Register2(_)
                 | Instr::Register3(_)
                 | Instr::RegisterAndImm32 { .. }
+                | Instr::RegisterSpan(_)
                 | Instr::RegisterList(_)
                 | Instr::CallIndirectParams(_)
                 | Instr::CallIndirectParamsImm16(_) => self.invalid_instruction_word()?,
@@ -1160,6 +1178,20 @@ impl<'engine> Executor<'engine> {
         self.set_register(instr.result, op(lhs, rhs)?);
         self.try_next_instr()
     }
+
+    /// Skips all [`Instruction`]s belonging to an [`Instruction::RegisterList`] encoding.
+    fn skip_register_list(ip: InstructionPtr) -> InstructionPtr {
+        let mut ip = ip;
+        while let Instruction::RegisterList(_) = *ip.get() {
+            ip.add(1);
+        }
+        // We skip an additional `Instruction` because we know that `Instruction::RegisterList` is always followed by one of:
+        // - `Instruction::Register`
+        // - `Instruction::Register2`
+        // - `Instruction::Register3`.
+        ip.add(1);
+        ip
+    }
 }
 
 impl<'engine> Executor<'engine> {