Convert all flat type construction to zero-copy iterators

Author: arjunr2

crates/wasmtime/src/runtime/func.rs

@@ -1291,12 +1291,11 @@ impl Func {
         let flat_params = ty
             .params()
             .into_iter()
-            .map(|x| x.to_wasm_type().byte_size())
-            .collect::<Vec<_>>();
+            .map(|x| x.to_wasm_type().byte_size());
 
         rr::core_hooks::record_and_replay_validate_host_func_entry(
             values_vec,
-            &flat_params,
+            flat_params,
             &mut caller.store.0,
         )?;
 
@@ -1322,13 +1321,8 @@ impl Func {
             let flat_results = ty
                 .results()
                 .into_iter()
-                .map(|x| x.to_wasm_type().byte_size())
-                .collect::<Vec<_>>();
-            rr::core_hooks::record_host_func_return(
-                values_vec,
-                &flat_results,
-                &mut caller.store.0,
-            )?;
+                .map(|x| x.to_wasm_type().byte_size());
+            rr::core_hooks::record_host_func_return(values_vec, flat_results, &mut caller.store.0)?;
         } else {
             rr::core_hooks::replay_host_func_return(values_vec, &mut caller.store.0)?;
         }
@@ -2430,30 +2424,23 @@ impl HostContext {
             };
             let func = &state.func;
 
-            let (flat_size_params, flat_size_results) = {
-                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()
-                        .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) = (flat_size_params.len(), flat_size_results.len());
+            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();
+            let (num_params, flat_size_params) = (
+                wasm_func_type.params().len(),
+                wasm_func_type.params().into_iter().map(|x| x.byte_size()),
+            );
+            let (num_results, flat_size_results) = (
+                wasm_func_type.returns().len(),
+                wasm_func_type.returns().into_iter().map(|x| x.byte_size()),
+            );
 
             // 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_and_replay_validate_host_func_entry(
                 unsafe { &args.as_ref()[..num_params] },
-                flat_size_params.as_slice(),
+                flat_size_params,
                 caller.store.0,
             )?;
 
