Added flat type encoding/decoding for events; reformatted validation events

Author: arjunr2

crates/environ/src/component/types.rs

@@ -1343,7 +1343,8 @@ pub enum FlatType {
 }
 
 impl FlatType {
-    const fn byte_size(&self) -> u8 {
+    /// Return the size in bytes for this flat type
+    pub const fn byte_size(&self) -> u8 {
         match self {
             FlatType::I32 | FlatType::F32 => 4,
             FlatType::I64 | FlatType::F64 => 8,

crates/environ/src/component/types_builder.rs

@@ -879,6 +879,14 @@ impl FlatTypesStorage {
             false
         }
     }
+
+    /// Generate an iterator over the 32-bit flat encoding
+    pub fn iter32(&self) -> impl Iterator<Item = u8> {
+        self.memory32
+            .iter()
+            .take(self.len as usize)
+            .map(|f| f.byte_size())
+    }
 }
 
 impl FlatType {

crates/environ/src/types.rs

@@ -233,6 +233,15 @@ impl WasmValType {
         }
     }
 
+    /// Return the number of bytes needed to represent this value
+    pub fn byte_size(&self) -> u8 {
+        match self {
+            WasmValType::I32 | WasmValType::F32 => 4,
+            WasmValType::I64 | WasmValType::F64 => 8,
+            WasmValType::V128 | WasmValType::Ref(_) => 16,
+        }
+    }
+
     /// Returns the contained reference type.
     ///
     /// Panics if the value type is not a vmgcref

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

@@ -596,7 +596,8 @@ impl Func {
                 |store| crate::Func::call_unchecked_raw(store, export, params_and_returns),
                 params_and_returns.as_ref(),
                 self.index,
-                *self.instance.id().get(store.0).component().checksum(),
+                self.abi_info(store.0).2,
+                self.instance.id(),
                 &mut store,
             )?;
         }

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

@@ -257,11 +257,11 @@ where
 
     let types = vminstance.component().types().clone();
 
-    rr::component_hooks::record_replay_host_func_entry(storage, &ty, store.0)?;
+    rr::component_hooks::record_replay_host_func_entry(storage, &types, &ty, store.0)?;
 
-    let ty = &types[ty];
-    let param_tys = InterfaceType::Tuple(ty.params);
-    let result_tys = InterfaceType::Tuple(ty.results);
+    let func_ty = &types[ty];
+    let param_tys = InterfaceType::Tuple(func_ty.params);
+    let result_tys = InterfaceType::Tuple(func_ty.results);
 
     let storage_type = if async_ {
         #[cfg(feature = "component-model-async")]
@@ -628,6 +628,8 @@ where
                     );
                     rr::component_hooks::record_host_func_return(
                         unsafe { storage_as_slice_mut(storage) },
+                        cx.types,
+                        &ty,
                         cx.store.0,
                     )?;
                     result
@@ -641,7 +643,7 @@ where
                         ptr,
                     );
                     // Recording here is just for marking the return event
-                    rr::component_hooks::record_host_func_return(&[], cx.store.0)?;
+                    rr::component_hooks::record_host_func_return(&[], cx.types, &ty, cx.store.0)?;
                     result
                 }
             }
@@ -811,7 +813,7 @@ where
 
     let types = instance.id().get(store.0).component().types().clone();
 
-    rr::component_hooks::record_replay_host_func_entry(storage, &ty, store.0)?;
+    rr::component_hooks::record_replay_host_func_entry(storage, &types, &ty, store.0)?;
 
     let func_ty = &types[ty];
     let param_tys = &types[func_ty.params];
@@ -931,7 +933,12 @@ where
                     )?;
                 }
                 assert!(dst.next().is_none());
-                rr::component_hooks::record_host_func_return(storage, cx.store.0)?;
+                rr::component_hooks::record_host_func_return(
+                    storage,
+                    cx.types,
+                    &InterfaceType::Tuple(func_ty.results),
+                    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)?;
@@ -945,7 +952,12 @@ where
                     )?;
                 }
                 // Recording here is just for marking the return event
-                rr::component_hooks::record_host_func_return(&[], cx.store.0)?;
+                rr::component_hooks::record_host_func_return(
+                    &[],
+                    cx.types,
+                    &InterfaceType::Tuple(func_ty.results),
+                    cx.store.0,
+                )?;
             }
 
             unsafe {

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

@@ -4,15 +4,17 @@ use crate::component::matching::InstanceType;
 use crate::component::resources::{HostResourceData, HostResourceIndex, HostResourceTables};
 use crate::component::{Instance, ResourceType};
 use crate::prelude::*;
+#[cfg(all(feature = "rr-component", feature = "rr-validate"))]
+use crate::rr::component_events::ResultEvent;
+#[cfg(feature = "rr-component")]
+use crate::rr::component_hooks::ReplayLoweringPhase;
+use crate::rr::{ConstMemorySliceCell, MemorySliceCell};
 #[cfg(feature = "rr-component")]
 use crate::rr::{
     RREvent, RecordBuffer, ReplayError, Replayer, component_events::ReallocEntryEvent,
 };
 #[cfg(all(feature = "rr-component", feature = "rr-validate"))]
 use crate::rr::{Validate, component_events::ReallocReturnEvent};
-#[cfg(feature = "rr-component")]
-use crate::rr::component_hooks::ReplayLoweringPhase;
-use crate::rr::{ConstMemorySliceCell, MemorySliceCell};
 use crate::runtime::vm::component::{
     CallContexts, ComponentInstance, InstanceFlags, ResourceTable, ResourceTables,
 };
@@ -404,9 +406,9 @@ impl<'a, T: 'static> LowerContext<'a, T> {
         })?;
         let result = self.realloc_inner(old, old_size, old_align, new_size);
         #[cfg(all(feature = "rr-component", feature = "rr-validate"))]
