factors: Merge InitContext and ModuleInitContext

Signed-off-by: Lann Martin <[email protected]>

Author: lann

crates/factor-wasi/src/preview1.rs

@@ -1,5 +1,5 @@
 use spin_factors::{
-    Factor, FactorInstancePreparer, ModuleInitContext, PrepareContext, Result, SpinFactors,
+    Factor, FactorInstancePreparer, InitContext, PrepareContext, Result, SpinFactors,
 };
 use wasmtime_wasi::{preview1::WasiP1Ctx, WasiCtxBuilder};
 
@@ -9,11 +9,8 @@ impl Factor for WasiPreview1Factor {
     type InstancePreparer = InstancePreparer;
     type InstanceState = WasiP1Ctx;
 
-    fn module_init<Factors: SpinFactors>(
-        &mut self,
-        mut ctx: ModuleInitContext<Factors, Self>,
-    ) -> Result<()> {
-        ctx.link_bindings(wasmtime_wasi::preview1::add_to_linker_async)
+    fn init<Factors: SpinFactors>(&mut self, mut ctx: InitContext<Factors, Self>) -> Result<()> {
+        ctx.link_module_bindings(wasmtime_wasi::preview1::add_to_linker_async)
     }
 }
 

crates/factors-derive/src/lib.rs

@@ -63,30 +63,15 @@ fn expand_factors(input: &DeriveInput) -> syn::Result<TokenStream> {
         impl #name {
             pub fn init(
                 &mut self,
-                linker: &mut #wasmtime::component::Linker<#state_name>
+                mut linker: Option<&mut #wasmtime::component::Linker<#state_name>>,
+                mut module_linker: Option<&mut #wasmtime::Linker<#state_name>>,
             ) -> #Result<()> {
                 #(
                     #Factor::init::<Self>(
                         &mut self.#factor_names,
                         #factors_path::InitContext::<Self, #factor_types>::new(
-                            linker,
-                            |state| &mut state.#factor_names,
-                        )
-                    )?;
-                )*
-                Ok(())
-            }
-
-            #[allow(dead_code)]
-            pub fn module_init(
-                &mut self,
-                linker: &mut #wasmtime::Linker<#state_name>
-            ) -> #Result<()> {
-                #(
-                    #Factor::module_init::<Self>(
-                        &mut self.#factor_names,
-                        #factors_path::ModuleInitContext::<Self, #factor_types>::new(
-                            linker,
+                            linker.as_deref_mut(),
+                            module_linker.as_deref_mut(),
                             |state| &mut state.#factor_names,
                         )
                     )?;

crates/factors/src/lib.rs

@@ -21,14 +21,6 @@ pub trait Factor: Any + Sized {
         Ok(())
     }
 
