c: Use drop-flags when auto-dropping borrows (#809)

* Use drop-flags when auto-dropping borrows

* Fix codegen errors in the Rust generator

---------

Co-authored-by: Alex Crichton <[email protected]>

Author: elliottt

crates/c/src/lib.rs

@@ -1967,6 +1967,11 @@ impl InterfaceGenerator<'_> {
     }
 }
 
+struct DroppableBorrow {
+    name: String,
+    ty: TypeId,
+}
+
 struct FunctionBindgen<'a, 'b> {
     gen: &'a mut InterfaceGenerator<'b>,
     locals: Ns,
@@ -1984,7 +1989,10 @@ struct FunctionBindgen<'a, 'b> {
 
     /// Borrows observed during lifting an export, that will need to be dropped when the guest
     /// function exits.
-    borrows: Vec<(String, TypeId)>,
+    borrows: Vec<DroppableBorrow>,
+
+    /// Forward declarations for temporary storage of borrow copies.
+    borrow_decls: wit_bindgen_core::Source,
 }
 
 impl<'a, 'b> FunctionBindgen<'a, 'b> {
@@ -2007,6 +2015,7 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
             ret_store_cnt: 0,
             import_return_pointer_area_size: 0,
             import_return_pointer_area_align: 0,
+            borrow_decls: Default::default(),
             borrows: Vec::new(),
         }
     }
@@ -2264,8 +2273,13 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                         if !self.gen.in_import && self.gen.autodrop_enabled() {
                             // Here we've received a borrow of an imported resource, which is the
                             // kind we'll need to drop when the exported function is returning.
-                            let target = dealias(self.gen.resolve, *id);
-                            self.borrows.push((op.clone(), target));
+                            let ty = dealias(self.gen.resolve, *id);
+
+                            let name = self.locals.tmp("borrow");
+                            uwriteln!(self.borrow_decls, "int32_t {name} = 0;");
+                            uwriteln!(self.src, "{name} = {op};");
+
+                            self.borrows.push(DroppableBorrow { name, ty });
                         }
                     }
                 }
@@ -2372,10 +2386,6 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 for (i, (case, (block, block_results))) in
                     variant.cases.iter().zip(blocks).enumerate()
                 {
-                    if let Some(ty) = case.ty.as_ref() {
-                        self.assert_no_droppable_borrows("variant", ty);
-                    }
-
                     uwriteln!(self.src, "case {}: {{", i);
                     self.src.push_str(&block);
                     assert!(block_results.len() == (case.ty.is_some() as usize));
@@ -2431,8 +2441,6 @@ impl Bindgen for FunctionBindgen<'_, '_> {
             }
 
             Instruction::OptionLift { ty, .. } => {
-                self.assert_no_droppable_borrows("option", &Type::Id(*ty));
-
                 let (mut some, some_results) = self.blocks.pop().unwrap();
                 let (mut none, none_results) = self.blocks.pop().unwrap();
                 assert!(none_results.len() == 0);
@@ -2520,8 +2528,6 @@ impl Bindgen for FunctionBindgen<'_, '_> {
             }
 
             Instruction::ResultLift { result, ty, .. } => {
-                self.assert_no_droppable_borrows("result", &Type::Id(*ty));
-
                 let (mut err, err_results) = self.blocks.pop().unwrap();
                 assert!(err_results.len() == (result.err.is_some() as usize));
                 let (mut ok, ok_results) = self.blocks.pop().unwrap();
@@ -2837,8 +2843,15 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 }
             },
             Instruction::Return { amt, .. } => {
-                for (op, ty) in self.borrows.iter() {
-                    uwriteln!(self.src, "{}({op});", self.gen.gen.resources[ty].drop_fn);
+                // Emit all temporary borrow decls
+                let src = std::mem::replace(&mut self.src, std::mem::take(&mut self.borrow_decls));
+                self.src.append_src(&src);
+
+                for DroppableBorrow { name, ty } in self.borrows.iter() {
+                    let drop_fn = self.gen.gen.resources[ty].drop_fn.as_str();
+                    uwriteln!(self.src, "if ({name} != 0) {{");
+                    uwriteln!(self.src, "  {drop_fn}({name});");
+                    uwriteln!(self.src, "}}");
                 }
 
                 assert!(*amt <= 1);

crates/rust/src/bindgen.rs

@@ -16,6 +16,7 @@ pub(super) struct FunctionBindgen<'a, 'b> {
     cleanup: Vec<(String, String)>,
     pub import_return_pointer_area_size: usize,
     pub import_return_pointer_area_align: usize,
+    pub handle_decls: Vec<String>,
 }
 
 impl<'a, 'b> FunctionBindgen<'a, 'b> {
@@ -34,6 +35,7 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
             cleanup: Vec::new(),
             import_return_pointer_area_size: 0,
             import_return_pointer_area_align: 0,
+            handle_decls: Vec::new(),
         }
     }
 
@@ -451,9 +453,19 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                                 format!("{name}::from_handle({op} as u32)")
                             }
                         }
