feat: Add Swift-side per-class identity cache with shared memory flag

Replace per-crossing deinit(pointer) WASM call on identity cache hits with
a Swift-side per-class Set that tracks exported pointers. On cache hit, Swift
uses passUnretained (no retain) and writes a flag to WASM linear memory. JS
reads the flag and skips the deinit call.

On cache miss, normal passRetained path. On FinalizationRegistry race (stale
JS cache but Swift thinks object is known), JS calls bjs_identity_retain to
recover.

Trade-off vs JS-only cache: ~40% slower on roundtrip (Set.contains overhead
on WASM due to SipHash), ~15% faster on create-heavy paths (fewer retain/
release cycles).

Author: krodak

Benchmarks/Sources/Generated/BridgeJS.swift

@@ -1531,12 +1531,24 @@ fileprivate func _bjs_IdentityCacheBenchmark_wrap_extern(_ pointer: UnsafeMutabl
     return _bjs_IdentityCacheBenchmark_wrap_extern(pointer)
 }
 
+nonisolated(unsafe) var _SimpleClassIdentity_identityExported: Set<UnsafeMutableRawPointer> = []
+
 @_expose(wasm, "bjs_SimpleClassIdentity_init")
 @_cdecl("bjs_SimpleClassIdentity_init")
 public func _bjs_SimpleClassIdentity_init(_ nameBytes: Int32, _ nameLength: Int32, _ count: Int32, _ flag: Int32, _ rate: Float32, _ precise: Float64) -> UnsafeMutableRawPointer {
     #if arch(wasm32)
     let ret = SimpleClassIdentity(name: String.bridgeJSLiftParameter(nameBytes, nameLength), count: Int.bridgeJSLiftParameter(count), flag: Bool.bridgeJSLiftParameter(flag), rate: Float.bridgeJSLiftParameter(rate), precise: Double.bridgeJSLiftParameter(precise))
-    return ret.bridgeJSLowerReturn()
+    return withExtendedLifetime(ret) {
+        let pointer = Unmanaged.passUnretained(ret).toOpaque()
+        if _SimpleClassIdentity_identityExported.contains(pointer) {
+            _bridgeJSIdentityFlag = 1
+            return pointer
+        }
+        _SimpleClassIdentity_identityExported.insert(pointer)
+        _ = Unmanaged.passRetained(ret)
+        _bridgeJSIdentityFlag = 0
+        return pointer
+    }
     #else
     fatalError("Only available on WebAssembly")
     #endif
@@ -1651,6 +1663,7 @@ public func _bjs_SimpleClassIdentity_precise_set(_ _self: UnsafeMutableRawPointe
 @_cdecl("bjs_SimpleClassIdentity_deinit")
 public func _bjs_SimpleClassIdentity_deinit(_ pointer: UnsafeMutableRawPointer) -> Void {
     #if arch(wasm32)
+    _SimpleClassIdentity_identityExported.remove(pointer)
     Unmanaged<SimpleClassIdentity>.fromOpaque(pointer).release()
     #else
     fatalError("Only available on WebAssembly")
@@ -1678,12 +1691,24 @@ fileprivate func _bjs_SimpleClassIdentity_wrap_extern(_ pointer: UnsafeMutableRa
     return _bjs_SimpleClassIdentity_wrap_extern(pointer)
 }
 
+nonisolated(unsafe) var _ClassRoundtripIdentity_identityExported: Set<UnsafeMutableRawPointer> = []
+
 @_expose(wasm, "bjs_ClassRoundtripIdentity_init")
 @_cdecl("bjs_ClassRoundtripIdentity_init")
 public func _bjs_ClassRoundtripIdentity_init() -> UnsafeMutableRawPointer {
     #if arch(wasm32)
     let ret = ClassRoundtripIdentity()
-    return ret.bridgeJSLowerReturn()
+    return withExtendedLifetime(ret) {
+        let pointer = Unmanaged.passUnretained(ret).toOpaque()
+        if _ClassRoundtripIdentity_identityExported.contains(pointer) {
+            _bridgeJSIdentityFlag = 1
+            return pointer
+        }
+        _ClassRoundtripIdentity_identityExported.insert(pointer)
+        _ = Unmanaged.passRetained(ret)
+        _bridgeJSIdentityFlag = 0
+        return pointer
+    }
     #else
     fatalError("Only available on WebAssembly")
     #endif
@@ -1725,6 +1750,7 @@ public func _bjs_ClassRoundtripIdentity_takeSimpleClassIdentity(_ _self: UnsafeM
 @_cdecl("bjs_ClassRoundtripIdentity_deinit")
 public func _bjs_ClassRoundtripIdentity_deinit(_ pointer: UnsafeMutableRawPointer) -> Void {
     #if arch(wasm32)
+    _ClassRoundtripIdentity_identityExported.remove(pointer)
     Unmanaged<ClassRoundtripIdentity>.fromOpaque(pointer).release()
     #else
     fatalError("Only available on WebAssembly")
@@ -1752,12 +1778,24 @@ fileprivate func _bjs_ClassRoundtripIdentity_wrap_extern(_ pointer: UnsafeMutabl
     return _bjs_ClassRoundtripIdentity_wrap_extern(pointer)
 }
 
+nonisolated(unsafe) var _IdentityCacheBenchmarkIdentity_identityExported: Set<UnsafeMutableRawPointer> = []
+
 @_expose(wasm, "bjs_IdentityCacheBenchmarkIdentity_init")
 @_cdecl("bjs_IdentityCacheBenchmarkIdentity_init")
 public func _bjs_IdentityCacheBenchmarkIdentity_init() -> UnsafeMutableRawPointer {
     #if arch(wasm32)
     let ret = IdentityCacheBenchmarkIdentity()
-    return ret.bridgeJSLowerReturn()
+    return withExtendedLifetime(ret) {
+        let pointer = Unmanaged.passUnretained(ret).toOpaque()
+        if _IdentityCacheBenchmarkIdentity_identityExported.contains(pointer) {
+            _bridgeJSIdentityFlag = 1
+            return pointer
+        }
+        _IdentityCacheBenchmarkIdentity_identityExported.insert(pointer)
+        _ = Unmanaged.passRetained(ret)
+        _bridgeJSIdentityFlag = 0
+        return pointer
+    }
     #else
     fatalError("Only available on WebAssembly")
     #endif
@@ -1788,6 +1826,7 @@ public func _bjs_IdentityCacheBenchmarkIdentity_getPoolRepeated(_ _self: UnsafeM
 @_cdecl("bjs_IdentityCacheBenchmarkIdentity_deinit")
 public func _bjs_IdentityCacheBenchmarkIdentity_deinit(_ pointer: UnsafeMutableRawPointer) -> Void {
     #if arch(wasm32)
+    _IdentityCacheBenchmarkIdentity_identityExported.remove(pointer)
     Unmanaged<IdentityCacheBenchmarkIdentity>.fromOpaque(pointer).release()
     #else
     fatalError("Only available on WebAssembly")

Plugins/BridgeJS/Sources/BridgeJSCore/ExportSwift.swift

@@ -31,6 +31,14 @@ public class ExportSwift {
         self.skeleton = skeleton
     }
 
+    /// Whether a class should use identity caching based on its annotation and the config default.
+    private func shouldUseIdentityCache(for klass: ExportedClass) -> Bool {
+        if let classOverride = klass.identityMode {
+            return classOverride
+        }
+        return skeleton.identityMode == "pointer"
+    }
+
     /// Finalizes the export process and generates the bridge code
     ///
     /// - Parameters:
@@ -103,6 +111,9 @@ public class ExportSwift {
         var abiReturnType: WasmCoreType?
         var externDecls: [DeclSyntax] = []
         let effects: Effects
+        /// When set, heap object returns use the per-class identity Set pattern
+        /// instead of the default `bridgeJSLowerReturn()`.
+        var identityClassName: String?
 
         init(effects: Effects) {
             self.effects = effects
@@ -323,6 +334,26 @@ public class ExportSwift {
                     }
                     """
                 )
+            case .swiftHeapObject(let className):
+                if let identityClass = identityClassName, identityClass == className {
+                    append(
+                        """
+                        return withExtendedLifetime(ret) {
+                            let pointer = Unmanaged.passUnretained(ret).toOpaque()
+                            if _\(raw: identityClass)_identityExported.contains(pointer) {
+                                _bridgeJSIdentityFlag = 1
+                                return pointer
+                            }
+                            _\(raw: identityClass)_identityExported.insert(pointer)
+                            _ = Unmanaged.passRetained(ret)
+                            _bridgeJSIdentityFlag = 0
+                            return pointer
+                        }
+                        """
+                    )
+                } else {
+                    append("return ret.bridgeJSLowerReturn()")
+                }
             default:
                 append("return ret.bridgeJSLowerReturn()")
             }
@@ -534,9 +565,11 @@ public class ExportSwift {
     private func renderSingleExportedConstructor(
         constructor: ExportedConstructor,
         callName: String,
-        returnType: BridgeType
+        returnType: BridgeType,
+        identityClassName: String? = nil
     ) throws -> DeclSyntax {
         let builder = ExportedThunkBuilder(effects: constructor.effects)
+        builder.identityClassName = identityClassName
         for param in constructor.parameters {
             try builder.liftParameter(param: param)
         }
@@ -548,9 +581,11 @@ public class ExportSwift {
     private func renderSingleExportedMethod(
         method: ExportedFunction,
         ownerTypeName: String,
-        instanceSelfType: BridgeType
+        instanceSelfType: BridgeType,
+        identityClassName: String? = nil
     ) throws -> DeclSyntax {
         let builder = ExportedThunkBuilder(effects: method.effects)
+        builder.identityClassName = identityClassName
         if !method.effects.isStatic {
             try builder.liftParameter(param: Parameter(label: nil, name: "_self", type: instanceSelfType))
         }
@@ -652,21 +687,31 @@ public class ExportSwift {
     func renderSingleExportedClass(klass: ExportedClass) throws -> [DeclSyntax] {
         var decls: [DeclSyntax] = []
 
-        if let constructor = klass.constructor {
+        let useIdentity = shouldUseIdentityCache(for: klass)
+
+        // For identity-mode classes, emit the per-class Set
+        if useIdentity {
             decls.append(
-                try renderSingleExportedConstructor(
-                    constructor: constructor,
-                    callName: klass.swiftCallName,
-                    returnType: .swiftHeapObject(klass.name)
-                )
+                "nonisolated(unsafe) var _\(raw: klass.name)_identityExported: Set<UnsafeMutableRawPointer> = []"
             )
         }
+
+        if let constructor = klass.constructor {
+            let constructorDecl = try renderSingleExportedConstructor(
+                constructor: constructor,
+                callName: klass.swiftCallName,
+                returnType: .swiftHeapObject(klass.name),
+                identityClassName: useIdentity ? klass.name : nil
+            )
+            decls.append(constructorDecl)
+        }
         for method in klass.methods {
             decls.append(
                 try renderSingleExportedMethod(
                     method: method,
                     ownerTypeName: klass.swiftCallName,
-                    instanceSelfType: .swiftHeapObject(klass.swiftCallName)
+                    instanceSelfType: .swiftHeapObject(klass.swiftCallName),
+                    identityClassName: useIdentity ? klass.name : nil
                 )
             )
         }
@@ -686,6 +731,9 @@ public class ExportSwift {
                     returnType: nil
                 )
             ) { printer in
+                if useIdentity {
+                    printer.write("_\(klass.name)_identityExported.remove(pointer)")
+                }
                 printer.write("Unmanaged<\(klass.swiftCallName)>.fromOpaque(pointer).release()")
             }
             decls.append(DeclSyntax(funcDecl))

Plugins/BridgeJS/Sources/BridgeJSLink/BridgeJSLink.swift

@@ -83,6 +83,13 @@ public struct BridgeJSLink {
         };
         """
 
+    /// Whether any class across all skeletons uses identity caching.
+    private var hasAnyIdentityClass: Bool {
+        skeletons.compactMap(\.exported).contains { skeleton in
+            skeleton.classes.contains { shouldUseIdentityCache(for: $0) }
+        }
+    }
+
     var swiftHeapObjectClassJs: String {
         var output = ""
         if enableLifetimeTracking {
@@ -131,13 +138,40 @@ public struct BridgeJSLink {
 
                     const cached = identityCache.get(pointer)?.deref();
                     if (cached && !cached.__swiftHeapObjectState.hasReleased) {
-                        deinit(pointer);
+
+            """
+        if hasAnyIdentityClass {
+            output += """
+                            const swiftRetained = __bjs_readIdentityFlag() === 0;
+                            if (swiftRetained) {
+                                deinit(pointer);
+                            }
+
+                """
+        } else {
+            output += "                        deinit(pointer);\n"
+        }
+        output += """
                         return cached;
                     }
                     if (identityCache.has(pointer)) {
                         identityCache.delete(pointer);
                     }
 
+            """
+        if hasAnyIdentityClass {
+            output += """
+
+                        {
+                            const swiftRetained = __bjs_readIdentityFlag() === 0;
+                            if (!swiftRetained) {
+                                \(JSGlueVariableScope.reservedInstance).exports.bjs_identity_retain(pointer);
+                            }
+                        }
+                """
+        }
+        output += """
+
                     return makeFresh(identityCache);
                 }
 
@@ -350,6 +384,9 @@ public struct BridgeJSLink {
             "",
             "let _exports = null;",
             "let bjs = null;",
+            "let __bjs_identity_flag_ptr = 0;",
+            "let __bjs_identity_flag_buffer = null;",
+            "let __bjs_identity_flag_view = null;",
         ]
     }
 
@@ -981,6 +1018,9 @@ public struct BridgeJSLink {
         printer.write(lines: data.topLevelTypeLines)
 
         let exportedSkeletons = skeletons.compactMap(\.exported)
+        let hasAnyIdentityClass = exportedSkeletons.contains { skeleton in
+            skeleton.classes.contains { shouldUseIdentityCache(for: $0) }
+        }
         let topLevelNamespaceCode = namespaceBuilder.buildTopLevelNamespaceInitialization(
             exportedSkeletons: exportedSkeletons
         )
@@ -998,6 +1038,20 @@ public struct BridgeJSLink {
         try printer.indent {
             printer.write(lines: generateVariableDeclarations())
 
+            if hasAnyIdentityClass {
+                printer.write(lines: [
+                    "",
+                    "function __bjs_readIdentityFlag() {",
+                    "    if (!__bjs_identity_flag_ptr) return 0;",
+                    "    if (__bjs_identity_flag_buffer !== \(JSGlueVariableScope.reservedMemory).buffer) {",
+                    "        __bjs_identity_flag_buffer = \(JSGlueVariableScope.reservedMemory).buffer;",
+                    "        __bjs_identity_flag_view = new DataView(__bjs_identity_flag_buffer, __bjs_identity_flag_ptr, 4);",
+                    "    }",
+                    "    return __bjs_identity_flag_view.getInt32(0, true);",
+                    "}",
+                ])
+            }
+
             let bodyPrinter = CodeFragmentPrinter()
             let allStructs = exportedSkeletons.flatMap { $0.structs }
             for structDef in allStructs {
@@ -1087,6 +1141,12 @@ public struct BridgeJSLink {
                     )
                 }
                 printer.write("}")
+                // Initialize identity flag pointer if any class uses identity caching
+                if hasAnyIdentityClass {
+                    printer.write(
+                        "if (\(JSGlueVariableScope.reservedInstance).exports.bjs_get_identity_flag_ptr) { __bjs_identity_flag_ptr = \(JSGlueVariableScope.reservedInstance).exports.bjs_get_identity_flag_ptr(); }"
+                    )
+                }
             }
             printer.write("},")
         }

Plugins/BridgeJS/Tests/BridgeJSToolTests/__Snapshots__/BridgeJSCodegenTests/IdentityModeClass.swift

@@ -1,9 +1,21 @@
+nonisolated(unsafe) var _CachedModel_identityExported: Set<UnsafeMutableRawPointer> = []
+
 @_expose(wasm, "bjs_CachedModel_init")
 @_cdecl("bjs_CachedModel_init")
 public func _bjs_CachedModel_init(_ nameBytes: Int32, _ nameLength: Int32) -> UnsafeMutableRawPointer {
     #if arch(wasm32)
     let ret = CachedModel(name: String.bridgeJSLiftParameter(nameBytes, nameLength))
-    return ret.bridgeJSLowerReturn()
+    return withExtendedLifetime(ret) {
+        let pointer = Unmanaged.passUnretained(ret).toOpaque()
+        if _CachedModel_identityExported.contains(pointer) {
+            _bridgeJSIdentityFlag = 1
+            return pointer
+        }
+        _CachedModel_identityExported.insert(pointer)
+        _ = Unmanaged.passRetained(ret)
+        _bridgeJSIdentityFlag = 0
+        return pointer
+    }
     #else
     fatalError("Only available on WebAssembly")
     #endif
@@ -34,6 +46,7 @@ public func _bjs_CachedModel_name_set(_ _self: UnsafeMutableRawPointer, _ valueB
 @_cdecl("bjs_CachedModel_deinit")
 public func _bjs_CachedModel_deinit(_ pointer: UnsafeMutableRawPointer) -> Void {
     #if arch(wasm32)
+    _CachedModel_identityExported.remove(pointer)
     Unmanaged<CachedModel>.fromOpaque(pointer).release()
     #else
     fatalError("Only available on WebAssembly")

Plugins/BridgeJS/Tests/BridgeJSToolTests/__Snapshots__/BridgeJSLinkTests/ArrayTypes.js

@@ -43,6 +43,9 @@ export async function createInstantiator(options, swift) {
 
     let _exports = null;
     let bjs = null;
+    let __bjs_identity_flag_ptr = 0;
+    let __bjs_identity_flag_buffer = null;
+    let __bjs_identity_flag_view = null;
     const __bjs_createPointHelpers = () => ({
         lower: (value) => {
             f64Stack.push(value.x);
@@ -373,7 +376,7 @@ export async function createInstantiator(options, swift) {
 
                     const cached = identityCache.get(pointer)?.deref();
                     if (cached && !cached.__swiftHeapObjectState.hasReleased) {
-                        deinit(pointer);
+                                    deinit(pointer);
                         return cached;
                     }
                     if (identityCache.has(pointer)) {

Plugins/BridgeJS/Tests/BridgeJSToolTests/__Snapshots__/BridgeJSLinkTests/Async.js

@@ -30,6 +30,9 @@ export async function createInstantiator(options, swift) {
 
     let _exports = null;
     let bjs = null;
+    let __bjs_identity_flag_ptr = 0;
+    let __bjs_identity_flag_buffer = null;
+    let __bjs_identity_flag_view = null;
 
     return {
         /**