Optimize shadow stack instruction sequences

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

@@ -334,6 +334,30 @@ impl Instruction {
         }
     }
 
+    /// Creates a new [`Instruction::I32AddImmIntoGlobal0`].
+    pub fn i32_add_imm_into_global_0(lhs: Register, rhs: impl Into<Const32<i32>>) -> Self {
+        Self::I32AddImmIntoGlobal0 {
+            lhs,
+            rhs: rhs.into(),
+        }
+    }
+
+    /// Creates a new [`Instruction::I32AddImmFromGlobal0`].
+    pub fn i32_add_imm_from_global_0(result: Register, rhs: impl Into<Const32<i32>>) -> Self {
+        Self::I32AddImmFromGlobal0 {
+            result,
+            rhs: rhs.into(),
+        }
+    }
+
+    /// Creates a new [`Instruction::I32AddImmInoutGlobal0`].
+    pub fn i32_add_imm_inout_global_0(result: Register, rhs: impl Into<Const32<i32>>) -> Self {
+        Self::I32AddImmInoutGlobal0 {
+            result,
+            rhs: rhs.into(),
+        }
+    }
+
     /// Creates a new [`Instruction::F32CopysignImm`] instruction.
     pub fn f32_copysign_imm(result: Register, lhs: Register, rhs: Sign) -> Self {
         Self::F32CopysignImm(BinInstrImm::new(result, lhs, rhs))

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

@@ -1030,6 +1030,91 @@ pub enum Instruction {
         input: Const16<i64>,
     },
 
+    /// Fused instruction to add a register and constant `i32` value to the `global` variable at index 0.
+    ///
+    /// This instruction represents the following Wasm instruction sequence:
+    ///
+    /// ```wat
+    /// local.get $v
+    /// i32.const $n
+    /// i32.add
+    /// global.set 0
+    /// ```
+    ///
+    /// Where `$v` is a local variable index and `$n` is a `i32` literal.
+    ///
+    /// # Note
+    ///
+    /// - This is an instruction primarily intended to optimize Wasm instruction
+    ///   sequences that emulate the shadow stack pointer arithmetic.
+    /// - We operate on the global variable at index 0 because that is the index
+    ///   that LLVM tends to use for its shadow stack pointer. Other Wasm producers
+    ///   have not been checked and might differ.
+    I32AddImmIntoGlobal0 {
+        /// The `lhs` register for the addition.
+        lhs: Register,
+        /// The `rhs` constant value for the addition.
+        rhs: Const32<i32>,
+    },
+    /// Fused instruction to add a constant `i32` value to the `global` variable at index 0.
+    ///
+    /// This instruction represents the following Wasm instruction sequence:
+    ///
+    /// ```wat
+    /// global.get 0
+    /// i32.const $n
+    /// i32.add
+    /// ```
+    ///
+    /// Where `$n` is a `i32` literal.
+    ///
+    /// # Note
+    ///
+    /// - This is an instruction primarily intended to optimize Wasm instruction
+    ///   sequences that emulate the shadow stack pointer arithmetic.
+    /// - We operate on the global variable at index 0 because that is the index
+    ///   that LLVM tends to use for its shadow stack pointer. Other Wasm producers
+    ///   have not been checked and might differ.
+    I32AddImmFromGlobal0 {
+        /// The register storing the result of the instruction.
+        result: Register,
+        /// The `rhs` constant value for the addition.
+        rhs: Const32<i32>,
+    },
+    /// Fused instruction to add a constant `i32` value to the `global` variable at index 0.
+    ///
+    /// The result of this addition is stored both into
+    ///
+    /// - the `result` register
+    /// - the `global` variable at index 0.
+    ///
+    /// This instruction represents the following Wasm instruction sequence:
+    ///
+    /// ```wat
+    /// global.get 0
+    /// i32.const $n
+    /// i32.add
+    /// local.tee $v
+    /// global.set 0
+    /// ```
+    ///
+    /// Where `$v` is a local variable index and `$n` is a `i32` literal.
+    ///
+    /// # Note
+    ///
+    /// - This is a refinement variant of [`Instruction::I32AddImmFromGlobal0`].
+    /// - This is an instruction primarily intended to optimize Wasm instruction
+    ///   sequences that emulate the shadow stack pointer arithmetic.
+    /// - We operate on the global variable at index 0 because that is the index
+    ///   that LLVM tends to use for its shadow stack pointer. Other Wasm producers
+    ///   have not been checked and might differ.
+    I32AddImmInoutGlobal0 {
+        /// The register storing the result of the instruction.
+        result: Register,
+        /// The `rhs` constant value for the addition.
+        rhs: Const32<i32>,
+    },
+
     /// Wasm `i32.load` equivalent Wasmi instruction.
     ///
     /// # Encoding

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

@@ -373,6 +373,15 @@
                 Instr::GlobalSetI64Imm16 { global, input } => {
                     self.execute_global_set_i64imm16(&mut store.inner, global, input)
                 }
+                Instr::I32AddImmIntoGlobal0 { lhs, rhs } => {
+                    self.execute_i32_add_imm_into_global_0(lhs, rhs)
+                }
+                Instr::I32AddImmFromGlobal0 { result, rhs } => {
+                    self.execute_i32_add_imm_from_global_0(result, rhs)
+                }
+                Instr::I32AddImmInoutGlobal0 { result, rhs } => {
+                    self.execute_i32_add_imm_inout_global_0(result, rhs)
+                }
                 Instr::I32Load(instr) => self.execute_i32_load(instr)?,
                 Instr::I32LoadAt(instr) => self.execute_i32_load_at(instr)?,
                 Instr::I32LoadOffset16(instr) => self.execute_i32_load_offset16(instr)?,

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

@@ -1,7 +1,7 @@
 use super::Executor;
 use crate::{
     core::{hint, UntypedVal},
-    engine::bytecode::{Const16, GlobalIdx, Register},
+    engine::bytecode::{Const16, Const32, GlobalIdx, Register},
     store::StoreInner,
 };
 
@@ -78,4 +78,38 @@
         };
         self.next_instr()
     }
