Optimize fast path for recording by avoiding flat type computation

Author: arjunr2

crates/wasmtime/src/runtime/rr/core/events/component_events.rs

@@ -224,9 +224,9 @@ pub struct LowerMemoryReturnEvent(pub ResultEvent<(), LowerMemoryError>);
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct WasmFuncReturnEvent(pub ResultEvent<RRFuncArgVals, WasmFuncReturnError>);
 
-impl Validate<Result<RRFuncArgVals>> for WasmFuncReturnEvent {
-    fn validate(&self, expect: &Result<RRFuncArgVals>) -> Result<(), ReplayError> {
-        self.0.validate(expect)
+impl Validate<&Result<RRFuncArgVals>> for WasmFuncReturnEvent {
+    fn validate(&self, expect: &&Result<RRFuncArgVals>) -> Result<(), ReplayError> {
+        self.0.validate(*expect)
     }
 }
 

crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs

@@ -61,7 +61,9 @@ where
         })?;
         store
             .0
-            .next_replay_event_validation::<WasmFuncReturnEvent, Result<RRFuncArgVals>>(&result)?;
+            .next_replay_event_validation::<WasmFuncReturnEvent, _, &Result<RRFuncArgVals>>(
+                || &result,
+            )?;
         result?;
         return Ok(());
     }
@@ -80,10 +82,13 @@ pub fn record_replay_host_func_entry(
     #[cfg(all(feature = "rr-component", feature = "rr-validate"))]
     {
         use crate::rr::component_events::HostFuncEntryEvent;
-        let flat_params = types.flat_types_storage(&InterfaceType::Tuple(types[*type_idx].params));
-        let event = HostFuncEntryEvent::new(args, flat_params, type_idx.clone());
-        store.record_event_validation(|| event.clone())?;
-        store.next_replay_event_validation::<HostFuncEntryEvent, _>(&event)?;
+        let event = || {
+            let flat_params =
+                types.flat_types_storage(&InterfaceType::Tuple(types[*type_idx].params));
+            HostFuncEntryEvent::new(args, flat_params, type_idx.clone())
+        };
+        store.record_event_validation(|| event())?;
+        store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| event())?;
     }
     let _ = (args, types, type_idx, store);
     Ok(())
@@ -98,10 +103,10 @@ pub fn record_host_func_return(
     store: &mut StoreOpaque,
 ) -> Result<()> {
     #[cfg(feature = "rr-component")]
-    {
+    store.record_event(|| {
         let flat_results = types.flat_types_storage(&ty);
-        store.record_event(|| HostFuncReturnEvent::new_from_flat_storage(args, flat_results))?;
-    }
+        HostFuncReturnEvent::new_from_flat_storage(args, flat_results)
+    })?;
     let _ = (args, types, ty, store);
     Ok(())
 }

crates/wasmtime/src/runtime/rr/hooks/core_hooks.rs

@@ -18,9 +18,10 @@ pub fn record_replay_host_func_entry(
     {
         // Record/replay the raw parameter args
         use crate::rr::core_events::HostFuncEntryEvent;
-        let event = HostFuncEntryEvent::new(&args, flat, ty.clone());
-        store.record_event_validation(|| event.clone())?;
-        store.next_replay_event_validation::<HostFuncEntryEvent, _>(&event)?;
+        store.record_event_validation(|| HostFuncEntryEvent::new(&args, flat, ty.clone()))?;
+        store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| {
+            HostFuncEntryEvent::new(&args, flat, ty.clone())
+        })?;
     }
     let _ = (args, flat, ty, store);
     Ok(())

crates/wasmtime/src/runtime/store.rs

@@ -1576,16 +1576,17 @@ impl StoreOpaque {
     /// Convenience wrapper around [`Replayer::next_event_validation`]
     #[cfg(feature = "rr-validate")]
     #[inline]
-    pub(crate) fn next_replay_event_validation<T, Y>(
+    pub(crate) fn next_replay_event_validation<T, F, Y>(
         &mut self,
-        expect: &Y,
+        expect: F,
     ) -> Result<(), ReplayError>
     where
         T: TryFrom<RREvent> + Validate<Y>,
+        F: FnOnce() -> Y,
         ReplayError: From<<T as TryFrom<RREvent>>::Error>,
     {
         if let Some(buf) = self.replay_buffer_mut() {
-            buf.next_event_validation::<T, Y>(expect)
+            buf.next_event_validation::<T, Y>(&expect())
         } else {
             Ok(())
         }