BridgeJS: Fix memory management of SwiftHeapObject

Plugins/BridgeJS/Sources/BridgeJSLink/BridgeJSLink.swift

@@ -10,12 +10,16 @@ import BridgeJSUtilities
 public struct BridgeJSLink {
     var skeletons: [BridgeJSSkeleton] = []
     let sharedMemory: Bool
+    /// Whether to track the lifetime of Swift objects.
+    ///
+    /// This is useful for debugging memory issues.
+    let enableLifetimeTracking: Bool = false
     private let namespaceBuilder = NamespaceBuilder()
     private let intrinsicRegistry = JSIntrinsicRegistry()
 
     public init(
         skeletons: [BridgeJSSkeleton] = [],
-        sharedMemory: Bool
+        sharedMemory: Bool = false
     ) {
         self.skeletons = skeletons
         self.sharedMemory = sharedMemory
@@ -50,37 +54,76 @@ public struct BridgeJSLink {
         }
         """
 
-    let swiftHeapObjectClassJs = """
-        const swiftHeapObjectFinalizationRegistry = (typeof FinalizationRegistry === "undefined") ? { register: () => {}, unregister: () => {} } : new FinalizationRegistry((state) => {
-            if (state.hasReleased) {
-                return;
-            }
-            state.hasReleased = true;
-            state.deinit(state.pointer);
-        });
+    let lifetimeTrackingClassJs = """
+        const TRACKING = {
+            wrap: (pointer, deinit, prototype, state) => {
+                console.log(JSON.stringify({ DEBUG: true, event: "WRP", class: prototype.constructor.name, state }));
+            },
+            release: (obj) => {
+                console.log(JSON.stringify({ DEBUG: true, event: "REL", class: obj.constructor.name, state: obj.__swiftHeapObjectState }));
+            },
+            finalization: (state) => {
+                console.log(JSON.stringify({ DEBUG: true, event: "FIN", state }));
+            }
+        };
+        """
 
-        /// Represents a Swift heap object like a class instance or an actor instance.
-        class SwiftHeapObject {
-            static __wrap(pointer, deinit, prototype) {
-                const obj = Object.create(prototype);
-                const state = { pointer, deinit, hasReleased: false };
-                obj.pointer = pointer;
-                obj.__swiftHeapObjectState = state;
-                swiftHeapObjectFinalizationRegistry.register(obj, state, state);
-                return obj;
-            }
-
-            release() {
-                const state = this.__swiftHeapObjectState;
+    var swiftHeapObjectClassJs: String {
+        var output = ""
+        if enableLifetimeTracking {
+            output += lifetimeTrackingClassJs + "\n"
+        }
+        output += """
+            const swiftHeapObjectFinalizationRegistry = (typeof FinalizationRegistry === "undefined") ? { register: () => {}, unregister: () => {} } : new FinalizationRegistry((state) => {
+
+            """
+        if enableLifetimeTracking {
+            output += "        TRACKING.finalization(state);\n"
+        }
+        output += """
                 if (state.hasReleased) {
                     return;
                 }
                 state.hasReleased = true;
-                swiftHeapObjectFinalizationRegistry.unregister(state);
                 state.deinit(state.pointer);
-            }
+            });
+
+            /// Represents a Swift heap object like a class instance or an actor instance.
+            class SwiftHeapObject {
+                static __wrap(pointer, deinit, prototype) {
+                    const obj = Object.create(prototype);
+                    const state = { pointer, deinit, hasReleased: false };
+
+            """
+        if enableLifetimeTracking {
+            output += "        TRACKING.wrap(pointer, deinit, prototype, state);\n"
         }
-        """
+        output += """
+                    obj.pointer = pointer;
+                    obj.__swiftHeapObjectState = state;
+                    swiftHeapObjectFinalizationRegistry.register(obj, state, state);
+                    return obj;
+                }
+
+                release() {
+
+            """
+        if enableLifetimeTracking {
+            output += "        TRACKING.release(this);\n"
+        }
+        output += """
+                    const state = this.__swiftHeapObjectState;
+                    if (state.hasReleased) {
+                        return;
+                    }
+                    state.hasReleased = true;
+                    swiftHeapObjectFinalizationRegistry.unregister(state);
+                    state.deinit(state.pointer);
+                }
+            }
+            """
+        return output
+    }
 
     fileprivate struct LinkData {
         var exportsLines: [String] = []

Sources/JavaScriptKit/BridgeJSIntrinsics.swift

@@ -124,7 +124,6 @@ internal func _swift_js_closure_unregister(_ id: Int32) {
 // - `func bridgeJSStackPush()`: push the value onto the return stack (used by _BridgedSwiftStackType for array elements)
 //
 // Optional types (ExportSwift only) additionally define:
-// - `func bridgeJSLowerParameterWithRetain()`: lower optional heap object with ownership transfer for escaping closures
 // - `func bridgeJSLiftReturnFromSideChannel()`: lift optional from side-channel storage for protocol property getters
 //
 // See JSGlueGen.swift in BridgeJSLink for JS-side lowering/lifting implementation.
@@ -467,9 +466,10 @@ public protocol _BridgedSwiftHeapObject: AnyObject, _BridgedSwiftStackType {}
 extension _BridgedSwiftHeapObject {
 
     // MARK: ImportTS
-    @_spi(BridgeJS) @_transparent public consuming func bridgeJSLowerParameter() -> UnsafeMutableRawPointer {
-        // For protocol parameters, we pass the unretained pointer since JS already has a reference
-        return Unmanaged.passUnretained(self).toOpaque()
+    @_spi(BridgeJS) @_transparent public func bridgeJSLowerParameter() -> UnsafeMutableRawPointer {
+        // Transfer ownership to JS for imported SwiftHeapObject parameters.
+        // JS side must eventually release (via release() or FinalizationRegistry).
+        return Unmanaged.passRetained(self).toOpaque()
     }
     @_spi(BridgeJS) @_transparent public static func bridgeJSLiftReturn(_ pointer: UnsafeMutableRawPointer) -> Self {
         // For protocol returns, take an unretained value since JS manages the lifetime
@@ -1489,11 +1489,10 @@ extension Optional where Wrapped: _BridgedSwiftHeapObject {
     // MARK: ExportSwift
     /// Lowers optional Swift heap object as (isSome, pointer) tuple for protocol parameters.
     ///
-    /// This method uses `passUnretained()` because the caller (JavaScript protocol implementation)
-    /// already owns the object and will not retain it. The pointer is only valid for the
-    /// duration of the call.
+    /// Transfer ownership to JavaScript for imported optional heap-object parameters; JS must
+    /// release via `release()` or finalizer.
     ///
-    /// - Returns: A tuple containing presence flag (0 for nil, 1 for some) and unretained pointer
+    /// - Returns: A tuple containing presence flag (0 for nil, 1 for some) and retained pointer
     @_spi(BridgeJS) @_transparent public consuming func bridgeJSLowerParameter() -> (
         isSome: Int32, pointer: UnsafeMutableRawPointer
     ) {
@@ -1505,25 +1504,6 @@ extension Optional where Wrapped: _BridgedSwiftHeapObject {
         }
     }
 
-    /// Lowers optional Swift heap object with ownership transfer for escaping closures.
-    ///
-    /// This method uses `passRetained()` to transfer ownership to JavaScript, ensuring the
-    /// object remains valid even if the JavaScript closure escapes and stores the parameter.
-    /// JavaScript must wrap the pointer with `__construct()` to create a managed reference
-    /// that will be cleaned up via FinalizationRegistry.
-    ///
-    /// - Returns: A tuple containing presence flag (0 for nil, 1 for some) and retained pointer
-    @_spi(BridgeJS) @_transparent public consuming func bridgeJSLowerParameterWithRetain() -> (
-        isSome: Int32, pointer: UnsafeMutableRawPointer
-    ) {
-        switch consume self {
-        case .none:
-            return (isSome: 0, pointer: UnsafeMutableRawPointer(bitPattern: 1)!)
-        case .some(let value):
-            return (isSome: 1, pointer: Unmanaged.passRetained(value).toOpaque())
-        }
-    }
-
     @_spi(BridgeJS) @_transparent public static func bridgeJSLiftReturn(_ pointer: UnsafeMutableRawPointer) -> Wrapped?
     {
         if pointer == UnsafeMutableRawPointer(bitPattern: 0) {
@@ -2008,12 +1988,6 @@ extension _BridgedAsOptional where Wrapped: _BridgedSwiftHeapObject {
         asOptional.bridgeJSLowerParameter()
     }
 
-    @_spi(BridgeJS) public consuming func bridgeJSLowerParameterWithRetain() -> (
-        isSome: Int32, pointer: UnsafeMutableRawPointer
-    ) {
-        asOptional.bridgeJSLowerParameterWithRetain()
-    }
-
     @_spi(BridgeJS) public static func bridgeJSLiftParameter(
         _ isSome: Int32,
         _ pointer: UnsafeMutableRawPointer

Tests/BridgeJSRuntimeTests/Generated/BridgeJS.swift

@@ -9190,6 +9190,45 @@ fileprivate func _bjs_Container_wrap_extern(_ pointer: UnsafeMutableRawPointer)
     return _bjs_Container_wrap_extern(pointer)
 }
 
+@_expose(wasm, "bjs_LeakCheck_init")
+@_cdecl("bjs_LeakCheck_init")
+public func _bjs_LeakCheck_init() -> UnsafeMutableRawPointer {
+    #if arch(wasm32)
+    let ret = LeakCheck()
+    return ret.bridgeJSLowerReturn()
+    #else
+    fatalError("Only available on WebAssembly")
+    #endif
+}
+
+@_expose(wasm, "bjs_LeakCheck_deinit")
+@_cdecl("bjs_LeakCheck_deinit")
+public func _bjs_LeakCheck_deinit(_ pointer: UnsafeMutableRawPointer) -> Void {
+    #if arch(wasm32)
+    Unmanaged<LeakCheck>.fromOpaque(pointer).release()
+    #else
+    fatalError("Only available on WebAssembly")
+    #endif
+}
+
+extension LeakCheck: ConvertibleToJSValue, _BridgedSwiftHeapObject {
+    public var jsValue: JSValue {
+        return .object(JSObject(id: UInt32(bitPattern: _bjs_LeakCheck_wrap(Unmanaged.passRetained(self).toOpaque()))))
+    }
+}
+
+#if arch(wasm32)
+@_extern(wasm, module: "BridgeJSRuntimeTests", name: "bjs_LeakCheck_wrap")
+fileprivate func _bjs_LeakCheck_wrap_extern(_ pointer: UnsafeMutableRawPointer) -> Int32
+#else
+fileprivate func _bjs_LeakCheck_wrap_extern(_ pointer: UnsafeMutableRawPointer) -> Int32 {
+    fatalError("Only available on WebAssembly")
+}
+#endif
+@inline(never) fileprivate func _bjs_LeakCheck_wrap(_ pointer: UnsafeMutableRawPointer) -> Int32 {
+    return _bjs_LeakCheck_wrap_extern(pointer)
+}
+
 #if arch(wasm32)
 @_extern(wasm, module: "BridgeJSRuntimeTests", name: "bjs_ClosureSupportImports_jsApplyVoid_static")
 fileprivate func bjs_ClosureSupportImports_jsApplyVoid_static_extern(_ callback: Int32) -> Void
@@ -11252,6 +11291,30 @@ fileprivate func bjs_SwiftClassSupportImports_jsRoundTripOptionalGreeter_static_
     return bjs_SwiftClassSupportImports_jsRoundTripOptionalGreeter_static_extern(greeterIsSome, greeterPointer)
 }
 
+#if arch(wasm32)
+@_extern(wasm, module: "BridgeJSRuntimeTests", name: "bjs_SwiftClassSupportImports_jsConsumeLeakCheck_static")
+fileprivate func bjs_SwiftClassSupportImports_jsConsumeLeakCheck_static_extern(_ value: UnsafeMutableRawPointer) -> Void
+#else
+fileprivate func bjs_SwiftClassSupportImports_jsConsumeLeakCheck_static_extern(_ value: UnsafeMutableRawPointer) -> Void {
+    fatalError("Only available on WebAssembly")
+}
+#endif
+@inline(never) fileprivate func bjs_SwiftClassSupportImports_jsConsumeLeakCheck_static(_ value: UnsafeMutableRawPointer) -> Void {
+    return bjs_SwiftClassSupportImports_jsConsumeLeakCheck_static_extern(value)
+}
+
+#if arch(wasm32)
+@_extern(wasm, module: "BridgeJSRuntimeTests", name: "bjs_SwiftClassSupportImports_jsConsumeOptionalLeakCheck_static")
+fileprivate func bjs_SwiftClassSupportImports_jsConsumeOptionalLeakCheck_static_extern(_ valueIsSome: Int32, _ valuePointer: UnsafeMutableRawPointer) -> Void
+#else
+fileprivate func bjs_SwiftClassSupportImports_jsConsumeOptionalLeakCheck_static_extern(_ valueIsSome: Int32, _ valuePointer: UnsafeMutableRawPointer) -> Void {
+    fatalError("Only available on WebAssembly")
+}
+#endif
+@inline(never) fileprivate func bjs_SwiftClassSupportImports_jsConsumeOptionalLeakCheck_static(_ valueIsSome: Int32, _ valuePointer: UnsafeMutableRawPointer) -> Void {
+    return bjs_SwiftClassSupportImports_jsConsumeOptionalLeakCheck_static_extern(valueIsSome, valuePointer)
+}
+
 func _$SwiftClassSupportImports_jsRoundTripGreeter(_ greeter: Greeter) throws(JSException) -> Greeter {
     let greeterPointer = greeter.bridgeJSLowerParameter()
     let ret = bjs_SwiftClassSupportImports_jsRoundTripGreeter_static(greeterPointer)
@@ -11268,4 +11331,20 @@ func _$SwiftClassSupportImports_jsRoundTripOptionalGreeter(_ greeter: Optional<G
         throw error
     }
     return Optional<Greeter>.bridgeJSLiftReturn(ret)
+}
+
+func _$SwiftClassSupportImports_jsConsumeLeakCheck(_ value: LeakCheck) throws(JSException) -> Void {
+    let valuePointer = value.bridgeJSLowerParameter()
+    bjs_SwiftClassSupportImports_jsConsumeLeakCheck_static(valuePointer)
+    if let error = _swift_js_take_exception() {
+        throw error
+    }
+}
+
+func _$SwiftClassSupportImports_jsConsumeOptionalLeakCheck(_ value: Optional<LeakCheck>) throws(JSException) -> Void {
+    let (valueIsSome, valuePointer) = value.bridgeJSLowerParameter()
+    bjs_SwiftClassSupportImports_jsConsumeOptionalLeakCheck_static(valueIsSome, valuePointer)
+    if let error = _swift_js_take_exception() {
+        throw error
+    }
 }

Tests/BridgeJSRuntimeTests/Generated/JavaScript/BridgeJS.json

@@ -3626,6 +3626,28 @@
           }
         ],
         "swiftCallName" : "Container"
+      },
+      {
+        "constructor" : {
+          "abiName" : "bjs_LeakCheck_init",
+          "effects" : {
+            "isAsync" : false,
+            "isStatic" : false,
+            "isThrows" : false
+          },
+          "parameters" : [
+
+          ]
+        },
+        "explicitAccessControl" : "public",
+        "methods" : [
+
+        ],
+        "name" : "LeakCheck",
+        "properties" : [
+
+        ],
+        "swiftCallName" : "LeakCheck"
       }
     ],
     "enums" : [
@@ -15643,6 +15665,47 @@
                     "_1" : "null"
                   }
                 }
+              },
+              {
+                "name" : "jsConsumeLeakCheck",
+                "parameters" : [
+                  {
+                    "name" : "value",
+                    "type" : {
+                      "swiftHeapObject" : {
+                        "_0" : "LeakCheck"
+                      }
+                    }
+                  }
+                ],
+                "returnType" : {
+                  "void" : {
+
+                  }
+                }
+              },
+              {
+                "name" : "jsConsumeOptionalLeakCheck",
+                "parameters" : [
+                  {
+                    "name" : "value",
+                    "type" : {
+                      "nullable" : {
+                        "_0" : {
+                          "swiftHeapObject" : {
+                            "_0" : "LeakCheck"
+                          }
+                        },
+                        "_1" : "null"
+                      }
+                    }
+                  }
+                ],
+                "returnType" : {
+                  "void" : {
+
+                  }
+                }
               }
             ]
           }

Tests/BridgeJSRuntimeTests/JavaScript/SwiftClassSupportTests.mjs

@@ -9,5 +9,14 @@ export function getImports(importsContext) {
         jsRoundTripOptionalGreeter: (greeter) => {
             return greeter;
         },
+        jsConsumeLeakCheck: (value) => {
+            // Explicitly release on JS side to mimic user cleanup.
+            value.release();
+        },
+        jsConsumeOptionalLeakCheck: (value) => {
+            if (value) {
+                value.release();
+            }
+        },
     };
-}
+}