Reuse func builder allocations (#411)

* make FunctionBuilder own an Engine

Fortunately cloning an Engine is very efficient.
This commit allows us to get rid of a lifetime on FunctionBuilder.

* make it possible to reuse FunctionBuilder allocations

* move LocalsRegistry into FunctionBuilderAllocations

wasmi_v1/src/engine/func_builder/control_stack.rs

@@ -8,6 +8,11 @@ pub struct ControlFlowStack {
 }
 
 impl ControlFlowStack {
+    /// Resets the [`ControlFlowStack`] to allow for reuse.
+    pub fn reset(&mut self) {
+        self.frames.clear()
+    }
+
     /// Returns `true` if `relative_depth` points to the first control flow frame.
     pub fn is_root(&self, relative_depth: u32) -> bool {
         debug_assert!(!self.is_empty());

wasmi_v1/src/engine/func_builder/inst_builder.rs

@@ -113,6 +113,12 @@ pub struct InstructionsBuilder {
 }
 
 impl InstructionsBuilder {
+    /// Resets the [`InstructionsBuilder`] to allow for reuse.
+    pub fn reset(&mut self) {
+        self.insts.clear();
+        self.labels.clear();
+    }
+
     /// Returns the current instruction pointer as index.
     pub fn current_pc(&self) -> InstructionIdx {
         InstructionIdx::from_usize(self.insts.len())

wasmi_v1/src/engine/func_builder/locals_registry.rs

@@ -68,6 +68,12 @@ impl LocalGroup {
 }
 
 impl LocalsRegistry {
+    /// Resets the [`LocalsRegistry`] to allow for reuse.
+    pub fn reset(&mut self) {
+        self.groups.clear();
+        self.max_index = 0;
+    }
+
     /// Returns the number of registered local variables.
     ///
     /// # Note

wasmi_v1/src/engine/func_builder/mod.rs

@@ -37,17 +37,35 @@ use crate::{
     Mutability,
 };
 use alloc::vec::Vec;
+use core::ops::{Deref, DerefMut};
 use wasmi_core::{Value, ValueType, F32, F64};
 
 /// The interface to translate a `wasmi` bytecode function using Wasm bytecode.
 #[derive(Debug)]
-pub struct FunctionBuilder<'engine, 'parser> {
+pub struct FunctionBuilder<'alloc, 'parser> {
     /// The [`Engine`] for which the function is translated.
-    engine: &'engine Engine,
+    engine: Engine,
     /// The function under construction.
     func: FuncIdx,
     /// The immutable `wasmi` module resources.
     res: ModuleResources<'parser>,
+    /// This represents the reachability of the currently translated code.
+    ///
+    /// - `true`: The currently translated code is reachable.
+    /// - `false`: The currently translated code is unreachable and can be skipped.
+    ///
+    /// # Note
+    ///
+    /// Visiting the Wasm `Else` or `End` control flow operator resets
+    /// reachability to `true` again.
+    reachable: bool,
+    /// The reusable data structures of the [`FunctionBuilder`].
+    allocations: &'alloc mut FunctionBuilderAllocations,
+}
+
+/// Reusable allocations of a [`FunctionBuilder`].
+#[derive(Debug, Default)]
+pub struct FunctionBuilderAllocations {
     /// The control flow frame stack that represents the Wasm control flow.
     control_frames: ControlFlowStack,
     /// The emulated value stack.
@@ -60,36 +78,55 @@ pub struct FunctionBuilder<'engine, 'parser> {
     inst_builder: InstructionsBuilder,
     /// Stores and resolves local variable types.
     locals: LocalsRegistry,
-    /// This represents the reachability of the currently translated code.
-    ///
-    /// - `true`: The currently translated code is reachable.
-    /// - `false`: The currently translated code is unreachable and can be skipped.
+}
+
+impl<'alloc, 'parser> Deref for FunctionBuilder<'alloc, 'parser> {
+    type Target = FunctionBuilderAllocations;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        self.allocations
+    }
+}
+
+impl<'alloc, 'parser> DerefMut for FunctionBuilder<'alloc, 'parser> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.allocations
+    }
+}
+
+impl FunctionBuilderAllocations {
+    /// Resets the data structures of the [`FunctionBuilderAllocations`].
     ///
     /// # Note
     ///
-    /// Visiting the Wasm `Else` or `End` control flow operator resets
-    /// reachability to `true` again.
-    reachable: bool,
+    /// This must be called before reusing this [`FunctionBuilderAllocations`]
+    /// by another [`FunctionBuilder`].
+    fn reset(&mut self) {
+        self.control_frames.reset();
+        self.value_stack.reset();
+        self.inst_builder.reset();
+        self.locals.reset();
+    }
 }
 
-impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
+impl<'alloc, 'parser> FunctionBuilder<'alloc, 'parser> {
     /// Creates a new [`FunctionBuilder`].
-    pub fn new(engine: &'engine Engine, func: FuncIdx, res: ModuleResources<'parser>) -> Self {
-        let mut inst_builder = InstructionsBuilder::default();
-        let mut control_frames = ControlFlowStack::default();
-        Self::register_func_body_block(func, res, &mut inst_builder, &mut control_frames);
-        let value_stack = ValueStack::default();
-        let mut locals = LocalsRegistry::default();
-        Self::register_func_params(func, res, &mut locals);
+    pub fn new(
+        engine: &Engine,
+        func: FuncIdx,
+        res: ModuleResources<'parser>,
+        allocations: &'alloc mut FunctionBuilderAllocations,
+    ) -> Self {
+        Self::register_func_body_block(func, res, allocations);
+        Self::register_func_params(func, res, allocations);
         Self {
-            engine,
+            engine: engine.clone(),
             func,
             res,
-            control_frames,
-            value_stack,
-            inst_builder,
-            locals,
             reachable: true,
+            allocations,
         }
     }
 
@@ -119,29 +156,29 @@ impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
     fn register_func_body_block(
         func: FuncIdx,
         res: ModuleResources<'parser>,
-        inst_builder: &mut InstructionsBuilder,
-        control_frames: &mut ControlFlowStack,
+        allocations: &mut FunctionBuilderAllocations,
     ) {
+        allocations.reset();
         let func_type = res.get_type_of_func(func);
         let block_type = BlockType::func_type(func_type);
-        let end_label = inst_builder.new_label();
+        let end_label = allocations.inst_builder.new_label();
         let block_frame = BlockControlFrame::new(block_type, end_label, 0);
-        control_frames.push_frame(block_frame);
+        allocations.control_frames.push_frame(block_frame);
     }
 
     /// Registers the function parameters in the emulated value stack.
     fn register_func_params(
         func: FuncIdx,
         res: ModuleResources<'parser>,
-        locals: &mut LocalsRegistry,
+        allocations: &mut FunctionBuilderAllocations,
     ) -> usize {
         let dedup_func_type = res.get_type_of_func(func);
         let func_type = res
             .engine()
             .resolve_func_type(dedup_func_type, Clone::clone);
         let params = func_type.params();
         for param_type in params {
-            locals.register_locals(*param_type, 1);
+            allocations.locals.register_locals(*param_type, 1);
         }
         params.len()
     }
@@ -177,9 +214,9 @@ impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
     }
 
     /// Finishes constructing the function and returns its [`FuncBody`].
-    pub fn finish(mut self) -> FuncBody {
-        self.inst_builder.finish(
-            self.engine,
+    pub fn finish(self) -> FuncBody {
+        self.allocations.inst_builder.finish(
+            &self.engine,
             self.len_locals(),
             self.value_stack.max_stack_height() as usize,
         )
@@ -216,9 +253,9 @@ impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
         // Find out how many values we need to keep (copy to the new stack location after the drop).
         let keep = match frame.kind() {
             ControlFrameKind::Block | ControlFrameKind::If => {
-                frame.block_type().len_results(self.engine)
+                frame.block_type().len_results(&self.engine)
             }
-            ControlFrameKind::Loop => frame.block_type().len_params(self.engine),
+            ControlFrameKind::Loop => frame.block_type().len_params(&self.engine),
         };
         // Find out how many values we need to drop.
         let current_height = self.value_stack.len();
@@ -320,7 +357,7 @@ pub enum AquiredTarget {
     Return(DropKeep),
 }
 
-impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
+impl<'alloc, 'parser> FunctionBuilder<'alloc, 'parser> {
     /// Translates a Wasm `unreachable` instruction.
     pub fn translate_unreachable(&mut self) -> Result<(), ModuleError> {
         self.translate_if_reachable(|builder| {
@@ -344,7 +381,7 @@ impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
     /// When the emulated value stack underflows. This should not happen
     /// since we have already validated the input Wasm prior.
     fn frame_stack_height(&self, block_type: BlockType) -> u32 {
-        let len_params = block_type.len_params(self.engine);
+        let len_params = block_type.len_params(&self.engine);
         let stack_height = self.value_stack.len();
         stack_height.checked_sub(len_params).unwrap_or_else(|| {
             panic!(
@@ -456,8 +493,8 @@ impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
         // when entering the `if` in the first place so that the `else`
         // block has the same parameters on top of the stack.
         self.value_stack.shrink_to(if_frame.stack_height());
-        if_frame.block_type().foreach_param(self.engine, |param| {
-            self.value_stack.push(param);
+        if_frame.block_type().foreach_param(&self.engine, |param| {
+            self.allocations.value_stack.push(param);
         });
         self.control_frames.push_frame(if_frame);
         // We can reset reachability now since the parent `if` block was reachable.
@@ -467,19 +504,22 @@ impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
 
     /// Translates a Wasm `end` control flow operator.
     pub fn translate_end(&mut self) -> Result<(), ModuleError> {
-        let frame = self.control_frames.last();
+        let frame = self.allocations.control_frames.last();
         if let ControlFrame::If(if_frame) = &frame {
             // At this point we can resolve the `Else` label.
             //
             // Note: The `Else` label might have already been resolved
             //       in case there was an `Else` block.
-            self.inst_builder
+            self.allocations
+                .inst_builder
                 .resolve_label_if_unresolved(if_frame.else_label());
         }
         if frame.is_reachable() && !matches!(frame.kind(), ControlFrameKind::Loop) {
             // At this point we can resolve the `End` labels.
             // Note that `loop` control frames do not have an `End` label.
-            self.inst_builder.resolve_label(frame.end_label());
+            self.allocations
+                .inst_builder
+                .resolve_label(frame.end_label());
         }
         // These bindings are required because of borrowing issues.
         let frame_reachable = frame.is_reachable();
@@ -498,9 +538,9 @@ impl<'engine, 'parser> FunctionBuilder<'engine, 'parser> {
             self.value_stack.shrink_to(frame_stack_height);
         }
         let frame = self.control_frames.pop_frame();
-        frame
-            .block_type()
-            .foreach_result(self.engine, |result| self.value_stack.push(result));
+        frame.block_type().foreach_result(&self.engine, |result| {
+            self.allocations.value_stack.push(result)
+        });
         Ok(())
     }
 

wasmi_v1/src/engine/func_builder/value_stack.rs

@@ -16,6 +16,12 @@ pub struct ValueStack {
 }
 
 impl ValueStack {
+    /// Resets the [`ValueStack`] to allow for reuse.
+    pub fn reset(&mut self) {
+        self.values.clear();
+        self.max_stack_height = 0;
+    }
+
     /// Returns the maximum value stack height.
     pub fn max_stack_height(&self) -> u32 {
         self.max_stack_height

wasmi_v1/src/engine/mod.rs

@@ -22,7 +22,14 @@ use self::{
 pub use self::{
     bytecode::{DropKeep, Target},
     code_map::FuncBody,
-    func_builder::{FunctionBuilder, InstructionIdx, LabelIdx, RelativeDepth, Reloc},
+    func_builder::{
+        FunctionBuilder,
+        FunctionBuilderAllocations,
+        InstructionIdx,
+        LabelIdx,
+        RelativeDepth,
+        Reloc,
+    },
     traits::{CallParams, CallResults},
 };
 use super::{func::FuncEntityInternal, AsContext, AsContextMut, Func};

wasmi_v1/src/module/compile/mod.rs

@@ -1,7 +1,7 @@
 pub use self::block_type::BlockType;
 use super::{utils::value_type_from_wasmparser, FuncIdx, ModuleResources};
 use crate::{
-    engine::{FuncBody, FunctionBuilder},
+    engine::{FuncBody, FunctionBuilder, FunctionBuilderAllocations},
     Engine,
     ModuleError,
 };
@@ -29,16 +29,17 @@ pub fn translate<'parser>(
     func_body: FunctionBody<'parser>,
     validator: FuncValidator<ValidatorResources>,
     res: ModuleResources<'parser>,
+    allocations: &mut FunctionBuilderAllocations,
 ) -> Result<FuncBody, ModuleError> {
-    FunctionTranslator::new(engine, func, func_body, validator, res).translate()
+    FunctionTranslator::new(engine, func, func_body, validator, res, allocations).translate()
 }
 
 /// Translates Wasm bytecode into `wasmi` bytecode for a single Wasm function.
-struct FunctionTranslator<'engine, 'parser> {
+struct FunctionTranslator<'alloc, 'parser> {
     /// The function body that shall be translated.
     func_body: FunctionBody<'parser>,
     /// The interface to incrementally build up the `wasmi` bytecode function.
-    func_builder: FunctionBuilder<'engine, 'parser>,
+    func_builder: FunctionBuilder<'alloc, 'parser>,
     /// The Wasm validator.
     validator: FuncValidator<ValidatorResources>,
     /// The `wasmi` module resources.
@@ -48,16 +49,17 @@ struct FunctionTranslator<'engine, 'parser> {
     res: ModuleResources<'parser>,
 }
 
-impl<'engine, 'parser> FunctionTranslator<'engine, 'parser> {
+impl<'alloc, 'parser> FunctionTranslator<'alloc, 'parser> {
     /// Creates a new Wasm to `wasmi` bytecode function translator.
     fn new(
-        engine: &'engine Engine,
+        engine: &Engine,
         func: FuncIdx,
         func_body: FunctionBody<'parser>,
         validator: FuncValidator<ValidatorResources>,
         res: ModuleResources<'parser>,
+        allocations: &'alloc mut FunctionBuilderAllocations,
     ) -> Self {
-        let func_builder = FunctionBuilder::new(engine, func, res);
+        let func_builder = FunctionBuilder::new(engine, func, res, allocations);
         Self {
             func_body,
             func_builder,

wasmi_v1/src/module/compile/operator.rs

@@ -6,7 +6,7 @@ use crate::{
 };
 use wasmparser::{Ieee32, Ieee64, TypeOrFuncType};
 
-impl<'engine, 'parser> FunctionTranslator<'engine, 'parser> {
+impl<'alloc, 'parser> FunctionTranslator<'alloc, 'parser> {
     /// Translate a Wasm `nop` (no operation) instruction.
     pub fn translate_nop(&mut self) -> Result<(), ModuleError> {
         // We can simply ignore Wasm `nop` instructions.
@@ -448,7 +448,7 @@ macro_rules! define_translate_fn {
     };
 }
 
-impl<'engine, 'parser> FunctionTranslator<'engine, 'parser> {
+impl<'alloc, 'parser> FunctionTranslator<'alloc, 'parser> {
     define_translate_fn! {
         /// Translate a Wasm `unreachable` instruction.
         fn translate_unreachable();