Remove module checksums on wasm function begin events; only rely on instance id

Author: arjunr2

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

@@ -552,14 +552,11 @@ impl Func {
         {
             use crate::rr::component_events::WasmFuncBeginEvent;
 
-            let component = *self.instance.id().get(store.0).component().checksum();
             let instance = self.instance.id().instance();
             let func_idx = self.index;
-            store.0.record_event(|| WasmFuncBeginEvent {
-                component,
-                instance,
-                func_idx,
-            })?;
+            store
+                .0
+                .record_event(|| WasmFuncBeginEvent { instance, func_idx })?;
         }
         let export = self.lifted_core_func(store.0);
 

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

@@ -13,8 +13,6 @@ use wasmtime_environ::{
 /// Beginning marker for a Wasm component function call from host
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct WasmFuncBeginEvent {
-    /// Checksum of component containing function
-    pub component: [u8; 32],
     /// Instance ID for the component instance
     pub instance: ComponentInstanceId,
     /// Export index for the invoked function

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

@@ -16,8 +16,6 @@ pub struct InstantiationEvent {
 /// A call event from Host into a core Wasm function
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct WasmFuncEntryEvent {
-    /// Checksum of module containing function
-    pub module: [u8; 32],
     /// Origin (instance + function index) for this function
     pub origin: WasmFuncOrigin,
     /// Raw values passed across call boundary

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

@@ -5,8 +5,6 @@ use crate::rr::core_events::{HostFuncReturnEvent, WasmFuncEntryEvent};
 use crate::rr::{
     RRFuncArgVals, ResultEvent, core_events::HostFuncEntryEvent, core_events::WasmFuncReturnEvent,
 };
-#[cfg(feature = "rr")]
-use crate::store::StoreInstanceId;
 use crate::store::StoreOpaque;
 use crate::{FuncType, StoreContextMut, ValRaw, WasmFuncOrigin, prelude::*};
 
@@ -28,15 +26,9 @@ where
     #[cfg(feature = "rr")]
     {
         if let Some(origin) = origin {
-            let checksum = *store
-                .0
-                .module_for_instance(StoreInstanceId::new(store.0.id(), origin.instance))
-                .unwrap()
-                .checksum();
             store.0.record_event(|| {
                 let flat = ty.params().map(|t| t.to_wasm_type().byte_size());
                 WasmFuncEntryEvent {
-                    module: checksum,
                     origin,
                     args: RRFuncArgVals::from_flat_iter(args, flat),
                 }

crates/wasmtime/src/runtime/rr/replay_driver.rs

@@ -1,10 +1,11 @@
 #[cfg(feature = "rr-component")]
-use crate::component::{self, Component};
+use crate::component::{self, Component, ComponentInstanceId};
 #[cfg(feature = "rr-component")]
 use crate::rr::component_events;
 use crate::rr::{
     RREvent, ReplayError, Validate, component_hooks::ReplayLoweringPhase, core_events,
 };
+use crate::store::InstanceId;
 use crate::{
     AsContextMut, Engine, Module, ReplayReader, ReplaySettings, Store, ValRaw, prelude::*,
 };
@@ -92,8 +93,9 @@ pub struct ReplayInstance<T: 'static> {
     store: Store<T>,
     component_linker: component::Linker<T>,
     module_linker: crate::Linker<T>,
-    module_instances: BTreeMap<core_events::InstantiationEvent, crate::Instance>,
-    component_instances: BTreeMap<component_events::InstantiationEvent, component::Instance>,
+    module_instances: BTreeMap<InstanceId, crate::Instance>,
+    #[cfg(feature = "rr-component")]
+    component_instances: BTreeMap<ComponentInstanceId, component::Instance>,
 }
 
 impl<T: 'static> ReplayInstance<T> {
@@ -119,6 +121,7 @@ impl<T: 'static> ReplayInstance<T> {
             component_linker,
             module_linker,
             module_instances: BTreeMap::new(),
+            #[cfg(feature = "rr-component")]
             component_instances: BTreeMap::new(),
         })
     }
@@ -161,9 +164,8 @@ impl<T: 'static> ReplayInstance<T> {
                         instance: instance.id().instance(),
                     })?;
 
-                    let ret = self.component_instances.insert(event, instance);
-                    // Ensures that an already-instantiated configuration is not re-instantiated
-                    assert!(ret.is_none());
+                    self.component_instances
+                        .insert(instance.id().instance(), instance);
                 }
                 #[cfg(not(feature = "rr-component"))]
                 {
@@ -176,14 +178,11 @@ impl<T: 'static> ReplayInstance<T> {
                 #[cfg(feature = "rr-component")]
                 {
                     // Grab the correct component instance
-                    let key = component_events::InstantiationEvent {
-                        component: event.component,
-                        instance: event.instance,
-                    };
+                    let key = event.instance;
                     let instance = self
                         .component_instances
                         .get_mut(&key)
-                        .ok_or(ReplayError::MissingComponentInstance(key.instance.as_u32()))?;
+                        .ok_or(ReplayError::MissingComponentInstance(key.as_u32()))?;
 
                     // Replay lowering steps and obtain raw value arguments to raw function call
                     let func = component::Func::from_lifted_func(*instance, event.func_idx);
@@ -247,20 +246,15 @@ impl<T: 'static> ReplayInstance<T> {
                     instance: instance.id(),
                 })?;
 
-                let ret = self.module_instances.insert(event, instance);
-                // Ensures that an already-instantiated configuration is not re-instantiated
-                assert!(ret.is_none());
+                self.module_instances.insert(instance.id(), instance);
             }
             RREvent::CoreWasmFuncEntry(event) => {
                 // Grab the correct module instance
-                let key = core_events::InstantiationEvent {
-                    module: event.module,
-                    instance: event.origin.instance,
-                };
+                let key = event.origin.instance;
                 let instance = self
                     .module_instances
                     .get_mut(&key)
-                    .ok_or(ReplayError::MissingModuleInstance(key.instance.as_u32()))?;
+                    .ok_or(ReplayError::MissingModuleInstance(key.as_u32()))?;
 
                 let entity = EntityIndex::from(event.origin.index);
                 let mut store = self.store.as_context_mut();
@@ -316,9 +310,8 @@ impl<T: 'static> ReplayInstance<T> {
                         instance: instance.id().instance(),
                     })?;
 
