Refactor WasmInteger trait (#1553)

Robbepop · web-flow · commit bf8e2274e6b9 · 2025-07-01T18:04:19.000+02:00
* add NonZero support to UntypedVal conversions

* refactor WasmInteger trait

* refactor WasmInteger users

* rename eq_zero -&gt; is_zero
diff --git a/crates/core/src/untyped.rs b/crates/core/src/untyped.rs
@@ -92,6 +92,12 @@ macro_rules! impl_write_as_for_int {
                     self.write_lo64(value as $as as _)
                 }
             }
+
+            impl WriteAs<::core::num::NonZero<$int>> for UntypedVal {
+                fn write_as(&mut self, value: ::core::num::NonZero<$int>) {
+                    <UntypedVal as WriteAs<$int>>::write_as(self, value.get())
+                }
+            }
         )*
     };
 }
@@ -106,10 +112,23 @@ macro_rules! impl_write_as_for_uint {
                     self.write_lo64(value as _)
                 }
             }
+
+            impl WriteAs<::core::num::NonZero<$int>> for UntypedVal {
+                fn write_as(&mut self, value: ::core::num::NonZero<$int>) {
+                    <UntypedVal as WriteAs<$int>>::write_as(self, value.get())
+                }
+            }
         )*
     };
 }
-impl_write_as_for_uint!(bool, u8, u16, u32, u64);
+impl_write_as_for_uint!(u8, u16, u32, u64);
+
+impl WriteAs<bool> for UntypedVal {
+    #[allow(clippy::cast_lossless)]
+    fn write_as(&mut self, value: bool) {
+        self.write_lo64(value as _)
+    }
+}
 
 macro_rules! impl_write_as_for_float {
     ( $( $float:ty ),* $(,)? ) => {
@@ -224,14 +243,27 @@ macro_rules! impl_from_unsigned_prim {
                     Self::from_bits64(value as _)
                 }
             }
+
+            impl From<::core::num::NonZero<$prim>> for UntypedVal {
+                fn from(value: ::core::num::NonZero<$prim>) -> Self {
+                    <_ as From<$prim>>::from(value.get())
+                }
+            }
         )*
     };
 }
 #[rustfmt::skip]
 impl_from_unsigned_prim!(
-    bool, u8, u16, u32, u64,
+    u8, u16, u32, u64,
 );
 
+impl From<bool> for UntypedVal {
+    #[allow(clippy::cast_lossless)]
+    fn from(value: bool) -> Self {
+        Self::from_bits64(value as _)
+    }
+}
+
 macro_rules! impl_from_signed_prim {
     ( $( $prim:ty as $base:ty ),* $(,)? ) => {
         $(
@@ -241,6 +273,12 @@ macro_rules! impl_from_signed_prim {
                     Self::from_bits64(u64::from(value as $base))
                 }
             }
+
+            impl From<::core::num::NonZero<$prim>> for UntypedVal {
+                fn from(value: ::core::num::NonZero<$prim>) -> Self {
+                    <_ as From<$prim>>::from(value.get())
+                }
+            }
         )*
     };
 }
