Reorganize RR module with hooks

Author: arjunr2

crates/wasmtime/src/runtime.rs

@@ -49,7 +49,6 @@ pub(crate) mod memory;
 pub(crate) mod module;
 pub(crate) mod resources;
 pub(crate) mod rr;
-pub(crate) mod rr_hooks;
 pub(crate) mod store;
 pub(crate) mod trampoline;
 pub(crate) mod trap;

crates/wasmtime/src/runtime/component/func.rs

@@ -4,7 +4,7 @@ use crate::component::storage::storage_as_slice;
 use crate::component::types::Type;
 use crate::component::values::Val;
 use crate::prelude::*;
-use crate::rr_hooks::component_hooks;
+use crate::rr::component_hooks;
 use crate::runtime::vm::component::{ComponentInstance, InstanceFlags, ResourceTables};
 use crate::runtime::vm::{Export, VMFuncRef};
 use crate::store::StoreOpaque;

crates/wasmtime/src/runtime/component/func/host.rs

@@ -5,9 +5,9 @@ use crate::component::matching::InstanceType;
 use crate::component::storage::{slice_to_storage_mut, storage_as_slice_mut};
 use crate::component::{ComponentNamedList, ComponentType, Instance, Lift, Lower, Val};
 use crate::prelude::*;
-use crate::rr_hooks;
+use crate::rr;
 #[cfg(feature = "rr-component")]
-use crate::rr_hooks::component_hooks::ReplayLoweringPhase;
+use crate::rr::component_hooks::ReplayLoweringPhase;
 use crate::runtime::vm::component::{
     ComponentInstance, VMComponentContext, VMLowering, VMLoweringCallee,
 };
@@ -257,7 +257,7 @@ where
 
     let types = vminstance.component().types().clone();
 
-    rr_hooks::component_hooks::record_replay_host_func_entry(storage, &ty, store.0)?;
+    rr::component_hooks::record_replay_host_func_entry(storage, &ty, store.0)?;
 
     let ty = &types[ty];
     let param_tys = InterfaceType::Tuple(ty.params);