@@ -2498,7 +2485,7 @@ impl HostContext {
                 // Record the return values
                 rr::core_hooks::record_host_func_return(
                     unsafe { &args.as_ref()[..num_results] },
-                    flat_size_results.as_slice(),
+                    flat_size_results,
                     caller.store.0,
                 )?;
             } else {

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

@@ -62,7 +62,8 @@ impl fmt::Debug for RRFuncArgVals {
 
 impl RRFuncArgVals {
     /// Construct [`RRFuncArgVals`] from raw value buffer and a flat size iterator
-    pub fn from_raw_slice<T>(args: &[T], flat: impl Iterator<Item = u8>) -> RRFuncArgVals
+    #[inline]
+    pub fn from_flat_iter<T>(args: &[T], flat: impl Iterator<Item = u8>) -> RRFuncArgVals
     where
         T: FlatBytes,
     {
@@ -77,22 +78,16 @@ impl RRFuncArgVals {
 
     /// Construct [`RRFuncArgVals`] from raw value buffer and a [`FlatTypesStorage`]
     #[cfg(feature = "rr-component")]
+    #[inline]
     pub fn from_flat_storage<T>(args: &[T], flat: FlatTypesStorage) -> RRFuncArgVals
     where
         T: FlatBytes,
     {
-        RRFuncArgVals::from_raw_slice(args, flat.iter32())
-    }
-
-    /// Construct [`RRFuncArgVals`] from raw value buffer and a `[&u8]` slice
-    pub fn from_flat_u8<T>(args: &[T], flat: &[u8]) -> RRFuncArgVals
-    where
-        T: FlatBytes,
-    {
-        RRFuncArgVals::from_raw_slice(args, flat.iter().copied())
+        RRFuncArgVals::from_flat_iter(args, flat.iter32())
     }
 
     /// Encode [`RRFuncArgVals`] back into raw value buffer
+    #[inline]
     pub fn into_raw_slice<T>(self, raw_args: &mut [T])
     where
         T: FlatBytes,
@@ -104,7 +99,8 @@ impl RRFuncArgVals {
         }
     }
 
-    /// Generate a vector of [`crate::Val`] from [`RRFuncArgVals`]
+    /// Generate a vector of [`crate::Val`] from [`RRFuncArgVals`] and [`ValType`]s
+    #[inline]
     pub fn to_val_vec(self, mut store: impl AsContextMut, val_types: Vec<ValType>) -> Vec<Val> {
         let mut pos = 0;
         let mut vals = Vec::new();

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

@@ -56,7 +56,7 @@ where
             &InterfaceType::Tuple(types[type_idx].results),
             MAX_FLAT_RESULTS,
         );
-        let result = result.map(|_| RRFuncArgVals::from_raw_slice(args, flat_results.iter32()));
+        let result = result.map(|_| RRFuncArgVals::from_flat_iter(args, flat_results.iter32()));
         store.0.record_event_validation(|| {
             WasmFuncReturnEvent(ResultEvent::from_anyhow_result(&result))
         })?;

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

@@ -34,14 +34,11 @@ where
                 .unwrap()
                 .checksum();
             store.0.record_event(|| {
-                let flat = ty
-                    .params()
-                    .map(|t| t.to_wasm_type().byte_size())
-                    .collect::<Vec<u8>>();
+                let flat = ty.params().map(|t| t.to_wasm_type().byte_size());
                 WasmFuncEntryEvent {
                     module: checksum,
                     origin,
-                    args: RRFuncArgVals::from_flat_u8(args, &flat),
+                    args: RRFuncArgVals::from_flat_iter(args, flat),
                 }
             })?;
         }
@@ -50,11 +47,8 @@ where
     #[cfg(all(feature = "rr", feature = "rr-validate"))]
     {
         if origin.is_some() {
-            let flat = ty
-                .results()
-                .map(|t| t.to_wasm_type().byte_size())
-                .collect::<Vec<u8>>();
-            let result = result.map(|_| RRFuncArgVals::from_raw_slice(args, flat.iter().copied()));
+            let flat = ty.results().map(|t| t.to_wasm_type().byte_size());
+            let result = result.map(|_| RRFuncArgVals::from_flat_iter(args, flat));
             store.0.record_event_validation(|| {
                 WasmFuncReturnEvent(ResultEvent::from_anyhow_result(&result))
             })?;
@@ -75,38 +69,47 @@ where
 #[inline]
 pub fn record_and_replay_validate_host_func_entry<T>(
     args: &[T],
-    flat: &[u8],
+    flat: impl Iterator<Item = u8>,
     store: &mut StoreOpaque,
 ) -> Result<()>
 where
     T: FlatBytes,
 {
+    let _ = (args, &flat, &store);
     #[cfg(all(feature = "rr", feature = "rr-validate"))]
     {
         // Record/replay the raw parameter args
-        store.record_event_validation(|| HostFuncEntryEvent {
-            args: RRFuncArgVals::from_flat_u8(args, flat),
-        })?;
-        store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| HostFuncEntryEvent {
-            args: RRFuncArgVals::from_flat_u8(args, flat),
-        })?;
+        if store.replay_enabled() {
+            store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| {
+                HostFuncEntryEvent {
+                    args: RRFuncArgVals::from_flat_iter(args, flat),
+                }
+            })?;
+        } else {
+            store.record_event_validation(|| HostFuncEntryEvent {
+                args: RRFuncArgVals::from_flat_iter(args, flat),
+            })?;
+        }
     }
-    let _ = (args, flat, store);
     Ok(())
 }
 
 /// Record hook operation for host function return events
 #[inline]
-pub fn record_host_func_return<T>(args: &[T], flat: &[u8], store: &mut StoreOpaque) -> Result<()>
+pub fn record_host_func_return<T>(
+    args: &[T],
+    flat: impl Iterator<Item = u8>,
+    store: &mut StoreOpaque,
+) -> Result<()>
 where
     T: FlatBytes,
 {
+    let _ = (args, &flat, &store);
     // Record the return values
     #[cfg(feature = "rr")]
     store.record_event(|| HostFuncReturnEvent {
-        args: RRFuncArgVals::from_flat_u8(args, flat),
+        args: RRFuncArgVals::from_flat_iter(args, flat),
     })?;
-    let _ = (args, flat, store);
     Ok(())
 }
 

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

@@ -6,7 +6,7 @@ use crate::rr::{
     RREvent, ReplayError, Validate, component_hooks::ReplayLoweringPhase, core_events,
 };
 use crate::{
-    AsContextMut, Engine, Module, ReplayReader, ReplaySettings, Store, ValRaw, ValType, prelude::*,
+    AsContextMut, Engine, Module, ReplayReader, ReplaySettings, Store, ValRaw, prelude::*,
 };
 use alloc::collections::BTreeMap;
 use core::mem::MaybeUninit;
@@ -251,7 +251,7 @@ impl<'a> ReplayInstance<'a> {
                     .into_func()
                     .ok_or(ReplayError::InvalidCoreFuncIndex(entity))?;
 
-                let params_ty = func.ty(&store).params().collect::<Vec<ValType>>();
+                let params_ty = func.ty(&store).params().collect::<Vec<_>>();
 
                 // Obtain the argument values for function call
                 let mut results = vec![crate::Val::I64(0); func.ty(&store).results().len()];