diff --git a/crates/wasmi/src/engine/translator/func/instr_encoder.rs b/crates/wasmi/src/engine/translator/func/instr_encoder.rs
@@ -840,7 +840,7 @@ impl InstrEncoder {
         lhs: Reg,
         rhs: T,
     ) -> Result<bool, Error> {
-        if !rhs.eq_zero() {
+        if !rhs.is_zero() {
             // Case: `rhs` needs to be zero to apply this optimization.
             return Ok(false);
         }
@@ -885,7 +885,7 @@ impl InstrEncoder {
         lhs: Reg,
         rhs: T,
     ) -> Result<bool, Error> {
-        if !rhs.eq_zero() {
+        if !rhs.is_zero() {
             // Case: `rhs` needs to be zero to apply this optimization.
             return Ok(false);
         }
diff --git a/crates/wasmi/src/engine/translator/func/mod.rs b/crates/wasmi/src/engine/translator/func/mod.rs
@@ -1328,7 +1328,7 @@ impl FuncTranslator {
                     // Custom optimization was applied: return early
                     return Ok(());
                 }
-                if T::from(lhs).eq_zero() {
+                if T::from(lhs).is_zero() {
                     // Optimization: Shifting or rotating a zero value is a no-op.
                     self.stack.push_const(lhs);
                     return Ok(());
@@ -1362,18 +1362,17 @@ impl FuncTranslator {
     ///
     /// - `{i32, i64}.{div_u, div_s, rem_u, rem_s}`
     #[allow(clippy::too_many_arguments)]
-    fn translate_divrem<T, NonZeroT>(
+    fn translate_divrem<T>(
         &mut self,
         make_instr: fn(result: Reg, lhs: Reg, rhs: Reg) -> Instruction,
-        make_instr_imm16: fn(result: Reg, lhs: Reg, rhs: Const16<NonZeroT>) -> Instruction,
+        make_instr_imm16: fn(result: Reg, lhs: Reg, rhs: Const16<T::NonZero>) -> Instruction,
         make_instr_imm16_rev: fn(result: Reg, lhs: Const16<T>, rhs: Reg) -> Instruction,
         consteval: fn(T, T) -> Result<T, TrapCode>,
         make_instr_opt: fn(&mut Self, lhs: Reg, rhs: Reg) -> Result<bool, Error>,
         make_instr_reg_imm_opt: fn(&mut Self, lhs: Reg, rhs: T) -> Result<bool, Error>,
     ) -> Result<(), Error>
     where
         T: WasmInteger,
-        NonZeroT: Copy + TryFrom<T> + TryInto<Const16<NonZeroT>>,
     {
         bail_unreachable!(self);
         match self.stack.pop2() {
@@ -1385,7 +1384,7 @@ impl FuncTranslator {
                 self.push_binary_instr(lhs, rhs, make_instr)
             }
             (TypedProvider::Register(lhs), TypedProvider::Const(rhs)) => {
-                let Some(non_zero_rhs) = NonZeroT::try_from(T::from(rhs)).ok() else {
+                let Some(non_zero_rhs) = <T as WasmInteger>::non_zero(T::from(rhs)) else {
                     // Optimization: division by zero always traps
                     self.translate_trap(TrapCode::IntegerDivisionByZero)?;
                     return Ok(());
diff --git a/crates/wasmi/src/engine/translator/func/visit.rs b/crates/wasmi/src/engine/translator/func/visit.rs
@@ -33,7 +33,6 @@ use crate::{
     FuncRef,
     Mutability,
 };
-use core::num::{NonZeroI32, NonZeroI64, NonZeroU32, NonZeroU64};
 use wasmparser::VisitOperator;
 
 /// Used to swap operands of binary [`Instruction`] constructor.
@@ -2037,7 +2036,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i32_div_s(&mut self) -> Self::Output {
-        self.translate_divrem::<i32, NonZeroI32>(
+        self.translate_divrem::<i32>(
             Instruction::i32_div_s,
             Instruction::i32_div_s_imm16_rhs,
             Instruction::i32_div_s_imm16_lhs,
@@ -2055,7 +2054,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i32_div_u(&mut self) -> Self::Output {
-        self.translate_divrem::<u32, NonZeroU32>(
+        self.translate_divrem::<u32>(
             Instruction::i32_div_u,
             Instruction::i32_div_u_imm16_rhs,
             Instruction::i32_div_u_imm16_lhs,
@@ -2073,7 +2072,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i32_rem_s(&mut self) -> Self::Output {
-        self.translate_divrem::<i32, NonZeroI32>(
+        self.translate_divrem::<i32>(
             Instruction::i32_rem_s,
             Instruction::i32_rem_s_imm16_rhs,
             Instruction::i32_rem_s_imm16_lhs,
@@ -2091,7 +2090,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i32_rem_u(&mut self) -> Self::Output {
-        self.translate_divrem::<u32, NonZeroU32>(
+        self.translate_divrem::<u32>(
             Instruction::i32_rem_u,
             Instruction::i32_rem_u_imm16_rhs,
             Instruction::i32_rem_u_imm16_lhs,
@@ -2354,7 +2353,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i64_div_s(&mut self) -> Self::Output {
-        self.translate_divrem::<i64, NonZeroI64>(
+        self.translate_divrem::<i64>(
             Instruction::i64_div_s,
             Instruction::i64_div_s_imm16_rhs,
             Instruction::i64_div_s_imm16_lhs,
@@ -2372,7 +2371,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i64_div_u(&mut self) -> Self::Output {
-        self.translate_divrem::<u64, NonZeroU64>(
+        self.translate_divrem::<u64>(
             Instruction::i64_div_u,
             Instruction::i64_div_u_imm16_rhs,
             Instruction::i64_div_u_imm16_lhs,
@@ -2390,7 +2389,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i64_rem_s(&mut self) -> Self::Output {
-        self.translate_divrem::<i64, NonZeroI64>(
+        self.translate_divrem::<i64>(
             Instruction::i64_rem_s,
             Instruction::i64_rem_s_imm16_rhs,
             Instruction::i64_rem_s_imm16_lhs,
@@ -2408,7 +2407,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     }
 
     fn visit_i64_rem_u(&mut self) -> Self::Output {
-        self.translate_divrem::<u64, NonZeroU64>(
+        self.translate_divrem::<u64>(
             Instruction::i64_rem_u,
             Instruction::i64_rem_u_imm16_rhs,
             Instruction::i64_rem_u_imm16_lhs,
diff --git a/crates/wasmi/src/engine/translator/utils.rs b/crates/wasmi/src/engine/translator/utils.rs
@@ -1,10 +1,11 @@
 use crate::{
-    core::{FuelCostsProvider, Typed, TypedVal, ValType},
+    core::{FuelCostsProvider, Typed, TypedVal, UntypedVal, ValType},
     ir::{Const16, Instruction, Sign},
     Error,
     ExternRef,
     FuncRef,
 };
+use core::num::NonZero;
 
 macro_rules! impl_typed_for {
     ( $( $ty:ident ),* $(,)? ) => {
@@ -39,35 +40,45 @@ impl_typed_for! {
 ///
 /// This trait provides some utility methods useful for translation.
 pub trait WasmInteger:
-    Copy + Eq + From<TypedVal> + Into<TypedVal> + TryInto<Const16<Self>>
+    Copy
+    + Eq
+    + Typed
+    + From<TypedVal>
+    + Into<TypedVal>
+    + From<UntypedVal>
+    + Into<UntypedVal>
+    + TryInto<Const16<Self>>
 {
-    /// Returns `true` if `self` is equal to zero (0).
-    fn eq_zero(self) -> bool;
-}
+    /// The non-zero type of the [`WasmInteger`].
+    type NonZero: Copy + Into<Self> + TryInto<Const16<Self::NonZero>> + Into<UntypedVal>;
 
-impl WasmInteger for i32 {
-    fn eq_zero(self) -> bool {
-        self == 0
-    }
-}
+    /// Returns `self` as [`Self::NonZero`] if possible.
+    ///
+    /// Returns `None` if `self` is zero.
+    fn non_zero(self) -> Option<Self::NonZero>;
 
-impl WasmInteger for u32 {
-    fn eq_zero(self) -> bool {
-        self == 0
-    }
+    /// Returns `true` if `self` is equal to zero (0).
+    fn is_zero(self) -> bool;
 }
 
-impl WasmInteger for i64 {
-    fn eq_zero(self) -> bool {
-        self == 0
-    }
-}
+macro_rules! impl_wasm_integer {
+    ($($ty:ty),*) => {
+        $(
+            impl WasmInteger for $ty {
+                type NonZero = NonZero<Self>;
 
-impl WasmInteger for u64 {
-    fn eq_zero(self) -> bool {
-        self == 0
-    }
+                fn non_zero(self) -> Option<Self::NonZero> {
+                    Self::NonZero::new(self)
+                }
+
+                fn is_zero(self) -> bool {
+                    self == 0
+                }
+            }
+        )*
+    };
 }
+impl_wasm_integer!(i32, u32, i64, u64);
 
 /// A WebAssembly float. Either `f32` or `f64`.
 ///