+
+    /// Executes an [`Instruction::I32AddImmIntoGlobal0`].
+    #[inline(always)]
+    pub fn execute_i32_add_imm_into_global_0(&mut self, lhs: Register, rhs: Const32<i32>) {
+        let lhs: i32 = self.get_register_as(lhs);
+        let rhs: i32 = i32::from(rhs);
+        let result = lhs.wrapping_add(rhs);
+        // Safety: TODO
+        unsafe { self.cache.global.set(result.into()) };
+        self.next_instr()
+    }
+
+    /// Executes an [`Instruction::I32AddImmFromGlobal0`].
+    #[inline(always)]
+    pub fn execute_i32_add_imm_from_global_0(&mut self, result: Register, rhs: Const32<i32>) {
+        // Safety: TODO
+        let lhs: i32 = unsafe { self.cache.global.get() }.into();
+        let rhs: i32 = rhs.into();
+        self.set_register(result, lhs.wrapping_add(rhs));
+        self.next_instr()
+    }
+
+    /// Executes an [`Instruction::I32AddImmInoutGlobal0`].
+    #[inline(always)]
+    pub fn execute_i32_add_imm_inout_global_0(&mut self, result: Register, rhs: Const32<i32>) {
+        // Safety: TODO
+        let lhs: i32 = unsafe { self.cache.global.get() }.into();
+        let rhs: i32 = rhs.into();
+        let sum = lhs.wrapping_add(rhs);
+        // Safety: TODO
+        unsafe { self.cache.global.set(sum.into()) };
+        self.set_register(result, sum);
+        self.next_instr()
+    }
 }

crates/wasmi/src/engine/translator/instr_encoder.rs

@@ -896,6 +896,122 @@
         Ok(())
     }
 
