Repository: bytecodealliance/wasmtime
Pull request #8791 (2024-06-13 00:44:49)

wasmtime: Make Store context available to call hooks
Rather than digging only the embedder's store-specific data out and passing that along, pass access to the entire Store. This is the same thing we do for epoch interruption hooks, including the annoyance of having to take the callback out of the store temporarily while calling it to avoid having multiple mutable borrows.

jameysharp: I think the only other example of this pattern in `Store` is [`Store::limiter`](https://docs.rs/wasmtime/latest/wasmtime/struct.Store.html#method.limiter) and its async variant. I don't want to change that today but I think it's worth discussing whether we should do that eventually.

Patch:
diff --git a/cranelift/isle/isle/src/codegen.rs b/cranelift/isle/isle/src/codegen.rs
index 76267e235..5230d96e0 100644
--- a/cranelift/isle/isle/src/codegen.rs
+++ b/cranelift/isle/isle/src/codegen.rs
@@ -639,8 +639,7 @@ impl<L: Length, C> Length for ContextIterWrapper<L, C> {{
                     )?;
                     write!(ctx.out, "{}", &ctx.indent)?;
                     match ret_kind {
-                        ReturnKind::Plain => write!(ctx.out, "return ")?,
-                        ReturnKind::Option => write!(ctx.out, "return Some(")?,
+                        ReturnKind::Plain | ReturnKind::Option => write!(ctx.out, "return ")?,
                         ReturnKind::Iterator => write!(ctx.out, "returns.extend(Some(")?,
                     }
                     self.emit_expr(ctx, result)?;
@@ -648,8 +647,7 @@ impl<L: Length, C> Length for ContextIterWrapper<L, C> {{
                         write!(ctx.out, ".clone()")?;
                     }
                     match ret_kind {
-                        ReturnKind::Plain => writeln!(ctx.out, ";")?,
-                        ReturnKind::Option => writeln!(ctx.out, ");")?,
+                        ReturnKind::Plain | ReturnKind::Option => writeln!(ctx.out, ";")?,
                         ReturnKind::Iterator => {
                             writeln!(ctx.out, "));")?;
                             writeln!(
@@ -753,6 +751,11 @@ impl<L: Length, C> Length for ContextIterWrapper<L, C> {{
                 Ok(())
             }
 
+            &Binding::MakeSome { inner } => {
+                write!(ctx.out, "Some(")?;
+                self.emit_expr(ctx, inner)?;
+                write!(ctx.out, ")")
+            }
             &Binding::MatchSome { source } => {
                 self.emit_expr(ctx, source)?;
                 write!(ctx.out, "?")
diff --git a/cranelift/isle/isle/src/sema.rs b/cranelift/isle/isle/src/sema.rs
index 91511d46d..4ba2b69fe 100644
--- a/cranelift/isle/isle/src/sema.rs
+++ b/cranelift/isle/isle/src/sema.rs
@@ -339,6 +339,17 @@ impl Term {
         matches!(self.kind, TermKind::EnumVariant { .. })
     }
 
+    /// Is this term partial?
+    pub fn is_partial(&self) -> bool {
+        matches!(
+            self.kind,
+            TermKind::Decl {
+                flags: TermFlags { partial: true, .. },
+                ..
+            }
+        )
+    }
+
     /// Does this term have a constructor?
     pub fn has_constructor(&self) -> bool {
         matches!(
diff --git a/cranelift/isle/isle/src/trie_again.rs b/cranelift/isle/isle/src/trie_again.rs
index b4fc37396..a6bb8255e 100644
--- a/cranelift/isle/isle/src/trie_again.rs
+++ b/cranelift/isle/isle/src/trie_again.rs
@@ -111,6 +111,11 @@ pub enum Binding {
         /// get the field names.
         field: TupleIndex,
     },
+    /// The result of constructing an Option::Some variant.
+    MakeSome {
+        /// Contained expression.
+        inner: BindingId,
+    },
     /// Pattern-match one of the previous bindings against `Option::Some` and produce a new binding
     /// from its contents. There must be a corresponding [Constraint::Some] for each `source` that
     /// appears in a `MatchSome` binding. (This currently only happens with external extractors.)
@@ -250,6 +255,7 @@ impl Binding {
             Binding::Iterator { source } => std::slice::from_ref(source),
             Binding::MakeVariant { fields, .. } => &fields[..],
             Binding::MatchVariant { source, .. } => std::slice::from_ref(source),
+            Binding::MakeSome { inner } => std::slice::from_ref(inner),
             Binding::MatchSome { source } => std::slice::from_ref(source),
             Binding::MatchTuple { source, .. } => std::slice::from_ref(source),
         }
@@ -384,6 +390,11 @@ impl RuleSetBuilder {
         self.current_rule.pos = rule.pos;
         self.current_rule.prio = rule.prio;
         self.current_rule.result = rule.visit(self, termenv);
+        if termenv.terms[rule.root_term.index()].is_partial() {
+            self.current_rule.result = self.dedup_binding(Binding::MakeSome {
+                inner: self.current_rule.result,
+            });
+        }
         self.normalize_equivalence_classes();
         let rule = std::mem::take(&mut self.current_rule);
 
diff --git a/crates/wasmtime/src/runtime/store.rs b/crates/wasmtime/src/runtime/store.rs
index e66f38325..83aa99f68 100644
--- a/crates/wasmtime/src/runtime/store.rs
+++ b/crates/wasmtime/src/runtime/store.rs
@@ -239,11 +239,11 @@ enum ResourceLimiterInner<T> {
 pub trait CallHookHandler<T>: Send {
     /// A callback to run when wasmtime is about to enter a host call, or when about to
     /// exit the hostcall.
-    async fn handle_call_event(&self, t: &mut T, ch: CallHook) -> Result<()>;
+    async fn handle_call_event(&self, t: StoreContextMut<'_, T>, ch: CallHook) -> Result<()>;
 }
 
 enum CallHookInner<T> {
-    Sync(Box<dyn FnMut(&mut T, CallHook) -> Result<()> + Send + Sync>),
+    Sync(Box<dyn FnMut(StoreContextMut<'_, T>, CallHook) -> Result<()> + Send + Sync>),
     #[cfg(feature = "async")]
     Async(Box<dyn CallHookHandler<T> + Send + Sync>),
 }
@@ -772,7 +772,7 @@ impl<T> Store<T> {
     /// to host or wasm code as the trap propagates to the root call.
     pub fn call_hook(
         &mut self,
-        hook: impl FnMut(&mut T, CallHook) -> Result<()> + Send + Sync + 'static,
+        hook: impl FnMut(StoreContextMut<'_, T>, CallHook) -> Result<()> + Send + Sync + 'static,
     ) {
         self.inner.call_hook = Some(CallHookInner::Sync(Box::new(hook)));
     }
@@ -1143,19 +1143,32 @@ impl<T> StoreInner<T> {
             }
         }
 
-        match &mut self.call_hook {
-            Some(CallHookInner::Sync(hook)) => hook(&mut self.data, s),
+        // Temporarily take the configured behavior to avoid mutably borrowing
+        // multiple times.
+        if let Some(mut call_hook) = self.call_hook.take() {
+            let result = self.invoke_call_hook(&mut call_hook, s);
+            self.call_hook = Some(call_hook);
+            result
+        } else {
+            Ok(())
+        }
+    }
+
+    fn invoke_call_hook(&mut self, call_hook: &mut CallHookInner<T>, s: CallHook) -> Result<()> {
+        match call_hook {
+            CallHookInner::Sync(hook) => hook((&mut *self).as_context_mut(), s),
 
             #[cfg(feature = "async")]
-            Some(CallHookInner::Async(handler)) => unsafe {
-                Ok(self
-                    .inner
+            CallHookInner::Async(handler) => unsafe {
+                self.inner
                     .async_cx()
                     .ok_or_else(|| anyhow!("couldn't grab async_cx for call hook"))?
-                    .block_on(handler.handle_call_event(&mut self.data, s).as_mut())??)
+                    .block_on(
+                        handler
+                            .handle_call_event((&mut *self).as_context_mut(), s)
+                            .as_mut(),
+                    )?
             },
-
-            None => Ok(()),
         }
     }
 }


Test patch:
diff --git a/tests/all/call_hook.rs b/tests/all/call_hook.rs
index 5b7bc60da..4be45fcc7 100644
--- a/tests/all/call_hook.rs
+++ b/tests/all/call_hook.rs
@@ -10,7 +10,7 @@ use wasmtime::*;
 #[test]
 fn call_wrapped_func() -> Result<(), Error> {
     let mut store = Store::<State>::default();
-    store.call_hook(State::call_hook);
+    store.call_hook(sync_call_hook);
 
     fn verify(state: &State) {
         // Calling this func will switch context into wasm, then back to host:
@@ -120,7 +120,7 @@ async fn call_wrapped_async_func() -> Result<(), Error> {
     config.async_support(true);
     let engine = Engine::new(&config)?;
     let mut store = Store::new(&engine, State::default());
-    store.call_hook(State::call_hook);
+    store.call_hook(sync_call_hook);
     let f = Func::wrap_async(
         &mut store,
         |caller: Caller<State>, (a, b, c, d): (i32, i64, f32, f64)| {
@@ -175,7 +175,7 @@ async fn call_wrapped_async_func() -> Result<(), Error> {
 fn call_linked_func() -> Result<(), Error> {
     let engine = Engine::default();
     let mut store = Store::new(&engine, State::default());
-    store.call_hook(State::call_hook);
+    store.call_hook(sync_call_hook);
     let mut linker = Linker::new(&engine);
 
     linker.func_wrap(
@@ -243,7 +243,7 @@ async fn call_linked_func_async() -> Result<(), Error> {
     config.async_support(true);
     let engine = Engine::new(&config)?;
     let mut store = Store::new(&engine, State::default());
-    store.call_hook(State::call_hook);
+    store.call_hook(sync_call_hook);
 
     let f = Func::wrap_async(
         &mut store,
@@ -313,7 +313,7 @@ async fn call_linked_func_async() -> Result<(), Error> {
 #[test]
 fn instantiate() -> Result<(), Error> {
     let mut store = Store::<State>::default();
-    store.call_hook(State::call_hook);
+    store.call_hook(sync_call_hook);
 
     let m = Module::new(store.engine(), "(module)")?;
     Instance::new(&mut store, &m, &[])?;
@@ -334,7 +334,7 @@ async fn instantiate_async() -> Result<(), Error> {
     config.async_support(true);
     let engine = Engine::new(&config)?;
     let mut store = Store::new(&engine, State::default());
-    store.call_hook(State::call_hook);
+    store.call_hook(sync_call_hook);
 
     let m = Module::new(store.engine(), "(module)")?;
     Instance::new_async(&mut store, &m, &[]).await?;
@@ -355,7 +355,7 @@ fn recursion() -> Result<(), Error> {
 
     let engine = Engine::default();
     let mut store = Store::new(&engine, State::default());
-    store.call_hook(State::call_hook);
+    store.call_hook(sync_call_hook);
     let mut linker = Linker::new(&engine);
 
     linker.func_wrap("host", "f", |mut caller: Caller<State>, n: i32| {
@@ -474,7 +474,7 @@ fn trapping() -> Result<(), Error> {
 
     let run = |action: i32, recur: bool| -> (State, Option<Error>) {
         let mut store = Store::new(&engine, State::default());
-        store.call_hook(State::call_hook);
+        store.call_hook(sync_call_hook);
         let inst = linker
             .instantiate(&mut store, &module)
             .expect("instantiate");
@@ -556,8 +556,12 @@ async fn basic_async_hook() -> Result<(), Error> {
 
     #[async_trait::async_trait]
     impl CallHookHandler<State> for HandlerR {
-        async fn handle_call_event(&self, obj: &mut State, ch: CallHook) -> Result<()> {
-            State::call_hook(obj, ch)
+        async fn handle_call_event(
+            &self,
+            ctx: StoreContextMut<'_, State>,
+            ch: CallHook,
+        ) -> Result<()> {
+            sync_call_hook(ctx, ch)
         }
     }
     let mut config = Config::new();
@@ -630,7 +634,12 @@ async fn timeout_async_hook() -> Result<(), Error> {
 
     #[async_trait::async_trait]
     impl CallHookHandler<State> for HandlerR {
-        async fn handle_call_event(&self, obj: &mut State, ch: CallHook) -> Result<()> {
+        async fn handle_call_event(
+            &self,
+            mut ctx: StoreContextMut<'_, State>,
+            ch: CallHook,
+        ) -> Result<()> {
+            let obj = ctx.data_mut();
             if obj.calls_into_host > 200 {
                 bail!("timeout");
             }
@@ -706,7 +715,12 @@ async fn drop_suspended_async_hook() -> Result<(), Error> {
 
     #[async_trait::async_trait]
     impl CallHookHandler<u32> for Handler {
-        async fn handle_call_event(&self, state: &mut u32, _ch: CallHook) -> Result<()> {
+        async fn handle_call_event(
+            &self,
+            mut ctx: StoreContextMut<'_, u32>,
+            _ch: CallHook,
+        ) -> Result<()> {
+            let state = ctx.data_mut();
             assert_eq!(*state, 0);
             *state += 1;
             let _dec = Decrement(state);
@@ -891,3 +905,7 @@ impl State {
         Ok(())
     }
 }
+
+fn sync_call_hook(mut ctx: StoreContextMut<'_, State>, transition: CallHook) -> Result<()> {
+    ctx.data_mut().call_hook(transition)
+}