-        self.store
-            .0
-            .record_event_validation(|| ReallocReturnEvent::from_anyhow_result(&result))?;
+        self.store.0.record_event_validation(|| {
+            ReallocReturnEvent(ResultEvent::from_anyhow_result(&result))
+        })?;
         result
     }
 
@@ -636,7 +638,7 @@ impl<'a, T: 'static> LowerContext<'a, T> {
                     if run_validate {
                         _lower_stack.pop().ok_or(ReplayError::InvalidOrdering)?;
                     }
-                    lowering_error = e.ret().map_err(Into::into).err();
+                    lowering_error = e.0.ret().map_err(Into::into).err();
                 }
                 RREvent::ComponentLowerMemoryReturn(e) => {
                     #[cfg(feature = "rr-validate")]
@@ -645,7 +647,7 @@ impl<'a, T: 'static> LowerContext<'a, T> {
                             .pop()
                             .ok_or(ReplayError::InvalidOrdering)?;
                     }
-                    lowering_error = e.ret().map_err(Into::into).err();
+                    lowering_error = e.0.ret().map_err(Into::into).err();
                 }
                 RREvent::ComponentMemorySliceWrite(e) => {
                     // The bounds check is performed here is required here (in the absence of
@@ -665,7 +667,7 @@ impl<'a, T: 'static> LowerContext<'a, T> {
                 {
                     #[cfg(feature = "rr-validate")]
                     if run_validate {
-                        lowering_error = _e.validate(&_realloc_stack.pop().unwrap()).err()
+                        lowering_error = _e.0.validate(&_realloc_stack.pop().unwrap()).err()
                     }
                 }
                 RREvent::ComponentLowerFlatEntry(_) => {

crates/wasmtime/src/runtime/func.rs

@@ -2366,20 +2366,30 @@ impl HostContext {
             let func = &state.func;
 
             let func_type_index = state._ty.index();
-            let (num_params, num_results) = {
+            let (params_size, results_size) = {
                 let type_index = state._ty.index();
                 let wasm_func_subtype = caller.engine().signatures().borrow(type_index).unwrap();
                 let wasm_func_type = wasm_func_subtype.unwrap_func();
                 (
-                    wasm_func_type.params().len(),
-                    wasm_func_type.returns().len(),
+                    wasm_func_type
+                        .params()
+                        .into_iter()
+                        .map(|x| x.byte_size())
+                        .collect::<Vec<_>>(),
+                    wasm_func_type
+                        .returns()
+                        .into_iter()
+                        .map(|x| x.byte_size())
+                        .collect::<Vec<_>>(),
                 )
             };
+            let (num_params, num_results) = (params_size.len(), results_size.len());
 
             // 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::core_hooks::record_replay_host_func_entry(
                 unsafe { &args.as_ref()[..num_params] },
+                params_size.as_slice(),
                 &func_type_index,
                 caller.store.0,
             )?;
@@ -2425,6 +2435,7 @@ impl HostContext {
                 // Record the return values
                 rr::core_hooks::record_host_func_return(
                     unsafe { &args.as_ref()[..num_results] },
+                    results_size.as_slice(),
                     &func_type_index,
                     caller.store.0,
                 )?;

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

@@ -6,12 +6,12 @@ use serde::{Deserialize, Serialize};
 // Use component events internally even without feature flags enabled
 // so that [`RREvent`] has a well-defined serialization format, but export
 // it for other modules only when enabled
+#[cfg(all(feature = "rr-validate", feature = "rr-component"))]
+pub use events::RRFuncArgVals;
 #[cfg(any(feature = "rr-validate", feature = "rr-component"))]
 pub use events::Validate;
 #[cfg(feature = "rr-component")]
 pub use events::component_events;
-#[cfg(all(feature = "rr-validate", feature = "rr-component"))]
-pub use events::{RRFuncArgVals, func_argvals_from_raw_slice};
 use events::{common_events, component_events as __component_events};
 pub use events::{core_events, marker_events};
 pub use io::{RecordWriter, ReplayReader};
@@ -535,20 +535,24 @@ mod tests {
     #[test]
     #[cfg(all(feature = "rr", feature = "rr-component"))]
     fn rr_buffers() -> Result<()> {
+        use wasmtime_environ::component::FlatTypesStorage;
+
         let record_settings = RecordSettings::default();
         let tmp = NamedTempFile::new()?;
         let tmppath = tmp.path().to_str().expect("Filename should be UTF-8");
 
-        let values = vec![ValRaw::i32(1), ValRaw::f32(2), ValRaw::i64(3)]
-            .into_iter()
-            .map(|x| MaybeUninit::new(x))
-            .collect::<Vec<_>>();
+        let values = vec![ValRaw::i32(1), ValRaw::f32(2), ValRaw::i64(3)];
+        let flat = FlatTypesStorage::new();
+        flat.push(FlatType::I32, FlatType::I32);
+        flat.push(FlatType::F32, FlatType::F32);
+        flat.push(FlatType::I64, FlatType::I64);
 
         // Record values
         let mut recorder =
             RecordBuffer::new_recorder(Box::new(File::create(tmppath)?), record_settings)?;
-        recorder
-            .record_event(|| __component_events::HostFuncReturnEvent::new(values.as_slice()))?;
+        recorder.record_event(|| {
+            __component_events::HostFuncReturnEvent::new(values.as_slice(), flat)
+        })?;
         recorder.flush()?;
 
         let tmp = tmp.into_temp_path();