-    fn module_init<Factors: SpinFactors>(
-        &mut self,
-        mut ctx: ModuleInitContext<Factors, Self>,
-    ) -> Result<()> {
-        _ = &mut ctx;
-        Ok(())
-    }
-
     fn validate_app(&self, app: &App) -> Result<()> {
         _ = app;
         Ok(())
@@ -38,24 +30,32 @@ pub trait Factor: Any + Sized {
 type GetDataFn<Factors, Fact> =
     fn(&mut <Factors as SpinFactors>::InstanceState) -> &mut <Fact as Factor>::InstanceState;
 
-pub struct FactorInitContext<'a, Factors: SpinFactors, Fact: Factor, Linker> {
-    linker: &'a mut Linker,
+pub struct InitContext<'a, Factors: SpinFactors, Fact: Factor> {
+    linker: Option<&'a mut Linker<Factors>>,
+    module_linker: Option<&'a mut ModuleLinker<Factors>>,
     get_data: GetDataFn<Factors, Fact>,
 }
 
-pub type InitContext<'a, Factors, Fact> = FactorInitContext<'a, Factors, Fact, Linker<Factors>>;
-
-pub type ModuleInitContext<'a, Factors, Fact> =
-    FactorInitContext<'a, Factors, Fact, ModuleLinker<Factors>>;
-
-impl<'a, Factors: SpinFactors, Fact: Factor, Linker> FactorInitContext<'a, Factors, Fact, Linker> {
+impl<'a, Factors: SpinFactors, Fact: Factor> InitContext<'a, Factors, Fact> {
     #[doc(hidden)]
-    pub fn new(linker: &'a mut Linker, get_data: GetDataFn<Factors, Fact>) -> Self {
-        Self { linker, get_data }
+    pub fn new(
+        linker: Option<&'a mut Linker<Factors>>,
+        module_linker: Option<&'a mut ModuleLinker<Factors>>,
+        get_data: GetDataFn<Factors, Fact>,
+    ) -> Self {
+        Self {
+            linker,
+            module_linker,
+            get_data,
+        }
+    }
+
+    pub fn linker(&mut self) -> Option<&mut Linker<Factors>> {
+        self.linker.as_deref_mut()
     }
 
-    pub fn linker(&mut self) -> &mut Linker {
-        self.linker
+    pub fn module_linker(&mut self) -> Option<&mut ModuleLinker<Factors>> {
+        self.module_linker.as_deref_mut()
     }
 
     pub fn get_data_fn(&self) -> GetDataFn<Factors, Fact> {
@@ -65,12 +65,31 @@ impl<'a, Factors: SpinFactors, Fact: Factor, Linker> FactorInitContext<'a, Facto
     pub fn link_bindings(
         &mut self,
         add_to_linker: impl Fn(
-            &mut Linker,
+            &mut Linker<Factors>,
             fn(&mut Factors::InstanceState) -> &mut Fact::InstanceState,
         ) -> Result<()>,
     ) -> Result<()>
 where {
-        add_to_linker(self.linker, self.get_data)
+        if let Some(linker) = self.linker.as_deref_mut() {
+            add_to_linker(linker, self.get_data)
+        } else {
+            Ok(())
+        }
+    }
+
+    pub fn link_module_bindings(
+        &mut self,
+        add_to_linker: impl Fn(
+            &mut ModuleLinker<Factors>,
+            fn(&mut Factors::InstanceState) -> &mut Fact::InstanceState,
+        ) -> Result<()>,
+    ) -> Result<()>
+where {
+        if let Some(linker) = self.module_linker.as_deref_mut() {
+            add_to_linker(linker, self.get_data)
+        } else {
+            Ok(())
+        }
     }
 }
 

crates/factors/tests/smoke.rs

@@ -1,25 +1,35 @@
-use spin_factor_wasi::WasiFactor;
+use spin_factor_wasi::{preview1::WasiPreview1Factor, WasiFactor};
 use spin_factors::SpinFactors;
 
 #[derive(SpinFactors)]
 struct Factors {
     wasi: WasiFactor,
+    wasip1: WasiPreview1Factor,
 }
 
 fn main() -> anyhow::Result<()> {
     let engine = wasmtime::Engine::default();
     let mut linker = wasmtime::component::Linker::new(&engine);
+    let mut module_linker = wasmtime::Linker::new(&engine);
 
     let mut factors = Factors {
         wasi: WasiFactor,
+        wasip1: WasiPreview1Factor,
         // outbound_networking_factor: OutboundNetworkingFactor,
         // outbound_http_factor: OutboundHttpFactor,
     };
-    factors.init(&mut linker).unwrap();
+    factors
+        .init(Some(&mut linker), Some(&mut module_linker))
+        .unwrap();
     let data = factors.build_store_data().unwrap();
 
     let mut store = wasmtime::Store::new(&engine, data);
+
     let component = wasmtime::component::Component::new(&engine, b"(component)").unwrap();
     let _instance = linker.instantiate(&mut store, &component).unwrap();
+
+    let module = wasmtime::Module::new(&engine, b"(module)").unwrap();
+    let _module_instance = module_linker.instantiate(&mut store, &module).unwrap();
+
     Ok(())
 }