Allow to compile Wasm modules in host functions called from the Wasmi…

… executor (take 2) (#1122)

* add test to compile a Wasm module in a called host func

* refactor CodeMap

This makes it possible to compile Wasm modules from host functions called from Wasm.

* add #[inline] and #[cold] attributes where useful

* fix clippy warning

* further refactoring of CodeMap::get method

* add safety comment

* refactor and clean-up CodeMap

* simplify InternalFuncEntity::get_uncompiled

* add safety comment

* change inline annotation

* fix safety comment

* add note comment

* improve docs and comments

* fix doc link

* remove unused methods

* fix assert matches

* move CodeMap up in the source file

* remove unused From impls

* move UncompiledFuncEntity::new up

* move TypeIndex type up

* refactor UncompiledFuncEntity::new params

* improve (and kinda fix) doc comment

* add comment to explain how we chose MAX_INLINE_SIZE

* rename InternalFuncEntity -> FuncEntity

* rename CodeMap public methods

* apply rustfmt

* update and fix docs (and doc links)

crates/wasmi/src/engine/executor/instrs/call.rs

@@ -3,10 +3,9 @@ use crate::{
     core::TrapCode,
     engine::{
         bytecode::{FuncIdx, Instruction, Register, RegisterSpan, SignatureIdx, TableIdx},
-        code_map::InstructionPtr,
+        code_map::{CompiledFuncRef, InstructionPtr},
         executor::stack::{CallFrame, FrameParams, ValueStack},
         CompiledFunc,
-        CompiledFuncEntity,
         FuncParams,
     },
     func::{FuncEntity, HostFuncEntity},
@@ -203,7 +202,7 @@ impl<'engine> Executor<'engine> {
     fn dispatch_compiled_func<C: CallContext>(
         &mut self,
         results: RegisterSpan,
-        func: &CompiledFuncEntity,
+        func: CompiledFuncRef,
     ) -> Result<CallFrame, Error> {
         // We have to reinstantiate the `self.sp` [`FrameRegisters`] since we just called
         // [`ValueStack::alloc_call_frame`] which might invalidate all live [`FrameRegisters`].

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

@@ -49,9 +49,8 @@ impl EngineInner {
     where
         Results: CallResults,
     {
-        let res = self.code_map.read();
         let mut stack = self.stacks.lock().reuse_or_new();
-        let results = EngineExecutor::new(&res, &mut stack)
+        let results = EngineExecutor::new(&self.code_map, &mut stack)
             .execute_root_func(ctx.store, func, params, results)
             .map_err(|error| match error.into_resumable() {
                 Ok(error) => error.into_error(),
@@ -79,9 +78,8 @@ impl EngineInner {
         Results: CallResults,
     {
         let store = ctx.store;
-        let code_map = self.code_map.read();
         let mut stack = self.stacks.lock().reuse_or_new();
-        let results = EngineExecutor::new(&code_map, &mut stack)
+        let results = EngineExecutor::new(&self.code_map, &mut stack)
             .execute_root_func(store, func, params, results);
         match results {
             Ok(results) => {
@@ -127,10 +125,9 @@ impl EngineInner {
     where
         Results: CallResults,
     {
-        let code_map = self.code_map.read();
         let host_func = invocation.host_func();
         let caller_results = invocation.caller_results();
-        let results = EngineExecutor::new(&code_map, &mut invocation.stack).resume_func(
+        let results = EngineExecutor::new(&self.code_map, &mut invocation.stack).resume_func(
             ctx.store,
             host_func,
             params,

crates/wasmi/src/engine/executor/stack/values.rs

@@ -1,7 +1,7 @@
 use super::{err_stack_overflow, StackOffsets};
 use crate::{
     core::{TrapCode, UntypedVal},
-    engine::{bytecode::Register, CompiledFuncEntity},
+    engine::{bytecode::Register, code_map::CompiledFuncRef},
 };
 use core::{
     fmt::{self, Debug},
@@ -218,7 +218,7 @@ impl ValueStack {
     /// When trying to grow the [`ValueStack`] over its maximum size limit.
     pub fn alloc_call_frame(
         &mut self,
-        func: &CompiledFuncEntity,
+        func: CompiledFuncRef,
         on_resize: impl FnMut(&mut Self),
     ) -> Result<(FrameParams, StackOffsets), TrapCode> {
         let len_registers = func.len_registers();

crates/wasmi/src/engine/mod.rs

@@ -18,6 +18,9 @@ mod tests;
 #[cfg(test)]
 use self::bytecode::RegisterSpan;
 
+#[cfg(test)]
+use code_map::CompiledFuncRef;
+
 pub(crate) use self::{
     block_type::BlockType,
     config::FuelCosts,
@@ -475,7 +478,7 @@ pub struct EngineInner {
     /// The [`Config`] of the engine.
     config: Config,
     /// Stores information about all compiled functions.
-    code_map: RwLock<CodeMap>,
+    code_map: CodeMap,
     /// Deduplicated function types.
     ///
     /// # Note
@@ -600,7 +603,7 @@ impl EngineInner {
         let engine_idx = EngineIdx::new();
         Self {
             config: *config,
-            code_map: RwLock::new(CodeMap::new(config)),
+            code_map: CodeMap::new(config),
             func_types: RwLock::new(FuncTypeRegistry::new(engine_idx)),
             allocs: Mutex::new(ReusableAllocationStack::default()),
             stacks: Mutex::new(EngineStacks::new(config)),
@@ -634,7 +637,7 @@ impl EngineInner {
     ///
     /// Returns a [`CompiledFunc`] reference to allow accessing the allocated [`CompiledFunc`].
     fn alloc_func(&self) -> CompiledFunc {
-        self.code_map.write().alloc_func()
+        self.code_map.alloc_func()
     }
 
     /// Returns reusable [`FuncTranslatorAllocations`] from the [`Engine`].
@@ -699,7 +702,8 @@ impl EngineInner {
     /// - If `func` is an invalid [`CompiledFunc`] reference for this [`CodeMap`].
     /// - If `func` refers to an already initialized [`CompiledFunc`].
     fn init_func(&self, compiled_func: CompiledFunc, func_entity: CompiledFuncEntity) {
-        self.code_map.write().init_func(compiled_func, func_entity)
+        self.code_map
+            .init_func_as_compiled(compiled_func, func_entity)
     }
 
     /// Initializes the uninitialized [`CompiledFunc`] for the [`Engine`].
@@ -722,8 +726,7 @@ impl EngineInner {
         func_to_validate: Option<FuncToValidate<ValidatorResources>>,
     ) {
         self.code_map
-            .write()
-            .init_lazy_func(func, func_idx, bytes, module, func_to_validate)
+            .init_func_as_uncompiled(func, func_idx, bytes, module, func_to_validate)
     }
 
     /// Resolves the [`InternalFuncEntity`] for [`CompiledFunc`] and applies `f` to it.
@@ -732,12 +735,12 @@ impl EngineInner {
     ///
     /// If [`CompiledFunc`] is invalid for [`Engine`].
     #[cfg(test)]
-    pub(super) fn resolve_func<F, R>(&self, func: CompiledFunc, f: F) -> Result<R, Error>
+    pub(super) fn resolve_func<'a, F, R>(&'a self, func: CompiledFunc, f: F) -> Result<R, Error>
     where
-        F: FnOnce(&CompiledFuncEntity) -> R,
+        F: FnOnce(CompiledFuncRef<'a>) -> R,
     {
         // Note: We use `None` so this test-only function will never charge for compilation fuel.
-        Ok(f(self.code_map.read().get(None, func)?))
+        Ok(f(self.code_map.get(None, func)?))
     }
 
     /// Returns the [`Instruction`] of `func` at `index`.

crates/wasmi/tests/e2e/v1/host_call_compilation.rs

@@ -0,0 +1,43 @@
+//! This tests that a host function called from Wasm can compile Wasm modules and does not deadlock.
+
+use wasmi::{AsContextMut, Caller, Engine, Linker, Module, Store};
+
+/// Converts the given `.wat` into `.wasm`.
+fn wat2wasm(wat: &str) -> Result<Vec<u8>, wat::Error> {
+    wat::parse_str(wat)
+}
+
+fn compile_module(engine: &Engine) -> wasmi::Module {
+    let wasm = wat2wasm(include_str!("../wat/host_call_compilation.wat")).unwrap();
+    Module::new(engine, &wasm[..]).unwrap()
+}
+
+#[test]
+fn test_compile_in_host_call() {
+    let engine = Engine::default();
+    let mut store = <Store<()>>::new(&engine, ());
+    let module = compile_module(store.engine());
+    let mut linker = <Linker<()>>::new(&engine);
+    linker
+        .func_wrap(
+            "env",
+            "compile",
+            |mut caller: Caller<()>| -> Result<(), wasmi::Error> {
+                let store = caller.as_context_mut();
+                let engine = store.engine();
+                let _module = compile_module(engine);
+                Ok(())
+            },
+        )
+        .unwrap();
+    let instance = linker
+        .instantiate(&mut store, &module)
+        .unwrap()
+        .ensure_no_start(&mut store)
+        .unwrap();
+    instance
+        .get_typed_func::<(), ()>(&mut store, "run")
+        .unwrap()
+        .call(&mut store, ())
+        .unwrap();
+}

crates/wasmi/tests/e2e/v1/mod.rs

@@ -1,6 +1,7 @@
 mod fuel_consumption;
 mod fuel_metering;
 mod func;
+mod host_call_compilation;
 mod host_call_instantiation;
 mod host_calls_wasm;
 mod resource_limiter;

crates/wasmi/tests/e2e/wat/host_call_compilation.wat

@@ -0,0 +1,7 @@
+(module
+    (import "env" "compile" (func $compile))
+
+    (func (export "run")
+        (call $compile)
+    )
+)