rewrite Store::{entering,exiting}_native_code_hook into Store::call_hook (#3313)

which provides additional detail on what state transition is being made

crates/wasmtime/src/func.rs

@@ -1,7 +1,7 @@
 use crate::store::{StoreData, StoreOpaque, Stored};
 use crate::{
-    AsContext, AsContextMut, Engine, Extern, FuncType, Instance, InterruptHandle, StoreContext,
+    AsContext, AsContextMut, CallHook, Engine, Extern, FuncType, Instance, InterruptHandle,
-    StoreContextMut, Trap, Val, ValType,
+    StoreContext, StoreContextMut, Trap, Val, ValType,
 };
 use anyhow::{bail, Context as _, Result};
 use std::cmp::max;
@@ -845,7 +845,7 @@ impl Func {
         values_vec: *mut u128,
         func: &dyn Fn(Caller<'_, T>, &[Val], &mut [Val]) -> Result<(), Trap>,
     ) -> Result<(), Trap> {
-        caller.store.0.entering_native_hook()?;
+        caller.store.0.call_hook(CallHook::CallingHost)?;
 
         // Translate the raw JIT arguments in `values_vec` into a `Val` which
         // we'll be passing as a slice. The storage for our slice-of-`Val` we'll
@@ -895,7 +895,7 @@ impl Func {
         // hostcall to reuse our own storage.
         val_vec.truncate(0);
         caller.store.0.save_hostcall_val_storage(val_vec);
-        caller.store.0.exiting_native_hook()?;
+        caller.store.0.call_hook(CallHook::ReturningFromHost)?;
         Ok(())
     }
 
@@ -1044,7 +1044,7 @@ pub(crate) fn invoke_wasm_and_catch_traps<T>(
     unsafe {
         let exit = enter_wasm(store)?;
 
-        if let Err(trap) = store.0.exiting_native_hook() {
+        if let Err(trap) = store.0.call_hook(CallHook::CallingWasm) {
             exit_wasm(store, exit);
             return Err(trap);
         }
@@ -1055,7 +1055,7 @@ pub(crate) fn invoke_wasm_and_catch_traps<T>(
             closure,
         );
         exit_wasm(store, exit);
-        store.0.entering_native_hook()?;
+        store.0.call_hook(CallHook::ReturningFromWasm)?;
         result.map_err(Trap::from_runtime_box)
     }
 }
@@ -1741,7 +1741,7 @@ macro_rules! impl_into_func {
 
                         let ret = {
                             panic::catch_unwind(AssertUnwindSafe(|| {
-                                if let Err(trap) = caller.store.0.entering_native_hook() {
+                                if let Err(trap) = caller.store.0.call_hook(CallHook::CallingHost) {
                                     return R::fallible_from_trap(trap);
                                 }
                                 $(let $args = $args::from_abi($args, caller.store.0);)*
@@ -1749,7 +1749,7 @@ macro_rules! impl_into_func {
                                     caller.sub_caller(),
                                     $( $args, )*
                                 );
-                                if let Err(trap) = caller.store.0.exiting_native_hook() {
+                                if let Err(trap) = caller.store.0.call_hook(CallHook::ReturningFromHost) {
                                     return R::fallible_from_trap(trap);
                                 }
                                 r.into_fallible()

crates/wasmtime/src/lib.rs

@@ -401,7 +401,7 @@ pub use crate::memory::*;
 pub use crate::module::{FrameInfo, FrameSymbol, Module};
 pub use crate::r#ref::ExternRef;
 pub use crate::store::{
-    AsContext, AsContextMut, InterruptHandle, Store, StoreContext, StoreContextMut,
+    AsContext, AsContextMut, CallHook, InterruptHandle, Store, StoreContext, StoreContextMut,
 };
 pub use crate::trap::*;
 pub use crate::types::*;

crates/wasmtime/src/store.rs

@@ -157,6 +157,37 @@ pub struct Store<T> {
     inner: ManuallyDrop<Box<StoreInner<T>>>,
 }
 
+#[derive(Copy, Clone, Debug)]
+/// Passed to the argument of [`Store::call_hook`] to indicate a state transition in
+/// the WebAssembly VM.
+pub enum CallHook {
+    /// Indicates the VM is calling a WebAssembly function, from the host.
+    CallingWasm,
+    /// Indicates the VM is returning from a WebAssembly function, to the host.
+    ReturningFromWasm,
+    /// Indicates the VM is calling a host function, from WebAssembly.
+    CallingHost,
+    /// Indicates the VM is returning from a host function, to WebAssembly.
+    ReturningFromHost,
+}
+
+impl CallHook {
+    /// Indicates the VM is entering host code (exiting WebAssembly code)
+    pub fn entering_host(&self) -> bool {
+        match self {
+            CallHook::ReturningFromWasm | CallHook::CallingHost => true,
+            _ => false,
+        }
+    }
+    /// Indicates the VM is exiting host code (entering WebAssembly code)
+    pub fn exiting_host(&self) -> bool {
+        match self {
+            CallHook::ReturningFromHost | CallHook::CallingWasm => true,
+            _ => false,
+        }
+    }
+}
+
 /// Internal contents of a `Store<T>` that live on the heap.
 ///
 /// The members of this struct are those that need to be generic over `T`, the
@@ -167,8 +198,7 @@ pub struct StoreInner<T> {
     inner: StoreOpaque,
 
     limiter: Option<Box<dyn FnMut(&mut T) -> &mut (dyn crate::ResourceLimiter) + Send + Sync>>,
-    entering_native_hook: Option<Box<dyn FnMut(&mut T) -> Result<(), crate::Trap> + Send + Sync>>,
+    call_hook: Option<Box<dyn FnMut(&mut T, CallHook) -> Result<(), crate::Trap> + Send + Sync>>,
-    exiting_native_hook: Option<Box<dyn FnMut(&mut T) -> Result<(), crate::Trap> + Send + Sync>>,
     // for comments about `ManuallyDrop`, see `Store::into_data`
     data: ManuallyDrop<T>,
 }
@@ -340,8 +370,7 @@ impl<T> Store<T> {
                 hostcall_val_storage: Vec::new(),
             },
             limiter: None,
-            entering_native_hook: None,
+            call_hook: None,
-            exiting_native_hook: None,
             data: ManuallyDrop::new(data),
         });
 
@@ -433,56 +462,25 @@ impl<T> Store<T> {
         inner.limiter = Some(Box::new(limiter));
     }
 
-    /// Configure a function that runs each time the host resumes execution from
+    /// Configure a function that runs on calls and returns between WebAssembly
-    /// WebAssembly.
+    /// and host code.
     ///
-    /// This hook is called in two circumstances:
+    /// The function is passed a [`CallHook`] argument, which indicates which
-    ///
+    /// state transition the VM is making.
-    /// * When WebAssembly calls a function defined by the host, this hook is
-    ///   called before other host code runs.
-    /// * When WebAssembly returns back to the host after being called, this
-    ///   hook is called.
-    ///
-    /// This method can be used with [`Store::exiting_native_code_hook`] to track
-    /// execution time of WebAssembly, for example, by starting/stopping timers
-    /// in the enter/exit hooks.
     ///
     /// This function may return a [`Trap`]. If a trap is returned when an
     /// import was called, it is immediately raised as-if the host import had
     /// returned the trap. If a trap is returned after wasm returns to the host
     /// then the wasm function's result is ignored and this trap is returned
     /// instead.
-    pub fn entering_native_code_hook(
+    ///
+    /// After this function returns a trap, it may be called for subsequent returns
+    /// to host or wasm code as the trap propogates to the root call.
+    pub fn call_hook(
         &mut self,
-        hook: impl FnMut(&mut T) -> Result<(), Trap> + Send + Sync + 'static,
+        hook: impl FnMut(&mut T, CallHook) -> Result<(), Trap> + Send + Sync + 'static,
     ) {
-        self.inner.entering_native_hook = Some(Box::new(hook));
+        self.inner.call_hook = Some(Box::new(hook));
-    }
-
-    /// Configure a function that runs just before WebAssembly code starts
-    /// executing.
-    ///
-    /// The closure provided is called in two circumstances:
-    ///
-    /// * When the host calls a WebAssembly function, the hook is called just
-    ///   before WebAssembly starts executing.
-    /// * When a host function returns back to WebAssembly this hook is called
-    ///   just before the return.
-    ///
-    /// This method can be used with [`Store::entering_native_code_hook`] to track
-    /// execution time of WebAssembly, for example, by starting/stopping timers
-    /// in the enter/exit hooks.
-    ///
-    /// This function may return a [`Trap`]. If a trap is returned when an
-    /// imported host function is returning, then the imported host function's
-    /// result is ignored and the trap is raised. If a trap is returned when
-    /// the host is about to start executing WebAssembly, then no WebAssembly
-    /// code is run and the trap is returned instead.
-    pub fn exiting_native_code_hook(
-        &mut self,
-        hook: impl FnMut(&mut T) -> Result<(), Trap> + Send + Sync + 'static,
-    ) {
-        self.inner.exiting_native_hook = Some(Box::new(hook));
     }
 
     /// Returns the [`Engine`] that this store is associated with.
@@ -783,17 +781,9 @@ impl<T> StoreInner<T> {
         Some(accessor(&mut self.data))
     }
 
-    pub fn entering_native_hook(&mut self) -> Result<(), Trap> {
+    pub fn call_hook(&mut self, s: CallHook) -> Result<(), Trap> {
-        if let Some(hook) = &mut self.entering_native_hook {
+        if let Some(hook) = &mut self.call_hook {
-            hook(&mut self.data)
+            hook(&mut self.data, s)
-        } else {
-            Ok(())
-        }
-    }
-
-    pub fn exiting_native_hook(&mut self) -> Result<(), Trap> {
-        if let Some(hook) = &mut self.exiting_native_hook {
-            hook(&mut self.data)
         } else {
             Ok(())
         }

tests/all/call_hook.rs

@@ -0,0 +1,583 @@
+use anyhow::Error;
+use wasmtime::*;
+
+// Crate a synchronous Func, call it directly:
+#[test]
+fn call_wrapped_func() -> Result<(), Error> {
+    let mut store = Store::<State>::default();
+    store.call_hook(State::call_hook);
+    let f = Func::wrap(
+        &mut store,
+        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
+            // Calling this func will switch context into wasm, then back to host:
+            assert_eq!(caller.data().context, vec![Context::Wasm, Context::Host]);
+
+            assert_eq!(
+                caller.data().calls_into_host,
+                caller.data().returns_from_host + 1
+            );
+            assert_eq!(
+                caller.data().calls_into_wasm,
+                caller.data().returns_from_wasm + 1
+            );
+
+            assert_eq!(a, 1);
+            assert_eq!(b, 2);
+            assert_eq!(c, 3.0);
+            assert_eq!(d, 4.0);
+        },
+    );
+
+    f.call(
+        &mut store,
+        &[Val::I32(1), Val::I64(2), 3.0f32.into(), 4.0f64.into()],
+    )?;
+
+    // One switch from vm to host to call f, another in return from f.
+    assert_eq!(store.data().calls_into_host, 1);
+    assert_eq!(store.data().returns_from_host, 1);
+    assert_eq!(store.data().calls_into_wasm, 1);
+    assert_eq!(store.data().returns_from_wasm, 1);
+
+    f.typed::<(i32, i64, f32, f64), (), _>(&store)?
+        .call(&mut store, (1, 2, 3.0, 4.0))?;
+
+    assert_eq!(store.data().calls_into_host, 2);
+    assert_eq!(store.data().returns_from_host, 2);
+    assert_eq!(store.data().calls_into_wasm, 2);
+    assert_eq!(store.data().returns_from_wasm, 2);
+
+    Ok(())
+}
+
+// Create an async Func, call it directly:
+#[tokio::test]
+async fn call_wrapped_async_func() -> Result<(), Error> {
+    let mut config = Config::new();
+    config.async_support(true);
+    let engine = Engine::new(&config)?;
+    let mut store = Store::new(&engine, State::default());
+    store.call_hook(State::call_hook);
+    let f = Func::wrap4_async(
+        &mut store,
+        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
+            Box::new(async move {
+                // Calling this func will switch context into wasm, then back to host:
+                assert_eq!(caller.data().context, vec![Context::Wasm, Context::Host]);
+
+                assert_eq!(
+                    caller.data().calls_into_host,
+                    caller.data().returns_from_host + 1
+                );
+                assert_eq!(
+                    caller.data().calls_into_wasm,
+                    caller.data().returns_from_wasm + 1
+                );
+
+                assert_eq!(a, 1);
+                assert_eq!(b, 2);
+                assert_eq!(c, 3.0);
+                assert_eq!(d, 4.0);
+            })
+        },
+    );
+
+    f.call_async(
+        &mut store,
+        &[Val::I32(1), Val::I64(2), 3.0f32.into(), 4.0f64.into()],
+    )
+    .await?;
+
+    // One switch from vm to host to call f, another in return from f.
+    assert_eq!(store.data().calls_into_host, 1);
+    assert_eq!(store.data().returns_from_host, 1);
+    assert_eq!(store.data().calls_into_wasm, 1);
+    assert_eq!(store.data().returns_from_wasm, 1);
+
+    f.typed::<(i32, i64, f32, f64), (), _>(&store)?
+        .call_async(&mut store, (1, 2, 3.0, 4.0))
+        .await?;
+
+    assert_eq!(store.data().calls_into_host, 2);
+    assert_eq!(store.data().returns_from_host, 2);
+    assert_eq!(store.data().calls_into_wasm, 2);
+    assert_eq!(store.data().returns_from_wasm, 2);
+
+    Ok(())
+}
+
+// Use the Linker to define a sync func, call it through WebAssembly:
+#[test]
+fn call_linked_func() -> Result<(), Error> {
+    let engine = Engine::default();
+    let mut store = Store::new(&engine, State::default());
+    store.call_hook(State::call_hook);
+    let mut linker = Linker::new(&engine);
+
+    linker.func_wrap(
+        "host",
+        "f",
+        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
+            // Calling this func will switch context into wasm, then back to host:
+            assert_eq!(caller.data().context, vec![Context::Wasm, Context::Host]);
+
+            assert_eq!(
+                caller.data().calls_into_host,
+                caller.data().returns_from_host + 1
+            );
+            assert_eq!(
+                caller.data().calls_into_wasm,
+                caller.data().returns_from_wasm + 1
+            );
+
+            assert_eq!(a, 1);
+            assert_eq!(b, 2);
+            assert_eq!(c, 3.0);
+            assert_eq!(d, 4.0);
+        },
+    )?;
+
+    let wat = r#"
+        (module
+            (import "host" "f"
+                (func $f (param i32) (param i64) (param f32) (param f64)))
+            (func (export "export")
+                (call $f (i32.const 1) (i64.const 2) (f32.const 3.0) (f64.const 4.0)))
+        )
+    "#;
+    let module = Module::new(&engine, wat)?;
+
+    let inst = linker.instantiate(&mut store, &module)?;
+    let export = inst
+        .get_export(&mut store, "export")
+        .expect("get export")
+        .into_func()
+        .expect("export is func");
+
+    export.call(&mut store, &[])?;
+
+    // One switch from vm to host to call f, another in return from f.
+    assert_eq!(store.data().calls_into_host, 1);
+    assert_eq!(store.data().returns_from_host, 1);
+    assert_eq!(store.data().calls_into_wasm, 1);
+    assert_eq!(store.data().returns_from_wasm, 1);
+
+    export.typed::<(), (), _>(&store)?.call(&mut store, ())?;
+
+    assert_eq!(store.data().calls_into_host, 2);
+    assert_eq!(store.data().returns_from_host, 2);
+    assert_eq!(store.data().calls_into_wasm, 2);
+    assert_eq!(store.data().returns_from_wasm, 2);
+
+    Ok(())
+}
+
+// Use the Linker to define an async func, call it through WebAssembly:
+#[tokio::test]
+async fn call_linked_func_async() -> Result<(), Error> {
+    let mut config = Config::new();
+    config.async_support(true);
+    let engine = Engine::new(&config)?;
+    let mut store = Store::new(&engine, State::default());
+    store.call_hook(State::call_hook);
+
+    let f = Func::wrap4_async(
+        &mut store,
+        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
+            Box::new(async move {
+                // Calling this func will switch context into wasm, then back to host:
+                assert_eq!(caller.data().context, vec![Context::Wasm, Context::Host]);
+
+                assert_eq!(
+                    caller.data().calls_into_host,
+                    caller.data().returns_from_host + 1
+                );
+                assert_eq!(
+                    caller.data().calls_into_wasm,
+                    caller.data().returns_from_wasm + 1
+                );
+                assert_eq!(a, 1);
+                assert_eq!(b, 2);
+                assert_eq!(c, 3.0);
+                assert_eq!(d, 4.0);
+            })
+        },
+    );
+
+    let mut linker = Linker::new(&engine);
+
+    linker.define("host", "f", f)?;
+
+    let wat = r#"
+        (module
+            (import "host" "f"
+                (func $f (param i32) (param i64) (param f32) (param f64)))
+            (func (export "export")
+                (call $f (i32.const 1) (i64.const 2) (f32.const 3.0) (f64.const 4.0)))
+        )
+    "#;
+    let module = Module::new(&engine, wat)?;
+
+    let inst = linker.instantiate_async(&mut store, &module).await?;
+    let export = inst
+        .get_export(&mut store, "export")
+        .expect("get export")
+        .into_func()
+        .expect("export is func");
+
+    export.call_async(&mut store, &[]).await?;
+
+    // One switch from vm to host to call f, another in return from f.
+    assert_eq!(store.data().calls_into_host, 1);
+    assert_eq!(store.data().returns_from_host, 1);
+    assert_eq!(store.data().calls_into_wasm, 1);
+    assert_eq!(store.data().returns_from_wasm, 1);
+
+    export
+        .typed::<(), (), _>(&store)?
+        .call_async(&mut store, ())
+        .await?;
+
+    assert_eq!(store.data().calls_into_host, 2);
+    assert_eq!(store.data().returns_from_host, 2);
+    assert_eq!(store.data().calls_into_wasm, 2);
+    assert_eq!(store.data().returns_from_wasm, 2);
+
+    Ok(())
+}
+
+#[test]
+fn instantiate() -> Result<(), Error> {
+    let mut store = Store::<State>::default();
+    store.call_hook(State::call_hook);
+
+    let m = Module::new(store.engine(), "(module)")?;
+    Instance::new(&mut store, &m, &[])?;
+    assert_eq!(store.data().calls_into_wasm, 0);
+    assert_eq!(store.data().calls_into_host, 0);
+
+    let m = Module::new(store.engine(), "(module (func) (start 0))")?;
+    Instance::new(&mut store, &m, &[])?;
+    assert_eq!(store.data().calls_into_wasm, 1);
+    assert_eq!(store.data().calls_into_host, 0);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn instantiate_async() -> Result<(), Error> {
+    let mut config = Config::new();
+    config.async_support(true);
+    let engine = Engine::new(&config)?;
+    let mut store = Store::new(&engine, State::default());
+    store.call_hook(State::call_hook);
+
+    let m = Module::new(store.engine(), "(module)")?;
+    Instance::new_async(&mut store, &m, &[]).await?;
+    assert_eq!(store.data().calls_into_wasm, 0);
+    assert_eq!(store.data().calls_into_host, 0);
+
+    let m = Module::new(store.engine(), "(module (func) (start 0))")?;
+    Instance::new_async(&mut store, &m, &[]).await?;
+    assert_eq!(store.data().calls_into_wasm, 1);
+    assert_eq!(store.data().calls_into_host, 0);
+
+    Ok(())
+}
+
+#[test]
+fn recursion() -> Result<(), Error> {
+    // Make sure call hook behaves reasonably when called recursively
+
+    let engine = Engine::default();
+    let mut store = Store::new(&engine, State::default());
+    store.call_hook(State::call_hook);
+    let mut linker = Linker::new(&engine);
+
+    linker.func_wrap("host", "f", |mut caller: Caller<State>, n: i32| {
+        assert_eq!(caller.data().context.last(), Some(&Context::Host));
+
+        assert_eq!(caller.data().calls_into_host, caller.data().calls_into_wasm);
+
+        // Recurse
+        if n > 0 {
+            caller
+                .get_export("export")
+                .expect("caller exports \"export\"")
+                .into_func()
+                .expect("export is a func")
+                .typed::<i32, (), _>(&caller)
+                .expect("export typing")
+                .call(&mut caller, n - 1)
+                .unwrap()
+        }
+    })?;
+
+    let wat = r#"
+        (module
+            (import "host" "f"
+                (func $f (param i32)))
+            (func (export "export") (param i32)
+                (call $f (local.get 0)))
+        )
+    "#;
+    let module = Module::new(&engine, wat)?;
+
+    let inst = linker.instantiate(&mut store, &module)?;
+    let export = inst
+        .get_export(&mut store, "export")
+        .expect("get export")
+        .into_func()
+        .expect("export is func");
+
+    // Recursion depth:
+    let n: usize = 10;
+
+    export.call(&mut store, &[Val::I32(n as i32)])?;
+
+    // Recurse down to 0: n+1 calls
+    assert_eq!(store.data().calls_into_host, n + 1);
+    assert_eq!(store.data().returns_from_host, n + 1);
+    assert_eq!(store.data().calls_into_wasm, n + 1);
+    assert_eq!(store.data().returns_from_wasm, n + 1);
+
+    export
+        .typed::<i32, (), _>(&store)?
+        .call(&mut store, n as i32)?;
+
+    assert_eq!(store.data().calls_into_host, 2 * (n + 1));
+    assert_eq!(store.data().returns_from_host, 2 * (n + 1));
+    assert_eq!(store.data().calls_into_wasm, 2 * (n + 1));
+    assert_eq!(store.data().returns_from_wasm, 2 * (n + 1));
+
+    Ok(())
+}
+
+#[test]
+fn trapping() -> Result<(), Error> {
+    const TRAP_IN_F: i32 = 0;
+    const TRAP_NEXT_CALL_HOST: i32 = 1;
+    const TRAP_NEXT_RETURN_HOST: i32 = 2;
+    const TRAP_NEXT_CALL_WASM: i32 = 3;
+    const TRAP_NEXT_RETURN_WASM: i32 = 4;
+
+    let engine = Engine::default();
+
+    let mut linker = Linker::new(&engine);
+
+    linker.func_wrap(
+        "host",
+        "f",
+        |mut caller: Caller<State>, action: i32, recur: i32| -> Result<(), Trap> {
+            assert_eq!(caller.data().context.last(), Some(&Context::Host));
+            assert_eq!(caller.data().calls_into_host, caller.data().calls_into_wasm);
+
+            match action {
+                TRAP_IN_F => return Err(Trap::new("trapping in f")),
+                TRAP_NEXT_CALL_HOST => caller.data_mut().trap_next_call_host = true,
+                TRAP_NEXT_RETURN_HOST => caller.data_mut().trap_next_return_host = true,
+                TRAP_NEXT_CALL_WASM => caller.data_mut().trap_next_call_wasm = true,
+                TRAP_NEXT_RETURN_WASM => caller.data_mut().trap_next_return_wasm = true,
+                _ => {} // Do nothing
+            }
+
+            // recur so that we can trigger a next call.
+            // propogate its trap, if it traps!
+            if recur > 0 {
+                let _ = caller
+                    .get_export("export")
+                    .expect("caller exports \"export\"")
+                    .into_func()
+                    .expect("export is a func")
+                    .typed::<(i32, i32), (), _>(&caller)
+                    .expect("export typing")
+                    .call(&mut caller, (action, 0))?;
+            }
+
+            Ok(())
+        },
+    )?;
+
+    let wat = r#"
+        (module
+            (import "host" "f"
+                (func $f (param i32) (param i32)))
+            (func (export "export") (param i32) (param i32)
+                (call $f (local.get 0) (local.get 1)))
+        )
+    "#;
+    let module = Module::new(&engine, wat)?;
+
+    let run = |action: i32, recur: bool| -> (State, Option<Error>) {
+        let mut store = Store::new(&engine, State::default());
+        store.call_hook(State::call_hook);
+        let inst = linker
+            .instantiate(&mut store, &module)
+            .expect("instantiate");
+        let export = inst
+            .get_export(&mut store, "export")
+            .expect("get export")
+            .into_func()
+            .expect("export is func");
+
+        let r = export.call(
+            &mut store,
+            &[Val::I32(action), Val::I32(if recur { 1 } else { 0 })],
+        );
+        (store.into_data(), r.err())
+    };
+
+    let (s, e) = run(TRAP_IN_F, false);
+    assert!(e.unwrap().to_string().starts_with("trapping in f"));
+    assert_eq!(s.calls_into_host, 1);
+    assert_eq!(s.returns_from_host, 1);
+    assert_eq!(s.calls_into_wasm, 1);
+    assert_eq!(s.returns_from_wasm, 1);
+
+    // trap in next call to host. No calls after the bit is set, so this trap shouldn't happen
+    let (s, e) = run(TRAP_NEXT_CALL_HOST, false);
+    assert!(e.is_none());
+    assert_eq!(s.calls_into_host, 1);
+    assert_eq!(s.returns_from_host, 1);
+    assert_eq!(s.calls_into_wasm, 1);
+    assert_eq!(s.returns_from_wasm, 1);
+
+    // trap in next call to host. recur, so the second call into host traps:
+    let (s, e) = run(TRAP_NEXT_CALL_HOST, true);
+    assert!(e
+        .unwrap()
+        .to_string()
+        .starts_with("call_hook: trapping on CallingHost"));
+    assert_eq!(s.calls_into_host, 2);
+    assert_eq!(s.returns_from_host, 1);
+    assert_eq!(s.calls_into_wasm, 2);
+    assert_eq!(s.returns_from_wasm, 2);
+
+    // trap in the return from host. should trap right away, without recursion
+    let (s, e) = run(TRAP_NEXT_RETURN_HOST, false);
+    assert!(e
+        .unwrap()
+        .to_string()
+        .starts_with("call_hook: trapping on ReturningFromHost"));
+    assert_eq!(s.calls_into_host, 1);
+    assert_eq!(s.returns_from_host, 1);
+    assert_eq!(s.calls_into_wasm, 1);
+    assert_eq!(s.returns_from_wasm, 1);
+
+    // trap in next call to wasm. No calls after the bit is set, so this trap shouldnt happen:
+    let (s, e) = run(TRAP_NEXT_CALL_WASM, false);
+    assert!(e.is_none());
+    assert_eq!(s.calls_into_host, 1);
+    assert_eq!(s.returns_from_host, 1);
+    assert_eq!(s.calls_into_wasm, 1);
+    assert_eq!(s.returns_from_wasm, 1);
+
+    // trap in next call to wasm. recur, so the second call into wasm traps:
+    let (s, e) = run(TRAP_NEXT_CALL_WASM, true);
+    assert!(e
+        .unwrap()
+        .to_string()
+        .starts_with("call_hook: trapping on CallingWasm"));
+    assert_eq!(s.calls_into_host, 1);
+    assert_eq!(s.returns_from_host, 1);
+    assert_eq!(s.calls_into_wasm, 2);
+    assert_eq!(s.returns_from_wasm, 1);
+
+    // trap in the return from wasm. should trap right away, without recursion
+    let (s, e) = run(TRAP_NEXT_RETURN_WASM, false);
+    assert!(e
+        .unwrap()
+        .to_string()
+        .starts_with("call_hook: trapping on ReturningFromWasm"));
+    assert_eq!(s.calls_into_host, 1);
+    assert_eq!(s.returns_from_host, 1);
+    assert_eq!(s.calls_into_wasm, 1);
+    assert_eq!(s.returns_from_wasm, 1);
+
+    Ok(())
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum Context {
+    Host,
+    Wasm,
+}
+
+struct State {
+    context: Vec<Context>,
+
+    calls_into_host: usize,
+    returns_from_host: usize,
+    calls_into_wasm: usize,
+    returns_from_wasm: usize,
+
+    trap_next_call_host: bool,
+    trap_next_return_host: bool,
+    trap_next_call_wasm: bool,
+    trap_next_return_wasm: bool,
+}
+
+impl Default for State {
+    fn default() -> Self {
+        State {
+            context: Vec::new(),
+            calls_into_host: 0,
+            returns_from_host: 0,
+            calls_into_wasm: 0,
+            returns_from_wasm: 0,
+            trap_next_call_host: false,
+            trap_next_return_host: false,
+            trap_next_call_wasm: false,
+            trap_next_return_wasm: false,
+        }
+    }
+}
+
+impl State {
+    // This implementation asserts that hooks are always called in a stack-like manner.
+    fn call_hook(&mut self, s: CallHook) -> Result<(), Trap> {
+        match s {
+            CallHook::CallingHost => {
+                self.calls_into_host += 1;
+                if self.trap_next_call_host {
+                    return Err(Trap::new("call_hook: trapping on CallingHost"));
+                } else {
+                    self.context.push(Context::Host);
+                }
+            }
+            CallHook::ReturningFromHost => match self.context.pop() {
+                Some(Context::Host) => {
+                    self.returns_from_host += 1;
+                    if self.trap_next_return_host {
+                        return Err(Trap::new("call_hook: trapping on ReturningFromHost"));
+                    }
+                }
+                c => panic!(
+                    "illegal context: expected Some(Host), got {:?}. remaining: {:?}",
+                    c, self.context
+                ),
+            },
+            CallHook::CallingWasm => {
+                self.calls_into_wasm += 1;
+                if self.trap_next_call_wasm {
+                    return Err(Trap::new("call_hook: trapping on CallingWasm"));
+                } else {
+                    self.context.push(Context::Wasm);
+                }
+            }
+            CallHook::ReturningFromWasm => match self.context.pop() {
+                Some(Context::Wasm) => {
+                    self.returns_from_wasm += 1;
+                    if self.trap_next_return_wasm {
+                        return Err(Trap::new("call_hook: trapping on ReturningFromWasm"));
+                    }
+                }
+                c => panic!(
+                    "illegal context: expected Some(Wasm), got {:?}. remaining: {:?}",
+                    c, self.context
+                ),
+            },
+        }
+        Ok(())
+    }
+}

tests/all/main.rs

@@ -1,4 +1,5 @@
 mod async_functions;
+mod call_hook;
 mod cli_tests;
 mod custom_signal_handler;
 mod debug;
@@ -23,7 +24,6 @@ mod module;
 mod module_linking;
 mod module_serialize;
 mod name;
-mod native_hooks;
 mod pooling_allocator;
 mod relocs;
 mod stack_overflow;

tests/all/native_hooks.rs

@@ -1,281 +0,0 @@
-use anyhow::Error;
-use wasmtime::*;
-
-// Crate a synchronous Func, call it directly:
-#[test]
-fn call_wrapped_func() -> Result<(), Error> {
-    let mut store = Store::<State>::default();
-    store.entering_native_code_hook(State::entering_native);
-    store.exiting_native_code_hook(State::exiting_native);
-    let f = Func::wrap(
-        &mut store,
-        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
-            assert_eq!(
-                caller.data().switches_into_native % 2,
-                1,
-                "odd number of switches into native while in a Func"
-            );
-            assert_eq!(a, 1);
-            assert_eq!(b, 2);
-            assert_eq!(c, 3.0);
-            assert_eq!(d, 4.0);
-        },
-    );
-
-    f.call(
-        &mut store,
-        &[Val::I32(1), Val::I64(2), 3.0f32.into(), 4.0f64.into()],
-    )?;
-
-    // One switch from vm to native to call f, another in return from f.
-    assert_eq!(store.data().switches_into_native, 2);
-
-    f.typed::<(i32, i64, f32, f64), (), _>(&store)?
-        .call(&mut store, (1, 2, 3.0, 4.0))?;
-
-    assert_eq!(store.data().switches_into_native, 4);
-
-    Ok(())
-}
-
-// Create an async Func, call it directly:
-#[tokio::test]
-async fn call_wrapped_async_func() -> Result<(), Error> {
-    let mut config = Config::new();
-    config.async_support(true);
-    let engine = Engine::new(&config)?;
-    let mut store = Store::new(&engine, State::default());
-    store.entering_native_code_hook(State::entering_native);
-    store.exiting_native_code_hook(State::exiting_native);
-    let f = Func::wrap4_async(
-        &mut store,
-        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
-            Box::new(async move {
-                assert_eq!(
-                    caller.data().switches_into_native % 2,
-                    1,
-                    "odd number of switches into native while in a Func"
-                );
-                assert_eq!(a, 1);
-                assert_eq!(b, 2);
-                assert_eq!(c, 3.0);
-                assert_eq!(d, 4.0);
-            })
-        },
-    );
-
-    f.call_async(
-        &mut store,
-        &[Val::I32(1), Val::I64(2), 3.0f32.into(), 4.0f64.into()],
-    )
-    .await?;
-
-    // One switch from vm to native to call f, another in return from f.
-    assert_eq!(store.data().switches_into_native, 2);
-
-    f.typed::<(i32, i64, f32, f64), (), _>(&store)?
-        .call_async(&mut store, (1, 2, 3.0, 4.0))
-        .await?;
-
-    assert_eq!(store.data().switches_into_native, 4);
-
-    Ok(())
-}
-
-// Use the Linker to define a sync func, call it through WebAssembly:
-#[test]
-fn call_linked_func() -> Result<(), Error> {
-    let engine = Engine::default();
-    let mut store = Store::new(&engine, State::default());
-    store.entering_native_code_hook(State::entering_native);
-    store.exiting_native_code_hook(State::exiting_native);
-    let mut linker = Linker::new(&engine);
-
-    linker.func_wrap(
-        "host",
-        "f",
-        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
-            assert_eq!(
-                caller.data().switches_into_native % 2,
-                1,
-                "odd number of switches into native while in a Func"
-            );
-            assert_eq!(a, 1);
-            assert_eq!(b, 2);
-            assert_eq!(c, 3.0);
-            assert_eq!(d, 4.0);
-        },
-    )?;
-
-    let wat = r#"
-        (module
-            (import "host" "f"
-                (func $f (param i32) (param i64) (param f32) (param f64)))
-            (func (export "export")
-                (call $f (i32.const 1) (i64.const 2) (f32.const 3.0) (f64.const 4.0)))
-        )
-    "#;
-    let module = Module::new(&engine, wat)?;
-
-    let inst = linker.instantiate(&mut store, &module)?;
-    let export = inst
-        .get_export(&mut store, "export")
-        .expect("get export")
-        .into_func()
-        .expect("export is func");
-
-    export.call(&mut store, &[])?;
-
-    // One switch from vm to native to call f, another in return from f.
-    assert_eq!(store.data().switches_into_native, 2);
-
-    export.typed::<(), (), _>(&store)?.call(&mut store, ())?;
-
-    assert_eq!(store.data().switches_into_native, 4);
-
-    Ok(())
-}
-
-// Use the Linker to define an async func, call it through WebAssembly:
-#[tokio::test]
-async fn call_linked_func_async() -> Result<(), Error> {
-    let mut config = Config::new();
-    config.async_support(true);
-    let engine = Engine::new(&config)?;
-    let mut store = Store::new(&engine, State::default());
-    store.entering_native_code_hook(State::entering_native);
-    store.exiting_native_code_hook(State::exiting_native);
-
-    let f = Func::wrap4_async(
-        &mut store,
-        |caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
-            Box::new(async move {
-                assert_eq!(
-                    caller.data().switches_into_native % 2,
-                    1,
-                    "odd number of switches into native while in a Func"
-                );
-                assert_eq!(a, 1);
-                assert_eq!(b, 2);
-                assert_eq!(c, 3.0);
-                assert_eq!(d, 4.0);
-            })
-        },
-    );
-
-    let mut linker = Linker::new(&engine);
-
-    linker.define("host", "f", f)?;
-
-    let wat = r#"
-        (module
-            (import "host" "f"
-                (func $f (param i32) (param i64) (param f32) (param f64)))
-            (func (export "export")
-                (call $f (i32.const 1) (i64.const 2) (f32.const 3.0) (f64.const 4.0)))
-        )
-    "#;
-    let module = Module::new(&engine, wat)?;
-
-    let inst = linker.instantiate_async(&mut store, &module).await?;
-    let export = inst
-        .get_export(&mut store, "export")
-        .expect("get export")
-        .into_func()
-        .expect("export is func");
-
-    export.call_async(&mut store, &[]).await?;
-
-    // One switch from vm to native to call f, another in return from export.
-    assert_eq!(store.data().switches_into_native, 2);
-
-    export
-        .typed::<(), (), _>(&store)?
-        .call_async(&mut store, ())
-        .await?;
-
-    // 2 more switches.
-    assert_eq!(store.data().switches_into_native, 4);
-
-    Ok(())
-}
-
-#[test]
-fn instantiate() -> Result<(), Error> {
-    let mut store = Store::<State>::default();
-    store.entering_native_code_hook(State::entering_native);
-    store.exiting_native_code_hook(State::exiting_native);
-
-    let m = Module::new(store.engine(), "(module)")?;
-    Instance::new(&mut store, &m, &[])?;
-    assert_eq!(store.data().switches_into_native, 0);
-
-    let m = Module::new(store.engine(), "(module (func) (start 0))")?;
-    Instance::new(&mut store, &m, &[])?;
-    assert_eq!(store.data().switches_into_native, 1);
-
-    Ok(())
-}
-
-#[tokio::test]
-async fn instantiate_async() -> Result<(), Error> {
-    let mut config = Config::new();
-    config.async_support(true);
-    let engine = Engine::new(&config)?;
-    let mut store = Store::new(&engine, State::default());
-    store.entering_native_code_hook(State::entering_native);
-    store.exiting_native_code_hook(State::exiting_native);
-
-    let m = Module::new(store.engine(), "(module)")?;
-    Instance::new_async(&mut store, &m, &[]).await?;
-    assert_eq!(store.data().switches_into_native, 0);
-
-    let m = Module::new(store.engine(), "(module (func) (start 0))")?;
-    Instance::new_async(&mut store, &m, &[]).await?;
-    assert_eq!(store.data().switches_into_native, 1);
-
-    Ok(())
-}
-
-enum Context {
-    Native,
-    Vm,
-}
-
-struct State {
-    context: Context,
-    switches_into_native: usize,
-}
-
-impl Default for State {
-    fn default() -> Self {
-        State {
-            context: Context::Native,
-            switches_into_native: 0,
-        }
-    }
-}
-
-impl State {
-    fn entering_native(&mut self) -> Result<(), Trap> {
-        match self.context {
-            Context::Vm => {
-                println!("entering native");
-                self.context = Context::Native;
-                self.switches_into_native += 1;
-                Ok(())
-            }
-            Context::Native => Err(Trap::new("illegal state: exiting vm when in native")),
-        }
-    }
-    fn exiting_native(&mut self) -> Result<(), Trap> {
-        match self.context {
-            Context::Native => {
-                println!("entering vm");
-                self.context = Context::Vm;
-                Ok(())
-            }
-            Context::Vm => Err(Trap::new("illegal state: exiting native when in vm")),
-        }
-    }
-}