+    /// Fuses a `global.get 0` and an `i32.add_imm` if possible.
+    ///
+    /// Returns `true` if `Instruction` fusion was successful, `false` otherwise.
+    pub fn fuse_global_get_i32_add_imm(
+        &mut self,
+        lhs: Register,
+        rhs: i32,
+        stack: &mut ValueStack,
+    ) -> Result<bool, Error> {
+        let Some(last_instr) = self.last_instr else {
+            // Without a last instruction there is no way to fuse.
+            return Ok(false);
+        };
+        let &Instruction::GlobalGet { result, global } = self.instrs.get(last_instr) else {
+            // It is only possible to fuse an `GlobalGet` with a `I32AddImm` instruction.
+            return Ok(false);
+        };
+        if u32::from(global) != 0 {
+            // There only is an optimized instruction for a global index of 0.
+            // This is because most Wasm producers use the global at index 0 for their shadow stack.
+            return Ok(false);
+        }
+        if !matches!(stack.get_register_space(result), RegisterSpace::Dynamic) {
+            // Due to observable state it is impossible to fuse `GlobalGet` that has a non-`dynamic` result.
+            return Ok(false);
+        };
+        if result != lhs {
+            // The `input` to `I32AddImm` must be the same as the result of `GetGlobal`.
+            return Ok(false);
+        }
+        let rhs = <Const32<i32>>::from(rhs);
+        let fused_instr = Instruction::i32_add_imm_from_global_0(result, rhs);
+        _ = mem::replace(self.instrs.get_mut(last_instr), fused_instr);
+        stack.push_register(result)?;
+        Ok(true)
+    }
+
+    /// Fuses the `global.set` instruction with its previous instruction if possible.
+    ///
+    /// Returns `true` if `Instruction` fusion was successful, `false` otherwise.
+    pub fn fuse_global_set(
+        &mut self,
+        global_index: u32,
+        input: Register,
+        stack: &mut ValueStack,
+    ) -> bool {
+        /// Returns `true` if the previous [`Instruction`] and
+        /// its `result` can be fused with the `global.set` and its `input`.
+        fn is_fusable(stack: &mut ValueStack, result: Register, input: Register) -> bool {
+            if !matches!(stack.get_register_space(result), RegisterSpace::Dynamic) {
+                // Due to observable state it is impossible to fuse an instruction that has a non-`dynamic` result.
+                return false;
+            };
+            if result != input {
+                // It is only possible to fuse the instructions if the `input` of `global.set`
+                // matches the `result` of the previous to-be-fused instruction.
+                return false;
+            }
+            true
+        }
+
+        if global_index != 0 {
+            // There only is an optimized instruction for a global index of 0.
+            // This is because most Wasm producers use the global at index 0 for their shadow stack.
+            return false;
+        }
+        let Some(last_instr) = self.last_instr else {
+            // Without a last instruction there is no way to fuse.
+            return false;
+        };
+        let fused_instr = match self.instrs.get(last_instr) {
+            Instruction::I32Add(instr) => {
+                if !is_fusable(stack, instr.result, input) {
+                    return false;
+                }
+                let Some(value) = stack.resolve_func_local_const(instr.rhs).map(i32::from) else {
+                    // It is only possiblet o fuse `I32Add` if its `rhs` is a function local constant.
+                    return false;
+                };
+                let lhs = instr.lhs;
+                let rhs = <Const32<i32>>::from(value);
+                Instruction::i32_add_imm_into_global_0(lhs, rhs)
+            }
+            Instruction::I32Sub(instr) => {
+                if !is_fusable(stack, instr.result, input) {
+                    return false;
+                }
+                let Some(value) = stack.resolve_func_local_const(instr.rhs).map(i32::from) else {
+                    // It is only possiblet o fuse `I32Add` if its `rhs` is a function local constant.
+                    return false;
+                };
+                let lhs = instr.lhs;
+                let rhs = <Const32<i32>>::from(-value);
+                Instruction::i32_add_imm_into_global_0(lhs, rhs)
+            }
+            Instruction::I32AddImm16(instr) => {
+                if !is_fusable(stack, instr.result, input) {
+                    return false;
+                }
+                let lhs = instr.reg_in;
+                let rhs = <Const32<i32>>::from(i32::from(instr.imm_in));
+                Instruction::i32_add_imm_into_global_0(lhs, rhs)
+            }
+            &Instruction::I32AddImmFromGlobal0 { result, rhs } => {
+                if result != input {
+                    // The `input` to `GlobalSet` must be the same as the result of `I32AddImmFromGlobal0`.
+                    return false;
+                }
+                Instruction::i32_add_imm_inout_global_0(result, rhs)
+            }
+            _ => return false,
+        };
+        _ = mem::replace(self.instrs.get_mut(last_instr), fused_instr);
+        true
+    }
+
     /// Translates a Wasm `i32.eqz` instruction.
     ///
     /// Tries to fuse `i32.eqz` with a previous `i32.{and,or,xor}` instruction if possible.

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