+                    } else if prefix == "" {
+                        let name = self.gen.type_path(resource, true);
+                        format!("{name}::from_handle({op} as u32)")
                     } else {
+                        let tmp = format!("handle{}", self.tmp());
+                        self.handle_decls.push(format!("let {tmp};"));
                         let name = self.gen.type_path(resource, true);
-                        format!("{prefix}{name}::from_handle({op} as u32)")
+                        format!(
+                            "{{\n
+                                {tmp} = {name}::from_handle({op} as u32);
+                                {prefix}{tmp}
+                            }}"
+                        )
                     },
                 );
             }

crates/rust/src/interface.rs

@@ -284,12 +284,14 @@ impl InterfaceGenerator<'_> {
             src,
             import_return_pointer_area_size,
             import_return_pointer_area_align,
+            handle_decls,
             ..
         } = f;
 
         if needs_cleanup_list {
             self.src.push_str("let mut cleanup_list = Vec::new();\n");
         }
+        assert!(handle_decls.is_empty());
         if import_return_pointer_area_size > 0 {
             uwrite!(
                 self.src,
@@ -389,9 +391,14 @@ impl InterfaceGenerator<'_> {
         let FunctionBindgen {
             needs_cleanup_list,
             src,
+            handle_decls,
             ..
         } = f;
         assert!(!needs_cleanup_list);
+        for decl in handle_decls {
+            self.src.push_str(&decl);
+            self.src.push_str("\n");
+        }
         self.src.push_str(&String::from(src));
         self.src.push_str("}\n");
 
@@ -420,9 +427,11 @@ impl InterfaceGenerator<'_> {
             let FunctionBindgen {
                 needs_cleanup_list,
                 src,
+                handle_decls,
                 ..
             } = f;
             assert!(!needs_cleanup_list);
+            assert!(handle_decls.is_empty());
             self.src.push_str(&String::from(src));
             self.src.push_str("}\n");
             self.src.push_str("};\n");

tests/codegen/resources.wit

@@ -11,6 +11,31 @@ world resources {
 
   export add: func(a: borrow<z>, b: borrow<z>) -> own<z>;
 
+  export maybe-with-z: func(a: option<borrow<z>>);
+
+  variant includes-borrow {
+    a,
+    b(borrow<z>),
+  }
+
+  export variant-with-z: func(a: includes-borrow);
+  export maybe-variant-with-z: func(a: option<includes-borrow>);
+
+  record big {
+    x1: borrow<z>,
+    x2: borrow<z>,
+    x3: borrow<z>,
+    x4: borrow<z>,
+    x5: borrow<z>,
+    x6: borrow<z>,
+    x7: borrow<z>,
+    x8: borrow<z>,
+    x9: borrow<z>,
+    x10: borrow<z>,
+  }
+
+  export big-record: func(r: big);
+
   export exports: interface {
     resource x {
       constructor(a: float64);