@@ -621,27 +621,27 @@ where
         ) -> Result<()> {
             match self.lower_dst() {
                 Dst::Direct(storage) => {
-                    let result = rr_hooks::component_hooks::record_lower_flat(
+                    let result = rr::component_hooks::record_lower_flat(
                         |cx, ty| ret.linear_lower_to_flat(cx, ty, storage),
                         cx,
                         ty,
                     );
-                    rr_hooks::component_hooks::record_host_func_return(
+                    rr::component_hooks::record_host_func_return(
                         unsafe { storage_as_slice_mut(storage) },
                         cx.store.0,
                     )?;
                     result
                 }
                 Dst::Indirect(ptr) => {
                     let ptr = validate_inbounds::<R>(cx.as_slice(), ptr)?;
-                    let result = rr_hooks::component_hooks::record_lower_memory(
+                    let result = rr::component_hooks::record_lower_memory(
                         |cx, ty, ptr| ret.linear_lower_to_memory(cx, ty, ptr),
                         cx,
                         ty,
                         ptr,
                     );
                     // Recording here is just for marking the return event
-                    rr_hooks::component_hooks::record_host_func_return(&[], cx.store.0)?;
+                    rr::component_hooks::record_host_func_return(&[], cx.store.0)?;
                     result
                 }
             }
@@ -811,7 +811,7 @@ where
 
     let types = instance.id().get(store.0).component().types().clone();
 
-    rr_hooks::component_hooks::record_replay_host_func_entry(storage, &ty, store.0)?;
+    rr::component_hooks::record_replay_host_func_entry(storage, &ty, store.0)?;
 
     let func_ty = &types[ty];
     let param_tys = &types[func_ty.params];
@@ -924,28 +924,28 @@ where
             if let Some(cnt) = result_tys.abi.flat_count(MAX_FLAT_RESULTS) {
                 let mut dst = storage[..cnt].iter_mut();
                 for (val, ty) in result_vals.iter().zip(result_tys.types.iter()) {
-                    rr_hooks::component_hooks::record_lower_flat(
+                    rr::component_hooks::record_lower_flat(
                         |cx, ty| val.lower(cx, ty, &mut dst),
                         &mut cx,
                         *ty,
                     )?;
                 }
                 assert!(dst.next().is_none());
-                rr_hooks::component_hooks::record_host_func_return(storage, cx.store.0)?;
+                rr::component_hooks::record_host_func_return(storage, cx.store.0)?;
             } else {
                 let ret_ptr = unsafe { storage[ret_index].assume_init_ref() };
                 let mut ptr = validate_inbounds_dynamic(&result_tys.abi, cx.as_slice(), ret_ptr)?;
                 for (val, ty) in result_vals.iter().zip(result_tys.types.iter()) {
                     let offset = types.canonical_abi(ty).next_field32_size(&mut ptr);
-                    rr_hooks::component_hooks::record_lower_memory(
+                    rr::component_hooks::record_lower_memory(
                         |cx, ty, ptr| val.store(cx, ty, ptr),
                         &mut cx,
                         *ty,
                         offset,
                     )?;
                 }
                 // Recording here is just for marking the return event
-                rr_hooks::component_hooks::record_host_func_return(&[], cx.store.0)?;
+                rr::component_hooks::record_host_func_return(&[], cx.store.0)?;
             }
 
             unsafe {

crates/wasmtime/src/runtime/component/func/options.rs

@@ -11,8 +11,8 @@ use crate::rr::{
 #[cfg(all(feature = "rr-component", feature = "rr-validate"))]
 use crate::rr::{Validate, component_events::ReallocReturnEvent};
 #[cfg(feature = "rr-component")]
-use crate::rr_hooks::component_hooks::ReplayLoweringPhase;
-use crate::rr_hooks::{ConstMemorySliceCell, MemorySliceCell};
+use crate::rr::component_hooks::ReplayLoweringPhase;
+use crate::rr::{ConstMemorySliceCell, MemorySliceCell};
 use crate::runtime::vm::component::{
     CallContexts, ComponentInstance, InstanceFlags, ResourceTable, ResourceTables,
 };

crates/wasmtime/src/runtime/component/func/typed.rs

@@ -3,7 +3,7 @@ use crate::component::func::{Func, LiftContext, LowerContext, Options};
 use crate::component::matching::InstanceType;
 use crate::component::storage::{storage_as_slice, storage_as_slice_mut};
 use crate::prelude::*;
-use crate::rr_hooks::component_hooks;
+use crate::rr::component_hooks;
 use crate::{AsContextMut, StoreContext, StoreContextMut, ValRaw};
 use alloc::borrow::Cow;
 use core::fmt;

crates/wasmtime/src/runtime/func.rs

@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use crate::rr_hooks;
+use crate::rr;
 use crate::runtime::Uninhabited;
 use crate::runtime::vm::{
     InterpreterRef, SendSyncPtr, StoreBox, VMArrayCallHostFuncContext, VMCommonStackInformation,
@@ -2378,7 +2378,7 @@ impl HostContext {
 
             // Record/replay(validation) of the raw parameter arguments
             // Don't need auto-assert GC store here since we aren't using P, just raw args
-            rr_hooks::core_hooks::record_replay_host_func_entry(
+            rr::core_hooks::record_replay_host_func_entry(
                 unsafe { &args.as_ref()[..num_params] },
                 &func_type_index,
                 caller.store.0,
@@ -2423,14 +2423,14 @@ impl HostContext {
                     unsafe { ret.store(&mut store, args.as_mut())? };
                 }
                 // Record the return values
-                rr_hooks::core_hooks::record_host_func_return(
+                rr::core_hooks::record_host_func_return(
                     unsafe { &args.as_ref()[..num_results] },
                     &func_type_index,
                     caller.store.0,
                 )?;
             } else {
                 // Replay the return values
-                rr_hooks::core_hooks::replay_host_func_return(
+                rr::core_hooks::replay_host_func_return(
                     unsafe { &mut args.as_mut()[..num_results] },
                     &func_type_index,
                     caller.store.0,

crates/wasmtime/src/runtime/rr.rs

@@ -0,0 +1,13 @@
+//! Wasmtime's Record and Replay support.
+//!
+//! This feature is currently not optimized and under development
+
+/// Convenience method hooks for injecting event recording/replaying in the rest of the engine
+mod hooks;
+pub use hooks::{ConstMemorySliceCell, MemorySliceCell, component_hooks, core_hooks};
+
+#[cfg(feature = "rr")]
+/// Core infrastructure for RR support
+mod core;
+#[cfg(feature = "rr")]
+pub use core::*;

crates/wasmtime/src/runtime/rr/mod.rs crates/wasmtime/src/runtime/rr/core.rscrates/wasmtime/src/runtime/rr/mod.rs renamed to crates/wasmtime/src/runtime/rr/core.rs

@@ -1,12 +1,3 @@
-#![cfg(feature = "rr")]
-//! Wasmtime's Record and Replay support.
-//!
-//! This feature is currently not optimized and under development
-//!
-//! ## Notes
-//!
-//! This module does NOT support RR for component builtins yet.
-
 use crate::config::{ModuleVersionStrategy, ReplaySettings};
 use crate::prelude::*;
 use core::fmt;