@@ -43,6 +43,7 @@
     core::{TrapCode, UntypedVal, ValType},
     engine::{
         bytecode::{
+            self,
             AnyConst32,
             Const16,
             Const32,
@@ -57,7 +58,7 @@
         BlockType,
         CompiledFunc,
     },
-    module::{FuncIdx, FuncTypeIdx, ModuleHeader},
+    module::{self, FuncIdx, FuncTypeIdx, ModuleHeader},
     Engine,
     Error,
     FuncType,
@@ -2553,4 +2554,41 @@
             fuel_info,
         )
     }
+
+    /// Translates a `global.set` with an immutable input.
+    fn translate_global_set_imm(
+        &mut self,
+        global: bytecode::GlobalIdx,
+        input: TypedVal,
+    ) -> Result<(), Error> {
+        let global_index = u32::from(global);
+        let (global_type, _init_value) = self
+            .module
+            .get_global(module::GlobalIdx::from(global_index));
+        debug_assert_eq!(global_type.content(), input.ty());
+        match global_type.content() {
+            ValType::I32 => {
+                if let Ok(value) = Const16::try_from(i32::from(input)) {
+                    self.push_fueled_instr(
+                        Instruction::global_set_i32imm16(global, value),
+                        FuelCosts::entity,
+                    )?;
+                    return Ok(());
+                }
+            }
+            ValType::I64 => {
+                if let Ok(value) = Const16::try_from(i64::from(input)) {
+                    self.push_fueled_instr(
+                        Instruction::global_set_i64imm16(global, value),
+                        FuelCosts::entity,
+                    )?;
+                    return Ok(());
+                }
+            }
+            _ => {}
+        };
+        let cref = self.alloc.stack.alloc_const(input)?;
+        self.push_fueled_instr(Instruction::global_set(global, cref), FuelCosts::entity)?;
+        Ok(())
+    }
 }

crates/wasmi/src/engine/translator/relink_result.rs

@@ -233,6 +233,9 @@
             I::GlobalSet { .. } | I::GlobalSetI32Imm16 { .. } | I::GlobalSetI64Imm16 { .. } => {
                 Ok(false)
             }
+            I::I32AddImmIntoGlobal0 { .. } => Ok(false),
+            I::I32AddImmFromGlobal0 { result, .. } => relink_simple(result, new_result, old_result),
+            I::I32AddImmInoutGlobal0 { result, .. } => relink_simple(result, new_result, old_result),
             I::I32Load(instr) |
             I::I64Load(instr) |
             I::F32Load(instr) |

crates/wasmi/src/engine/translator/stack/mod.rs

@@ -189,6 +189,13 @@ impl ValueStack {
         self.consts.alloc(value.into())
     }
 
+    /// Returns the [`UntypedVal`] of the function local constant value referred to by `register` if any.
+    ///
+    /// Returns `None` if the `register` does not refer to a function local constant.
+    pub fn resolve_func_local_const(&self, register: Register) -> Option<UntypedVal> {
+        self.consts.get(register)
+    }
+
     /// Returns the allocated function local constant values in reversed allocation order.
     ///
     /// # Note