Recreate the Store from scratch on wasmi every time

lib/src/executor/vm/interpreter.rs

@@ -49,33 +49,34 @@ impl Module {
 
 /// See [`super::VirtualMachinePrototype`].
 pub struct InterpreterPrototype {
+    /// Base components that can be used to recreate a prototype later if desired.
+    base_components: BaseComponents,
+
     // TODO: doc
     store: wasmi::Store<()>,
 
-    /// Original module.
-    module: Arc<wasmi::Module>,
-
-    /// Linker used to create instances.
-    linker: wasmi::Linker<()>,
-
     /// An instance of the module.
     instance: wasmi::Instance,
 
     /// Memory of the module instantiation.
     memory: wasmi::Memory,
 }
 
+struct BaseComponents {
+    module: Arc<wasmi::Module>,
+
+    /// For each import of the module, either `None` if not a function, or `Some` containing the
+    /// `usize` of that function.
+    resolved_imports: Vec<Option<usize>>,
+}
+
 impl InterpreterPrototype {
     /// See [`super::VirtualMachinePrototype::new`].
     pub fn new(
         module: &Module,
         mut symbols: impl FnMut(&str, &str, &Signature) -> Result<usize, ()>,
     ) -> Result<Self, NewErr> {
-        let mut store = wasmi::Store::new(module.inner.engine(), ());
-
-        let mut linker = wasmi::Linker::new();
-        let mut import_memory = None;
-
+        let mut resolved_imports = Vec::with_capacity(module.inner.imports().len());
         for import in module.inner.imports() {
             match import.ty() {
                 wasmi::ExternType::Func(func_type) => {
@@ -97,6 +98,36 @@ impl InterpreterPrototype {
                             }
                         };
 
+                    resolved_imports.push(Some(function_index));
+                }
+                wasmi::ExternType::Memory(_) => resolved_imports.push(None),
+                wasmi::ExternType::Global(_) | wasmi::ExternType::Table(_) => {
+                    return Err(NewErr::ImportTypeNotSupported)
+                }
+            }
+        }
+
+        Self::from_base_components(BaseComponents {
+            module: module.inner.clone(),
+            resolved_imports,
+        })
+    }
+
+    fn from_base_components(base_components: BaseComponents) -> Result<Self, NewErr> {
+        let mut store = wasmi::Store::new(base_components.module.engine(), ());
+
+        let mut linker = wasmi::Linker::<()>::new();
+        let mut import_memory = None;
+
+        for (module_import, resolved_function) in base_components
+            .module
+            .imports()
+            .zip(base_components.resolved_imports.iter())
+        {
+            match module_import.ty() {
+                wasmi::ExternType::Func(func_type) => {
+                    let function_index = resolved_function.unwrap();
+
                     let func = wasmi::Func::new(
                         &mut store,
                         func_type.clone(),
@@ -113,10 +144,12 @@ impl InterpreterPrototype {
 
                     // `define` returns an error in case of duplicate definition. Since we
                     // enumerate over the imports, this can't happen.
-                    linker.define(import.module(), import.name(), func).unwrap();
+                    linker
+                        .define(module_import.module(), module_import.name(), func)
+                        .unwrap();
                 }
                 wasmi::ExternType::Memory(memory_type) => {
-                    if import.module() != "env" || import.name() != "memory" {
+                    if module_import.module() != "env" || module_import.name() != "memory" {
                         return Err(NewErr::MemoryNotNamedMemory);
                     }
 
@@ -131,25 +164,17 @@ impl InterpreterPrototype {
                     // `define` returns an error in case of duplicate definition. Since we
                     // enumerate over the imports, this can't happen.
                     linker
-                        .define(import.module(), import.name(), memory)
+                        .define(module_import.module(), module_import.name(), memory)
                         .unwrap();
                 }
                 wasmi::ExternType::Global(_) | wasmi::ExternType::Table(_) => {
-                    return Err(NewErr::ImportTypeNotSupported)
+                    unreachable!()
                 }
             }
         }
 
-        Self::from_components(module.inner.clone(), store, linker)
-    }
-
-    fn from_components(
-        module: Arc<wasmi::Module>,
-        mut store: wasmi::Store<()>,
-        linker: wasmi::Linker<()>,
-    ) -> Result<Self, NewErr> {
         let instance = linker
-            .instantiate(&mut store, &*module)
+            .instantiate(&mut store, &*base_components.module)
             .map_err(|err| NewErr::Other(err.to_string()))?
             .ensure_no_start(&mut store)
             .map_err(|_| NewErr::StartFunctionNotSupported)?;
@@ -169,10 +194,9 @@ impl InterpreterPrototype {
             };
 
         Ok(InterpreterPrototype {
+            base_components,
             store,
             instance,
-            linker,
-            module,
             memory,
         })
     }
@@ -222,7 +246,8 @@ pub struct Prepare {
 impl Prepare {
     /// See [`super::Prepare::into_prototype`].
     pub fn into_prototype(self) -> InterpreterPrototype {
-        self.inner
+        // Since creation has succeeded in the past, there is no reason for it to fail now.
+        InterpreterPrototype::from_base_components(self.inner.base_components).unwrap()
     }
 
     /// See [`super::Prepare::memory_size`].
@@ -353,10 +378,9 @@ impl Prepare {
         };
 
         Ok(Interpreter {
+            base_components: self.inner.base_components,
             store: self.inner.store,
-            module: self.inner.module,
             memory: self.inner.memory,
-            linker: self.inner.linker,
             dummy_output_value: dummy_output_value,
             execution: Some(Execution::NotStarted(
                 func_to_call,
@@ -393,18 +417,15 @@ impl wasmi::core::HostError for InterruptedTrap {}
 
 /// See [`super::VirtualMachine`].
 pub struct Interpreter {
+    /// Base components that can be used to recreate a prototype later if desired.
+    base_components: BaseComponents,
+
     // TODO: doc
     store: wasmi::Store<()>,
 
-    /// Original module.
-    module: Arc<wasmi::Module>,
-
     /// Memory of the module instantiation.
     memory: wasmi::Memory,
 
-    /// Linker used to create instances.
-    linker: wasmi::Linker<()>,
-
     /// Execution context of this virtual machine. This notably holds the program counter, state
     /// of the stack, and so on.
     ///
@@ -571,11 +592,8 @@ impl Interpreter {
 
     /// See [`super::VirtualMachine::into_prototype`].
     pub fn into_prototype(self) -> InterpreterPrototype {
-        // TODO: zero the memory
-
-        // Because we have successfully instantiated the module in the past, there's no reason
-        // why instantiating again could fail now and not before.
-        InterpreterPrototype::from_components(self.module, self.store, self.linker).unwrap()
+        // Since creation has succeeded in the past, there is no reason for it to fail now.
+        InterpreterPrototype::from_base_components(self.base_components).unwrap()
     }
 }