Improve storeN instruction encoding with immediates (#1194)

* replace utility method with map_const method

* add Wrap utility trait for translator

* apply rustfmt

* improve encoding of wrapping istore instructions

This gets rid of unnecessary allocated function local constants for the immediates of wrapping istore instructions by integrating the wrapping in the translator. Before this commit the translator did not wrap the `value` to store before trying to use it and eventually convert it into a function local constant. For instructions such as i32_store8 (and more) this will never again allocate function local constants now because the wrapped i8 always fits in the `Instruction::I32Store8` instruction's `value` field, etc..

* clean-up use loop

* add test_store_wrap_imm test method

* use new test method for i32_store8 translation test

* use for loops in tests instead of repeating code

* improve i32_store8 translation tests with immediate ptr+offset

* clean-up test code a bit

* fix last remaining i32_store8 translation tests

* update translation tests for i32_store16

* update translation tests for i64_store8

* update translation tests for i64_store16

* fix store_offset16_imm based translation tests

* extend values for some at_imm translation tests

* test both mem0 and mem1 in test_store_at_overflow

* update translation tests for i64_store32

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

@@ -60,6 +60,7 @@ use crate::{
 };
 use core::fmt;
 use std::vec::Vec;
+use utils::Wrap;
 use wasmparser::{
     BinaryReaderError,
     FuncToValidate,
@@ -1918,6 +1919,37 @@ impl FuncTranslator {
         Ok(())
     }
 
+    /// Translates non-wrapping Wasm integer `store` to Wasmi bytecode.
+    ///
+    /// # Note
+    ///
+    /// Convenience method that simply forwards to [`Self::translate_istore_wrap`].
+    #[allow(clippy::too_many_arguments)]
+    fn translate_istore<Src, Field>(
+        &mut self,
+        memarg: MemArg,
+        make_instr: fn(ptr: Reg, memory: index::Memory) -> Instruction,
+        make_instr_imm: fn(ptr: Reg, memory: index::Memory) -> Instruction,
+        make_instr_offset16: fn(ptr: Reg, offset: u16, value: Reg) -> Instruction,
+        make_instr_offset16_imm: fn(ptr: Reg, offset: u16, value: Field) -> Instruction,
+        make_instr_at: fn(value: Reg, address: u32) -> Instruction,
+        make_instr_at_imm: fn(value: Field, address: u32) -> Instruction,
+    ) -> Result<(), Error>
+    where
+        Src: Copy + From<TypedVal>,
+        Field: TryFrom<Src> + Into<AnyConst16>,
+    {
+        self.translate_istore_wrap::<Src, Src, Field>(
+            memarg,
+            make_instr,
+            make_instr_imm,
+            make_instr_offset16,
+            make_instr_offset16_imm,
+            make_instr_at,
+            make_instr_at_imm,
+        )
+    }
+
     /// Translates Wasm integer `store` and `storeN` instructions to Wasmi bytecode.
     ///
     /// # Note
@@ -1931,27 +1963,27 @@ impl FuncTranslator {
     ///
     /// - `{i32, i64}.{store, store8, store16, store32}`
     #[allow(clippy::too_many_arguments)]
-    fn translate_istore<T, U>(
+    fn translate_istore_wrap<Src, Wrapped, Field>(
         &mut self,
         memarg: MemArg,
         make_instr: fn(ptr: Reg, memory: index::Memory) -> Instruction,
         make_instr_imm: fn(ptr: Reg, memory: index::Memory) -> Instruction,
         make_instr_offset16: fn(ptr: Reg, offset: u16, value: Reg) -> Instruction,
-        make_instr_offset16_imm: fn(ptr: Reg, offset: u16, value: U) -> Instruction,
+        make_instr_offset16_imm: fn(ptr: Reg, offset: u16, value: Field) -> Instruction,
         make_instr_at: fn(value: Reg, address: u32) -> Instruction,
-        make_instr_at_imm: fn(value: U, address: u32) -> Instruction,
+        make_instr_at_imm: fn(value: Field, address: u32) -> Instruction,
     ) -> Result<(), Error>
     where
-        T: Copy + From<TypedVal>,
-        U: TryFrom<T> + Into<AnyConst16>,
+        Src: Copy + Wrap<Wrapped> + From<TypedVal>,
+        Field: TryFrom<Wrapped> + Into<AnyConst16>,
     {
         bail_unreachable!(self);
         let (memory, offset) = Self::decode_memarg(memarg);
         let (ptr, value) = self.alloc.stack.pop2();
         let ptr = match ptr {
             Provider::Register(ptr) => ptr,
             Provider::Const(ptr) => {
-                return self.translate_istore_at::<T, U>(
+                return self.translate_istore_wrap_at::<Src, Wrapped, Field>(
                     memory,
                     u32::from(ptr),
                     offset,
@@ -1962,7 +1994,7 @@ impl FuncTranslator {
             }
         };
         if memory.is_default() {
-            if let Some(_instr) = self.translate_istore_mem0::<T, U>(
+            if let Some(_instr) = self.translate_istore_wrap_mem0::<Src, Wrapped, Field>(
                 ptr,
                 offset,
                 value,
@@ -1977,7 +2009,7 @@ impl FuncTranslator {
                 make_instr(ptr, memory),
                 Instruction::register_and_imm32(value, offset),
             ),
-            TypedProvider::Const(value) => match U::try_from(T::from(value)).ok() {
+            TypedProvider::Const(value) => match Field::try_from(Src::from(value).wrap()).ok() {
                 Some(value) => (
                     make_instr_imm(ptr, memory),
                     Instruction::imm16_and_imm32(value, offset),
@@ -1998,18 +2030,18 @@ impl FuncTranslator {
     /// # Note
     ///
     /// This is used in cases where the `ptr` is a known constant value.
-    fn translate_istore_at<T, U>(
+    fn translate_istore_wrap_at<Src, Wrapped, Field>(
         &mut self,
         memory: index::Memory,
         ptr: u32,
         offset: u32,
         value: TypedProvider,
         make_instr_at: fn(value: Reg, address: u32) -> Instruction,
-        make_instr_at_imm: fn(value: U, address: u32) -> Instruction,
+        make_instr_at_imm: fn(value: Field, address: u32) -> Instruction,
     ) -> Result<(), Error>
     where
-        T: Copy + From<TypedVal>,
-        U: TryFrom<T>,
+        Src: Copy + From<TypedVal> + Wrap<Wrapped>,
+        Field: TryFrom<Wrapped>,
     {
         let Some(address) = Self::effective_address(ptr, offset) else {
             return self.translate_trap(TrapCode::MemoryOutOfBounds);
@@ -2019,7 +2051,7 @@ impl FuncTranslator {
                 self.push_fueled_instr(make_instr_at(value, address), FuelCosts::store)?;
             }
             Provider::Const(value) => {
-                if let Ok(value) = U::try_from(T::from(value)) {
+                if let Ok(value) = Field::try_from(Src::from(value).wrap()) {
                     self.push_fueled_instr(make_instr_at_imm(value, address), FuelCosts::store)?;
                 } else {
                     let value = self.alloc.stack.alloc_const(value)?;
@@ -2042,17 +2074,17 @@ impl FuncTranslator {
     /// This optimizes for cases where the Wasm linear memory that is operated on is known
     /// to be the default memory.
     /// Returns `Some` in case the optimized instructions have been encoded.
-    fn translate_istore_mem0<T, U>(
+    fn translate_istore_wrap_mem0<Src, Wrapped, Field>(
         &mut self,
         ptr: Reg,
         offset: u32,
         value: TypedProvider,
         make_instr_offset16: fn(Reg, u16, Reg) -> Instruction,
-        make_instr_offset16_imm: fn(Reg, u16, U) -> Instruction,
+        make_instr_offset16_imm: fn(Reg, u16, Field) -> Instruction,
     ) -> Result<Option<Instr>, Error>
     where
-        T: Copy + From<TypedVal>,
-        U: TryFrom<T>,
+        Src: Copy + From<TypedVal> + Wrap<Wrapped>,
+        Field: TryFrom<Wrapped>,
     {
         let Ok(offset16) = u16::try_from(offset) else {
             return Ok(None);
@@ -2061,7 +2093,7 @@ impl FuncTranslator {
             Provider::Register(value) => {
                 self.push_fueled_instr(make_instr_offset16(ptr, offset16, value), FuelCosts::store)?
             }
-            Provider::Const(value) => match U::try_from(T::from(value)) {
+            Provider::Const(value) => match Field::try_from(Src::from(value).wrap()) {
                 Ok(value) => self.push_fueled_instr(
                     make_instr_offset16_imm(ptr, offset16, value),
                     FuelCosts::store,

crates/wasmi/src/engine/translator/tests/op/store/i32_store16.rs

@@ -12,24 +12,23 @@ fn reg() {
 #[cfg_attr(miri, ignore)]
 fn imm() {
     let values = [
+        0,
+        1,
+        -1,
+        42,
         i32::from(i16::MIN) - 1,
+        i32::from(i16::MIN),
+        i32::from(i16::MIN + 1),
+        i32::from(i16::MAX - 1),
+        i32::from(i16::MAX),
         i32::from(i16::MAX) + 1,
         i32::MIN,
         i32::MIN + 1,
         i32::MAX,
         i32::MAX - 1,
     ];
     for value in values {
-        test_store_imm::<i32>(WASM_OP, value, Instruction::i32_store16);
-    }
-}
-
-#[test]
-#[cfg_attr(miri, ignore)]
-fn imm16() {
-    let values = [0, 1, -1, 42, i16::MIN + 1, i16::MIN, i16::MAX - 1, i16::MAX];
-    for value in values {
-        test_store_imm16::<i16>(WASM_OP, Instruction::i32_store16_imm, value);
+        test_store_wrap_imm::<i32, i16, i16>(WASM_OP, value, Instruction::i32_store16_imm);
     }
 }
 
@@ -42,32 +41,36 @@ fn offset16() {
 #[test]
 #[cfg_attr(miri, ignore)]
 fn offset16_imm() {
-    test_store_offset16_imm::<i32>(
-        WASM_OP,
+    let values = [
+        0,
+        1,
+        -1,
+        -1000,
+        1000,
+        i32::from(i8::MIN) - 1,
+        i32::from(i8::MIN),
+        i32::from(i8::MIN) + 1,
+        i32::from(i8::MAX) - 1,
+        i32::from(i8::MAX),
+        i32::from(i8::MAX) + 1,
         i32::from(i16::MIN) - 1,
-        Instruction::i32_store16_offset16,
-    );
-    test_store_offset16_imm::<i32>(
-        WASM_OP,
+        i32::from(i16::MIN),
+        i32::from(i16::MIN) + 1,
+        i32::from(i16::MAX) - 1,
+        i32::from(i16::MAX),
         i32::from(i16::MAX) + 1,
-        Instruction::i32_store16_offset16,
-    );
-    test_store_offset16_imm::<i32>(WASM_OP, i32::MIN + 1, Instruction::i32_store16_offset16);
-    test_store_offset16_imm::<i32>(WASM_OP, i32::MAX - 1, Instruction::i32_store16_offset16);
-    test_store_offset16_imm::<i32>(WASM_OP, i32::MIN, Instruction::i32_store16_offset16);
-    test_store_offset16_imm::<i32>(WASM_OP, i32::MAX, Instruction::i32_store16_offset16);
-}
-
-#[test]
-#[cfg_attr(miri, ignore)]
-fn offset16_imm16() {
-    test_store_offset16_imm16::<i16>(WASM_OP, 0, Instruction::i32_store16_offset16_imm);
-    test_store_offset16_imm16::<i16>(WASM_OP, 1, Instruction::i32_store16_offset16_imm);
-    test_store_offset16_imm16::<i16>(WASM_OP, -1, Instruction::i32_store16_offset16_imm);
-    test_store_offset16_imm16::<i16>(WASM_OP, i16::MIN + 1, Instruction::i32_store16_offset16_imm);
-    test_store_offset16_imm16::<i16>(WASM_OP, i16::MAX - 1, Instruction::i32_store16_offset16_imm);
-    test_store_offset16_imm16::<i16>(WASM_OP, i16::MIN, Instruction::i32_store16_offset16_imm);
-    test_store_offset16_imm16::<i16>(WASM_OP, i16::MAX, Instruction::i32_store16_offset16_imm);
+        i32::MIN,
+        i32::MIN + 1,
+        i32::MAX - 1,
+        i32::MAX,
+    ];
+    for value in values {
+        test_store_wrap_offset16_imm::<i32, i16, i16>(
+            WASM_OP,
+            value,
+            Instruction::i32_store16_offset16_imm,
+        );
+    }
 }
 
 #[test]
@@ -85,22 +88,30 @@ fn at_overflow() {
 #[test]
 #[cfg_attr(miri, ignore)]
 fn at_imm() {
-    test_store_at_imm::<i32>(
-        WASM_OP,
-        i32::from(i16::MAX) + 1,
-        Instruction::i32_store16_at,
-    );
-    test_store_at_imm::<i32>(WASM_OP, i32::MAX - 1, Instruction::i32_store16_at);
-    test_store_at_imm::<i32>(WASM_OP, i32::MAX, Instruction::i32_store16_at);
+    let values = [
+        0,
+        1,
+        -1000,
+        1000,
+        i32::from(i16::MIN),
+        i32::from(i16::MIN) + 1,
+        i32::from(i16::MAX) - 1,
+        i32::from(i16::MAX),
+        i32::MIN,
+        i32::MIN + 1,
+        i32::MAX - 1,
+        i32::MAX,
+    ];
+    for value in values {
+        test_store_wrap_at_imm::<i32, i16, i16>(WASM_OP, value, Instruction::i32_store16_at_imm);
+    }
 }
 
 #[test]
 #[cfg_attr(miri, ignore)]
 fn imm_at_overflow() {
-    test_store_at_imm_overflow(WASM_OP, 0);
-    test_store_at_imm_overflow(WASM_OP, 1);
-    test_store_at_imm_overflow(WASM_OP, -1);
-    test_store_at_imm_overflow(WASM_OP, 42);
-    test_store_at_imm_overflow(WASM_OP, i32::MIN);
-    test_store_at_imm_overflow(WASM_OP, i32::MAX);
+    let values = [0, 1, -1, i32::MIN, i32::MAX];
+    for value in values {
+        test_store_at_imm_overflow(WASM_OP, value);
+    }
 }