-                    let ret = self.component_instances.insert(event, instance);
-                    // Ensures that an already-instantiated configuration is not re-instantiated
-                    assert!(ret.is_none());
+                    self.component_instances
+                        .insert(instance.id().instance(), instance);
                 }
                 #[cfg(not(feature = "rr-component"))]
                 {
@@ -331,14 +324,11 @@ impl<T: 'static> ReplayInstance<T> {
                 #[cfg(feature = "rr-component")]
                 {
                     // Grab the correct component instance
-                    let key = component_events::InstantiationEvent {
-                        component: event.component,
-                        instance: event.instance,
-                    };
+                    let key = event.instance;
                     let instance = self
                         .component_instances
                         .get_mut(&key)
-                        .ok_or(ReplayError::MissingComponentInstance(key.instance.as_u32()))?;
+                        .ok_or(ReplayError::MissingComponentInstance(key.as_u32()))?;
 
                     // Replay lowering steps and obtain raw value arguments to raw function call
                     let func = component::Func::from_lifted_func(*instance, event.func_idx);
@@ -405,20 +395,15 @@ impl<T: 'static> ReplayInstance<T> {
                     instance: instance.id(),
                 })?;
 
-                let ret = self.module_instances.insert(event, instance);
-                // Ensures that an already-instantiated configuration is not re-instantiated
-                assert!(ret.is_none());
+                self.module_instances.insert(instance.id(), instance);
             }
             RREvent::CoreWasmFuncEntry(event) => {
                 // Grab the correct module instance
-                let key = core_events::InstantiationEvent {
-                    module: event.module,
-                    instance: event.origin.instance,
-                };
+                let key = event.origin.instance;
                 let instance = self
                     .module_instances
                     .get_mut(&key)
-                    .ok_or(ReplayError::MissingModuleInstance(key.instance.as_u32()))?;
+                    .ok_or(ReplayError::MissingModuleInstance(key.as_u32()))?;
 
                 let entity = EntityIndex::from(event.origin.index);
                 let mut store = self.store.as_context_mut();