wizer: Don't copy all memory out of a store

crates/wizer/src/component.rs

@@ -58,7 +58,8 @@ pub trait ComponentInstanceState: Send {
         &mut self,
         instance: &str,
         func: &str,
-    ) -> impl Future<Output = Vec<u8>> + Send;
+        contents: impl FnOnce(&[u8]) + Send,
+    ) -> impl Future<Output = ()> + Send;
 
     /// Same as [`Self::call_func_ret_list_u8`], but for the `s32` WIT type.
     fn call_func_ret_s32(&mut self, instance: &str, func: &str)

crates/wizer/src/component/snapshot.rs

@@ -115,7 +115,7 @@ where
         }
     }
 
-    async fn memory_contents(&mut self, name: &str) -> Vec<u8> {
+    async fn memory_contents(&mut self, name: &str, contents: impl FnOnce(&[u8]) + Send) {
         let Accessor::Memory {
             accessor_export_name,
             ..
@@ -124,7 +124,7 @@ where
             panic!("expected memory accessor for {name}");
         };
         self.ctx
-            .call_func_ret_list_u8(WIZER_INSTANCE, accessor_export_name)
+            .call_func_ret_list_u8(WIZER_INSTANCE, accessor_export_name, contents)
             .await
     }
 }

crates/wizer/src/component/wasmtime.rs

@@ -1,7 +1,9 @@
 use crate::Wizer;
 use crate::component::ComponentInstanceState;
 use anyhow::{Context, anyhow};
-use wasmtime::component::{Component, ComponentExportIndex, Instance, Lift, types::ComponentItem};
+use wasmtime::component::{
+    Component, ComponentExportIndex, Instance, Lift, WasmList, types::ComponentItem,
+};
 use wasmtime::{Result, Store};
 
 impl Wizer {
@@ -84,7 +86,12 @@ pub struct WasmtimeWizerComponent<'a, T: 'static> {
 }
 
 impl<T: Send> WasmtimeWizerComponent<'_, T> {
-    async fn call_func<R>(&mut self, instance: &str, func: &str) -> R
+    async fn call_func<R, R2>(
+        &mut self,
+        instance: &str,
+        func: &str,
+        use_ret: impl FnOnce(&mut Store<T>, R) -> R2,
+    ) -> R2
     where
         R: Lift + 'static,
     {
@@ -102,29 +109,40 @@ impl<T: Send> WasmtimeWizerComponent<'_, T> {
             .get_typed_func::<(), (R,)>(&mut *self.store, func_export)
             .unwrap();
         let ret = func.call_async(&mut *self.store, ()).await.unwrap().0;
+        let ret = use_ret(&mut *self.store, ret);
         func.post_return_async(&mut *self.store).await.unwrap();
         ret
     }
 }
 
 impl<T: Send> ComponentInstanceState for WasmtimeWizerComponent<'_, T> {
-    async fn call_func_ret_list_u8(&mut self, instance: &str, func: &str) -> Vec<u8> {
-        self.call_func(instance, func).await
+    async fn call_func_ret_list_u8(
+        &mut self,
+        instance: &str,
+        func: &str,
+        contents: impl FnOnce(&[u8]) + Send,
+    ) {
+        self.call_func(instance, func, |store, list: WasmList<u8>| {
+            contents(list.as_le_slice(&store));
+        })
+        .await
     }
 
     async fn call_func_ret_s32(&mut self, instance: &str, func: &str) -> i32 {
-        self.call_func(instance, func).await
+        self.call_func(instance, func, |_, r| r).await
     }
 
     async fn call_func_ret_s64(&mut self, instance: &str, func: &str) -> i64 {
-        self.call_func(instance, func).await
+        self.call_func(instance, func, |_, r| r).await
     }
 
     async fn call_func_ret_f32(&mut self, instance: &str, func: &str) -> u32 {
-        self.call_func::<f32>(instance, func).await.to_bits()
+        self.call_func(instance, func, |_, r: f32| r.to_bits())
+            .await
     }
 
     async fn call_func_ret_f64(&mut self, instance: &str, func: &str) -> u64 {
-        self.call_func::<f64>(instance, func).await.to_bits()
+        self.call_func(instance, func, |_, r: f64| r.to_bits())
+            .await
     }
 }

crates/wizer/src/lib.rs

@@ -386,5 +386,9 @@ pub trait InstanceState {
     /// # Panics
     ///
     /// This function panics if `name` isn't an exported memory.
-    fn memory_contents(&mut self, name: &str) -> impl Future<Output = Vec<u8>> + Send;
+    fn memory_contents(
+        &mut self,
+        name: &str,
+        contents: impl FnOnce(&[u8]) + Send,
+    ) -> impl Future<Output = ()> + Send;
 }

crates/wizer/src/rewrite.rs

@@ -25,11 +25,12 @@ impl Wizer {
         // than the original, uninitialized data segments.
         let add_data_segments = |data_section: &mut wasm_encoder::DataSection| {
             for seg in &snapshot.data_segments {
-                data_section.active(
-                    seg.memory_index,
-                    &ConstExpr::i32_const(seg.offset as i32),
-                    seg.data().iter().copied(),
-                );
+                let offset = if seg.is64 {
+                    ConstExpr::i64_const(seg.offset.cast_signed())
+                } else {
+                    ConstExpr::i32_const(u32::try_from(seg.offset).unwrap().cast_signed())
+                };
+                data_section.active(seg.memory_index, &offset, seg.data.iter().copied());
             }
         };