fix: support runtime borrow drop tracking (#429)

crates/js-component-bindgen/src/function_bindgen.rs

@@ -68,7 +68,7 @@ pub type ResourceMap = BTreeMap<TypeId, ResourceTable>;
 
 pub struct FunctionBindgen<'a> {
     pub resource_map: &'a ResourceMap,
-    pub cur_resource_borrows: Vec<String>,
+    pub cur_resource_borrows: bool,
     pub intrinsics: &'a mut BTreeSet<Intrinsic>,
     pub valid_lifting_optimization: bool,
     pub sizes: &'a SizeAlign,
@@ -1119,11 +1119,35 @@ impl Bindgen for FunctionBindgen<'_> {
                 }
 
                 // After a high level call, we need to deactivate the component resource borrows.
-                if !self.cur_resource_borrows.is_empty() {
-                    for resource_borrow in &self.cur_resource_borrows {
-                        uwriteln!(self.src, "{}", resource_borrow);
+                if self.cur_resource_borrows {
+                    let symbol_resource_handle = self.intrinsic(Intrinsic::SymbolResourceHandle);
+                    let cur_resource_borrows = self.intrinsic(Intrinsic::CurResourceBorrows);
+                    let host = matches!(
+                        self.resource_map.iter().nth(0).unwrap().1.data,
+                        ResourceData::Host { .. }
+                    );
+                    if host {
+                        uwriteln!(
+                            self.src,
+                            "for (const rsc of {cur_resource_borrows}) {{
+                                rsc[{symbol_resource_handle}] = null;
+                            }}
+                            {cur_resource_borrows} = [];"
+                        );
+                    } else {
+                        uwriteln!(
+                            self.src,
+                            "for (const {{ rsc, drop }} of {cur_resource_borrows}) {{
+                                if (rsc[{symbol_resource_handle}]) {{
+                                    drop(rsc[{symbol_resource_handle}]);
+                                    delete rsc[{symbol_resource_handle}];
+                                }}
+                                {cur_resource_borrows}[i][{symbol_resource_handle}] = null;
+                            }}
+                            {cur_resource_borrows} = [];"
+                        );
                     }
-                    self.cur_resource_borrows = Vec::new();
+                    self.cur_resource_borrows = false;
                 }
             }
 
@@ -1237,13 +1261,6 @@ impl Bindgen for FunctionBindgen<'_> {
                             } else {
                                 // Borrow handles of local resources have rep handles, which we carry through here.
                                 uwriteln!(self.src, "Object.defineProperty({rsc}, {symbol_resource_handle}, {{ writable: true, value: {handle} }});");
-                                // Borrow handles are tracked to release after the call by CallInterface.
-                                // Once may_enter and may_leave are properly implemented, we can be certain to avoid non-reentrancy of import calls
-                                // calling export calls with borrow handles, so can maintain the invariant that borrows cannot ever be passed in.
-                                self.cur_resource_borrows.push(format!(
-                                    "{}[{symbol_resource_handle}] = null;",
-                                    rsc.to_string()
-                                ));
                             }
                         } else {
                             let rep = format!("rep{}", self.tmp());
@@ -1270,14 +1287,16 @@ impl Bindgen for FunctionBindgen<'_> {
                                     }}
                                     {rsc_table_remove}(handleTable{tid}, {handle});"
                                 );
-                            } else {
-                                // If lifting a borrow, that was not previously captured, create the class.
-                                self.cur_resource_borrows.push(format!(
-                                    "{}[{symbol_resource_handle}] = null;",
-                                    rsc.to_string()
-                                ));
                             }
                         }
