Optimize wasmi::Linker host function setup

crates/wasmi/benches/benches.rs

@@ -16,6 +16,7 @@ use wasmi::{
     Engine,
     Extern,
     Func,
+    FuncType,
     Linker,
     Memory,
     Module,
@@ -63,6 +64,20 @@ criterion_group!(
         bench_overhead_call_untyped_0,
         bench_overhead_call_untyped_16,
 );
+criterion_group!(
+    name = bench_linker;
+    config = Criterion::default()
+        .sample_size(10)
+        .measurement_time(Duration::from_millis(1000))
+        .warm_up_time(Duration::from_millis(1000));
+    targets =
+        bench_linker_setup_same,
+        bench_linker_build_finish_same,
+        bench_linker_build_construct_same,
+        bench_linker_setup_unique,
+        bench_linker_build_finish_unique,
+        bench_linker_build_construct_unique,
+);
 criterion_group! {
     name = bench_execute;
     config = Criterion::default()
@@ -96,7 +111,8 @@ criterion_main!(
     bench_translate,
     bench_instantiate,
     bench_execute,
-    bench_overhead
+    bench_overhead,
+    bench_linker,
 );
 
 const WASM_KERNEL: &str =
@@ -244,6 +260,145 @@ fn bench_instantiate_wasm_kernel(c: &mut Criterion) {
     });
 }
 
+fn bench_linker_build_finish_same(c: &mut Criterion) {
+    let len_funcs = 50;
+    let bench_id = format!("linker/build/finish/same/{len_funcs}");
+    c.bench_function(&bench_id, |b| {
+        let func_names: Vec<String> = (0..len_funcs).map(|i| format!("{i}")).collect();
+        let mut builder = <Linker<()>>::build();
+        for func_name in &func_names {
+            builder.func_wrap("env", func_name, || ()).unwrap();
+        }
+        b.iter(|| {
+            let engine = Engine::default();
+            _ = builder.finish(&engine);
+        })
+    });
+}
+
+fn bench_linker_build_construct_same(c: &mut Criterion) {
+    let len_funcs = 50;
+    let bench_id = format!("linker/build/construct/same/{len_funcs}");
+    c.bench_function(&bench_id, |b| {
+        let func_names: Vec<String> = (0..len_funcs).map(|i| format!("{i}")).collect();
+        b.iter(|| {
+            let mut builder = <Linker<()>>::build();
+            for func_name in &func_names {
+                builder.func_wrap("env", func_name, || ()).unwrap();
+            }
+        })
+    });
+}
+
+fn bench_linker_setup_same(c: &mut Criterion) {
+    let len_funcs = 50;
+    let bench_id = format!("linker/setup/same/{len_funcs}");
+    c.bench_function(&bench_id, |b| {
+        let func_names: Vec<String> = (0..len_funcs).map(|i| format!("{i}")).collect();
+        b.iter(|| {
+            let engine = Engine::default();
+            let mut linker = <Linker<()>>::new(&engine);
+            for func_name in &func_names {
+                linker.func_wrap("env", func_name, || ()).unwrap();
+            }
+        })
+    });
+}
+
+/// Generates `count` host functions with different signatures.
+fn generate_unique_host_functions(count: usize) -> Vec<(String, FuncType)> {
+    let types = [
+        ValueType::I32,
+        ValueType::I64,
+        ValueType::F32,
+        ValueType::F64,
+        ValueType::FuncRef,
+        ValueType::ExternRef,
+    ];
+    (0..count)
+        .map(|i| {
+            let func_name = format!("{i}");
+            let (len_params, len_results) = if i % 2 == 0 {
+                ((i / (types.len() * 2)) + 1, 0)
+            } else {
+                (0, (i / (types.len() * 2)) + 1)
+            };
+            let chosen_type = types[i % 4];
+            let func_type = FuncType::new(
+                vec![chosen_type; len_params],
+                vec![chosen_type; len_results],
+            );
+            (func_name, func_type)
+        })
+        .collect()
+}
+
+fn bench_linker_setup_unique(c: &mut Criterion) {
+    let len_funcs = 50;
+    let bench_id = format!("linker/setup/unique/{len_funcs}");
+    c.bench_function(&bench_id, |b| {
+        let funcs = generate_unique_host_functions(len_funcs);
+        b.iter(|| {
+            let engine = Engine::default();
+            let mut linker = <Linker<()>>::new(&engine);
+            for (func_name, func_type) in &funcs {
+                linker
+                    .func_new(
+                        "env",
+                        func_name,
+                        func_type.clone(),
+                        move |_caller, _params, _results| Ok(()),
+                    )
+                    .unwrap();
+            }
+        })
+    });
+}
+
+fn bench_linker_build_finish_unique(c: &mut Criterion) {
+    let len_funcs = 50;
+    let bench_id = format!("linker/build/finish/unique/{len_funcs}");
+    c.bench_function(&bench_id, |b| {
+        let funcs = generate_unique_host_functions(len_funcs);
+        let mut builder = <Linker<()>>::build();
+        for (func_name, func_type) in &funcs {
+            builder
+                .func_new(
+                    "env",
+                    func_name,
+                    func_type.clone(),
+                    move |_caller, _params, _results| Ok(()),
+                )
+                .unwrap();
+        }
+        b.iter(|| {
+            let engine = Engine::default();
+            _ = builder.finish(&engine);
+        })
+    });
+}
+
+fn bench_linker_build_construct_unique(c: &mut Criterion) {
+    let len_funcs = 50;
+    let bench_id = format!("linker/build/construct/unique/{len_funcs}");
+    c.bench_function(&bench_id, |b| {
+        let funcs = generate_unique_host_functions(len_funcs);
+        b.iter(|| {
+            let mut builder = <Linker<()>>::build();
+            for (func_name, func_type) in &funcs {
+                builder
+                    .func_new(
+                        "env",
+                        func_name,
+                        func_type.clone(),
+                        move |_caller, _params, _results| Ok(()),
+                    )
+                    .unwrap();
+            }
+        })
+    });
+}
+
 #[allow(dead_code)]
 fn bench_instantiate_contract(c: &mut Criterion, name: &str, path: &str) {
     let bench_id = format!("instantiate/{name}");

crates/wasmi/src/func/mod.rs

@@ -22,7 +22,7 @@
     StoreContext,
     Stored,
 };
