Generalize comparator submodule logic

crates/wasmi/src/engine/translator/func/comparator.rs → crates/wasmi/src/engine/translator/comparator.rs

@@ -1,9 +1,27 @@
 use crate::{
-    engine::translator::func::ValueStack,
+    core::UntypedVal,
     ir::{BranchOffset, BranchOffset16, Comparator, ComparatorAndOffset, Instruction, Reg},
     Error,
 };
 
+/// Types able to allocate function local constant values.
+///
+/// # Note
+///
+/// This allows to cheaply convert immediate values to [`Reg`]s.
+///
+/// # Errors
+///
+/// If the function local constant allocation from immediate value to [`Reg`] failed.
+pub trait AllocConst {
+    /// Allocates a new function local constant value and returns its [`Reg`].
+    ///
+    /// # Note
+    ///
+    /// Constant values allocated this way are deduplicated and return shared [`Reg`].
+    fn alloc_const<T: Into<UntypedVal>>(&mut self, value: T) -> Result<Reg, Error>;
+}
+
 pub trait NegateCmpInstr: Sized {
     /// Negates the compare (`cmp`) [`Instruction`].
     fn negate_cmp_instr(&self) -> Option<Self>;
@@ -346,15 +364,15 @@ pub trait TryIntoCmpBranchInstr: Sized {
     fn try_into_cmp_branch_instr(
         &self,
         offset: BranchOffset,
-        stack: &mut ValueStack,
+        stack: &mut impl AllocConst,
     ) -> Result<Option<Self>, Error>;
 }
 
 impl TryIntoCmpBranchInstr for Instruction {
     fn try_into_cmp_branch_instr(
         &self,
         offset: BranchOffset,
-        stack: &mut ValueStack,
+        stack: &mut impl AllocConst,
     ) -> Result<Option<Self>, Error> {
         use Instruction as I;
         let Ok(offset) = BranchOffset16::try_from(offset) else {
@@ -458,15 +476,15 @@ pub trait TryIntoCmpBranchFallbackInstr {
     fn try_into_cmp_branch_fallback_instr(
         &self,
         offset: BranchOffset,
-        stack: &mut ValueStack,
+        stack: &mut impl AllocConst,
     ) -> Result<Option<Instruction>, Error>;
 }
 
 impl TryIntoCmpBranchFallbackInstr for Instruction {
     fn try_into_cmp_branch_fallback_instr(
         &self,
         offset: BranchOffset,
-        stack: &mut ValueStack,
+        stack: &mut impl AllocConst,
     ) -> Result<Option<Instruction>, Error> {
         use Instruction as I;
         debug_assert!(BranchOffset16::try_from(offset).is_err());

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

@@ -1,16 +1,18 @@
 use crate::{
     core::{FuelCostsProvider, UntypedVal, ValType},
     engine::translator::{
-        func::{
-            stack::RegisterSpace,
-            utils::FromProviders as _,
-            LabelRef,
-            LabelRegistry,
+        comparator::{
             LogicalizeCmpInstr,
             NegateCmpInstr,
             TryIntoCmpBranchFallbackInstr,
             TryIntoCmpBranchInstr,
             TryIntoCmpSelectInstr,
+        },
+        func::{
+            stack::RegisterSpace,
+            utils::FromProviders as _,
+            LabelRef,
+            LabelRegistry,
             TypedProvider,
             ValueStack,
         },

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

@@ -1,6 +1,5 @@
 //! Function translation for the register-machine bytecode based Wasmi engine.
 
-mod comparator;
 mod control_frame;
 mod control_stack;
 mod instr_encoder;
@@ -13,14 +12,13 @@ mod visit;
 #[cfg(feature = "simd")]
 mod simd;
 
+pub use self::{
+    control_frame::ControlFrame,
+    control_stack::ControlStack,
+    instr_encoder::{Instr, InstrEncoder},
+    stack::TypedProvider,
+};
 use self::{
-    comparator::{
-        LogicalizeCmpInstr,
-        NegateCmpInstr,
-        TryIntoCmpBranchFallbackInstr,
-        TryIntoCmpBranchInstr,
-        TryIntoCmpSelectInstr,
-    },
     control_frame::{
         BlockControlFrame,
         BlockHeight,
@@ -34,12 +32,6 @@ use self::{
     stack::ValueStack,
     utils::FromProviders as _,
 };
-pub use self::{
-    control_frame::ControlFrame,
-    control_stack::ControlStack,
-    instr_encoder::{Instr, InstrEncoder},
-    stack::TypedProvider,
-};
 use crate::{
     core::{FuelCostsProvider, TrapCode, TypedVal, UntypedVal, ValType},
     engine::{

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

@@ -13,7 +13,10 @@ use super::{PreservedLocal, TypedVal};
 use crate::{
     core::UntypedVal,
     engine::{
-        translator::func::{Provider, UntypedProvider},
+        translator::{
+            comparator::AllocConst,
+            func::{Provider, UntypedProvider},
+        },
         TranslationError,
     },
     ir::{Reg, RegSpan},
@@ -420,3 +423,9 @@ impl ValueStack {
         self.reg_alloc.dec_register_usage(register)
     }
 }
+
+impl AllocConst for ValueStack {
+    fn alloc_const<T: Into<UntypedVal>>(&mut self, value: T) -> Result<Reg, Error> {
+        self.consts.alloc(value.into())
+    }
+}

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

@@ -1,5 +1,6 @@
 //! Function translation for the register-machine bytecode based Wasmi engine.
 
+mod comparator;
 mod driver;
 mod error;
 mod func;