[Rust] use cleanup list as necessary for stream/future writes (#1240)

Previously, I had failed to consider (or test) the case where lowering values
for stream/future writes would require intermediate allocations that need to be
cleaned up.  That manifested as compiler errors due to generated code
referencing a non-existent `cleanup_list` variable.

Fixes #1153

Signed-off-by: Joel Dice <[email protected]>

Author: dicej

crates/rust/src/bindgen.rs

@@ -54,7 +54,22 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
         }
     }
 
-    fn emit_cleanup(&mut self) {
+    pub(crate) fn flush_cleanup(&mut self) {
+        if !self.cleanup.is_empty() {
+            self.needs_cleanup_list = true;
+            self.push_str("cleanup_list.extend_from_slice(&[");
+            for (ptr, layout) in mem::take(&mut self.cleanup) {
+                self.push_str("(");
+                self.push_str(&ptr);
+                self.push_str(", ");
+                self.push_str(&layout);
+                self.push_str("),");
+            }
+            self.push_str("]);\n");
+        }
+    }
+
+    pub(crate) fn emit_cleanup(&mut self) {
         if self.emitted_cleanup {
             return;
         }
@@ -244,18 +259,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
     }
 
     fn finish_block(&mut self, operands: &mut Vec<String>) {
-        if !self.cleanup.is_empty() {
-            self.needs_cleanup_list = true;
-            self.push_str("cleanup_list.extend_from_slice(&[");
-            for (ptr, layout) in mem::take(&mut self.cleanup) {
-                self.push_str("(");
-                self.push_str(&ptr);
-                self.push_str(", ");
-                self.push_str(&layout);
-                self.push_str("),");
-            }
-            self.push_str("]);\n");
-        }
+        self.flush_cleanup();
         let (prev_src, prev_cleanup) = self.block_storage.pop().unwrap();
         let src = mem::replace(&mut self.src, prev_src);
         self.cleanup = prev_cleanup;

crates/rust/src/interface.rs

@@ -535,34 +535,40 @@ macro_rules! {macro_name} {{
                         };
                         let size = size.size_wasm32();
                         let align = align.align_wasm32();
-                        let (lower, lift) = if let Some(payload_type) = payload_type {
-                            let lower =
+                        let (lower, cleanup, lift) = if let Some(payload_type) = payload_type {
+                            let (lower, cleanup) =
                                 self.lower_to_memory("address", "&value", &payload_type, &module);
                             let lift =
                                 self.lift_from_memory("address", "value", &payload_type, &module);
-                            (lower, lift)
+                            (lower, cleanup, lift)
                         } else {
-                            (String::new(), "let value = ();\n".into())
+                            (String::new(), None, "let value = ();\n".into())
                         };
 
+                        let (cleanup_start, cleanup_end) =
+                            cleanup.unwrap_or_else(|| (String::new(), String::new()));
+
                         let box_ = self.path_to_box();
                         let code = format!(
                             r#"
 #[doc(hidden)]
 pub mod vtable{ordinal} {{
     fn write(future: u32, value: {name}) -> ::core::pin::Pin<{box_}<dyn ::core::future::Future<Output = bool>>> {{
         {box_}::pin(async move {{
+            {cleanup_start}
             #[repr(align({align}))]
             struct Buffer([::core::mem::MaybeUninit::<u8>; {size}]);
             let mut buffer = Buffer([::core::mem::MaybeUninit::uninit(); {size}]);
             let address = buffer.0.as_mut_ptr() as *mut u8;
             {lower}
 
-            match unsafe {{ {async_support}::await_future_result(start_write, future, address).await }} {{
+            let result = match unsafe {{ {async_support}::await_future_result(start_write, future, address).await }} {{
                 {async_support}::AsyncWaitResult::Values(_) => true,
                 {async_support}::AsyncWaitResult::End => false,
                 {async_support}::AsyncWaitResult::Error(_) => unreachable!("received error while performing write"),
-            }}
+            }};
+            {cleanup_end}
+            result
         }})
     }}
 
@@ -660,14 +666,15 @@ pub mod vtable{ordinal} {{
                         let size = size.size_wasm32();
                         let align = align.align_wasm32();
                         let alloc = self.path_to_std_alloc_module();
-                        let (lower_address, lower, lift_address, lift) = match payload_type {
+                        let (lower_address, lower, cleanup, lift_address, lift) = match payload_type
+                        {
                             Some(payload_type) if !stream_direct(payload_type) => {
                                 let address = format!(
                                     "let address = unsafe {{ {alloc}::alloc\
                                      ({alloc}::Layout::from_size_align_unchecked\
                                      ({size} * values.len(), {align})) }};"
                                 );
-                                let lower = self.lower_to_memory(
+                                let (lower, cleanup) = self.lower_to_memory(
                                     "address",
                                     "value",
                                     &payload_type,
@@ -696,7 +703,7 @@ for (index, dst) in values.iter_mut().take(count).enumerate() {{
 }}
                                 "#
                                 );
-                                (address.clone(), lower, address, lift)
+                                (address.clone(), lower, cleanup, address, lift)
                             }
                             _ => {
                                 let lower_address =
@@ -706,25 +713,29 @@ for (index, dst) in values.iter_mut().take(count).enumerate() {{
                                 (
                                     lower_address,
                                     String::new(),
+                                    None,
                                     lift_address,
                                     "let value = ();\n".into(),
                                 )
                             }
                         };
 
+                        let (cleanup_start, cleanup_end) =
+                            cleanup.unwrap_or_else(|| (String::new(), String::new()));
+
                         let box_ = self.path_to_box();
                         let code = format!(
                             r#"
 #[doc(hidden)]
 pub mod vtable{ordinal} {{
     fn write(stream: u32, values: &[{name}]) -> ::core::pin::Pin<{box_}<dyn ::core::future::Future<Output = usize> + '_>> {{
         {box_}::pin(async move {{
+            {cleanup_start}
             {lower_address}
             {lower}
 
             let mut total = 0;
             while total < values.len() {{
-
                 match unsafe {{
                     {async_support}::await_stream_result(
                         start_write,
@@ -738,6 +749,7 @@ pub mod vtable{ordinal} {{
                     {async_support}::AsyncWaitResult::End => break,
                 }}
             }}
+            {cleanup_end}
             total
         }})
     }}
@@ -867,10 +879,30 @@ pub mod vtable{ordinal} {{
         }
     }
 
-    fn lower_to_memory(&mut self, address: &str, value: &str, ty: &Type, module: &str) -> String {
+    fn lower_to_memory(
+        &mut self,
+        address: &str,
+        value: &str,
+        ty: &Type,
+        module: &str,
+    ) -> (String, Option<(String, String)>) {
         let mut f = FunctionBindgen::new(self, Vec::new(), true, module, true);
         abi::lower_to_memory(f.r#gen.resolve, &mut f, address.into(), value.into(), ty);
-        format!("unsafe {{ {} }}", String::from(f.src))
+        f.flush_cleanup();
+        let lower = format!("unsafe {{ {} }}", String::from(f.src));
+        let cleanup = if f.needs_cleanup_list {
+            f.src = Default::default();
+            f.emit_cleanup();
+            let body = String::from(f.src);
+            let vec = self.path_to_vec();
+            Some((
+                format!("let mut cleanup_list = {vec}::new();\n"),
+                format!("unsafe {{ {body} }}"),
+            ))
+        } else {
+            None
+        };
+        (lower, cleanup)
     }
 
     fn lift_from_memory(&mut self, address: &str, value: &str, ty: &Type, module: &str) -> String {

tests/codegen/futures.wit

@@ -27,6 +27,10 @@ interface futures {
   future-f32-ret: func() -> future<f32>;
   future-f64-ret: func() -> future<f64>;
 
+  future-result-list-string-ret: func() -> future<result<list<string>>>;
+  future-result-list-list-u8-ret: func() -> future<result<list<list<u8>>>>;
+  future-list-list-list-u8-ret: func() -> future<list<list<list<u8>>>>;
+
   tuple-future: func(x: future<tuple<u8, s8>>) -> future<tuple<s64, u32>>;
   string-future-arg: func(a: future<string>);
   string-future-ret: func() -> future<string>;

tests/codegen/streams.wit

@@ -41,6 +41,10 @@ interface streams {
   stream-f32-ret: func() -> stream<f32>;
   stream-f64-ret: func() -> stream<f64>;
 
+  stream-result-list-string-ret: func() -> stream<result<list<string>>>;
+  stream-result-list-list-u8-ret: func() -> stream<result<list<list<u8>>>>;
+  stream-list-list-list-u8-ret: func() -> stream<list<list<list<u8>>>>;
+
   tuple-stream: func(x: stream<tuple<u8, s8>>) -> stream<tuple<s64, u32>>;
   string-stream-arg: func(a: stream<string>);
   string-stream-ret: func() -> stream<string>;