-use crate::{engine::ResumableCall, Engine, Error, Value};
+use crate::{engine::ResumableCall, Error, Value};
 use core::{fmt, fmt::Debug, num::NonZeroU32};
 use std::{boxed::Box, sync::Arc};
 use wasmi_arena::ArenaIndex;
@@ -175,15 +175,15 @@
 /// A host function instance.
 pub struct HostFuncTrampolineEntity<T> {
     /// The type of the associated host function.
-    ty: DedupFuncType,
+    ty: FuncType,
     /// The trampoline of the associated host function.
     trampoline: TrampolineEntity<T>,
 }
 
 impl<T> Clone for HostFuncTrampolineEntity<T> {
     fn clone(&self) -> Self {
         Self {
-            ty: self.ty,
+            ty: self.ty.clone(),
             trampoline: self.trampoline.clone(),
         }
     }
@@ -198,7 +198,7 @@
 impl<T> HostFuncTrampolineEntity<T> {
     /// Creates a new host function trampoline from the given dynamically typed closure.
     pub fn new(
-        engine: &Engine,
+        // engine: &Engine,
         ty: FuncType,
         func: impl Fn(Caller<'_, T>, &[Value], &mut [Value]) -> Result<(), Error>
             + Send
@@ -224,19 +224,19 @@
             func(caller, params, results)?;
             Ok(func_results.encode_results_from_slice(results).unwrap())
         });
-        let ty = engine.alloc_func_type(ty.clone());
+        // let ty = engine.alloc_func_type(ty.clone());
         Self { ty, trampoline }
     }
 
     /// Creates a new host function trampoline from the given statically typed closure.
-    pub fn wrap<Params, Results>(engine: &Engine, func: impl IntoFunc<T, Params, Results>) -> Self {
-        let (signature, trampoline) = func.into_func();
-        let ty = engine.alloc_func_type(signature);
+    pub fn wrap<Params, Results>(func: impl IntoFunc<T, Params, Results>) -> Self {
+        let (ty, trampoline) = func.into_func();
+        // let ty = engine.alloc_func_type(signature);
         Self { ty, trampoline }
     }
 
-    /// Returns the signature of the host function.
-    pub fn ty_dedup(&self) -> &DedupFuncType {
+    /// Returns the [`FuncType`] of the host function.
+    pub fn func_type(&self) -> &FuncType {
         &self.ty
     }
 
@@ -340,8 +340,8 @@
             + 'static,
     ) -> Self {
         let engine = ctx.as_context().store.engine();
-        let host_func = HostFuncTrampolineEntity::new(engine, ty, func);
-        let ty_dedup = *host_func.ty_dedup();
+        let host_func = HostFuncTrampolineEntity::new(ty, func);
+        let ty_dedup = engine.alloc_func_type(host_func.func_type().clone());
         let trampoline = host_func.trampoline().clone();
         let func = ctx.as_context_mut().store.alloc_trampoline(trampoline);
         ctx.as_context_mut()
@@ -356,8 +356,8 @@
         func: impl IntoFunc<T, Params, Results>,
     ) -> Self {
         let engine = ctx.as_context().store.engine();
-        let host_func = HostFuncTrampolineEntity::wrap(engine, func);
-        let ty_dedup = *host_func.ty_dedup();
+        let host_func = HostFuncTrampolineEntity::wrap(func);
+        let ty_dedup = engine.alloc_func_type(host_func.func_type().clone());
         let trampoline = host_func.trampoline().clone();
         let func = ctx.as_context_mut().store.alloc_trampoline(trampoline);
         ctx.as_context_mut()

crates/wasmi/src/lib.rs

@@ -151,7 +151,7 @@ pub use self::{
     global::{Global, GlobalType, Mutability},
     instance::{Export, ExportsIter, Extern, ExternType, Instance},
     limits::{ResourceLimiter, StoreLimits, StoreLimitsBuilder},
-    linker::Linker,
+    linker::{Linker, LinkerBuilder},
     memory::{Memory, MemoryType},
     module::{
         ExportType,