+
+                        // Borrow handles are tracked to release after the call by CallInterface.
+                        if !is_own {
+                            let cur_resource_borrows =
+                                self.intrinsic(Intrinsic::CurResourceBorrows);
+                            uwriteln!(self.src, "{cur_resource_borrows}.push({rsc});");
+                            self.cur_resource_borrows = true;
+                        }
                     }
 
                     ResourceData::Guest {
@@ -1318,11 +1337,10 @@ impl Bindgen for FunctionBindgen<'_> {
                             );
 
                             if !is_own {
-                                self.cur_resource_borrows.push(format!(
-                                    "if ({rsc}[{symbol_resource_handle}]) {{
-                                        $resource_import${prefix}drop${lower_camel}({rsc}[{symbol_resource_handle}]);
-                                        delete {rsc}[{symbol_resource_handle}];
-                                    }}"));
+                                let cur_resource_borrows =
+                                    self.intrinsic(Intrinsic::CurResourceBorrows);
+                                uwriteln!(self.src, "{cur_resource_borrows}.push({{ rsc: {rsc}, drop: $resource_import${prefix}drop${lower_camel} }});");
+                                self.cur_resource_borrows = true;
                             }
                         }
                     }

crates/js-component-bindgen/src/intrinsics.rs

@@ -8,6 +8,7 @@ pub enum Intrinsic {
     Base64Compile,
     ClampGuest,
     ComponentError,
+    CurResourceBorrows,
     DataView,
     EmptyFunc,
     F32ToI32,
@@ -135,6 +136,10 @@ pub fn render_intrinsics(
                 }
             "),
 
+            Intrinsic::CurResourceBorrows => output.push_str("
+                let curResourceBorrows = [];
+            "),
+
             Intrinsic::DataView => output.push_str("
                 let dv = new DataView(new ArrayBuffer());
                 const dataView = mem => dv.buffer === mem.buffer ? dv : dv = new DataView(mem.buffer);
@@ -588,6 +593,7 @@ impl Intrinsic {
             "base64Compile",
             "clampGuest",
             "ComponentError",
+            "curResourceBorrows",
             "dataView",
             "emptyFunc",
             "f32ToI32",
@@ -603,32 +609,32 @@ impl Intrinsic {
             "instantiateCore",
             "isLE",
             "resourceCallBorrows",
-            "T_FLAG",
+            "resourceTransferBorrow",
+            "resourceTransferBorrowValidLifting",
+            "resourceTransferOwn",
             "rscTableCreateBorrow",
             "rscTableCreateOwn",
             "rscTableGet",
-            "rscTableTryGet",
             "rscTableRemove",
-            "resourceTransferBorrow",
-            "resourceTransferBorrowValidLifting",
-            "resourceTransferOwn",
+            "rscTableTryGet",
             "scopeId",
             "symbolCabiDispose",
             "symbolCabiLower",
             "symbolDispose",
             "symbolRscHandle",
             "symbolRscRep",
+            "T_FLAG",
             "throwInvalidBool",
             "throwUninitialized",
-            "toInt64",
-            "toUint64",
             "toInt16",
             "toInt32",
+            "toInt64",
             "toInt8",
             "toResultString",
             "toString",
             "toUint16",
             "toUint32",
+            "toUint64",
             "toUint8",
             "utf16Decoder",
             "utf16Encode",
@@ -665,6 +671,7 @@ impl Intrinsic {
             Intrinsic::Base64Compile => "base64Compile",
             Intrinsic::ClampGuest => "clampGuest",
             Intrinsic::ComponentError => "ComponentError",
+            Intrinsic::CurResourceBorrows => "curResourceBorrows",
             Intrinsic::DataView => "dataView",
             Intrinsic::EmptyFunc => "emptyFunc",
             Intrinsic::F32ToI32 => "f32ToI32",

crates/js-component-bindgen/src/transpile_bindgen.rs

@@ -1496,7 +1496,7 @@ impl<'a> Instantiator<'a, '_> {
 
         let mut f = FunctionBindgen {
             resource_map: &resource_map,
-            cur_resource_borrows: Vec::new(),
+            cur_resource_borrows: false,
             intrinsics: &mut self.gen.all_intrinsics,
             valid_lifting_optimization: self.gen.opts.valid_lifting_optimization,
             sizes: &self.sizes,