BridgeJS: Add JSTypedClosure API

Package.swift

@@ -9,12 +9,15 @@ let useLegacyResourceBundling =
     Context.environment["JAVASCRIPTKIT_USE_LEGACY_RESOURCE_BUNDLING"].flatMap(Bool.init) ?? false
 
 let testingLinkerFlags: [LinkerSetting] = [
-    .unsafeFlags([
-        "-Xlinker", "--stack-first",
-        "-Xlinker", "--global-base=524288",
-        "-Xlinker", "-z",
-        "-Xlinker", "stack-size=524288",
-    ])
+    .unsafeFlags(
+        [
+            "-Xlinker", "--stack-first",
+            "-Xlinker", "--global-base=524288",
+            "-Xlinker", "-z",
+            "-Xlinker", "stack-size=524288",
+        ],
+        .when(platforms: [.wasi])
+    )
 ]
 
 let package = Package(

Package@swift-6.2.swift

@@ -16,12 +16,15 @@ let tracingTrait = Trait(
 )
 
 let testingLinkerFlags: [LinkerSetting] = [
-    .unsafeFlags([
-        "-Xlinker", "--stack-first",
-        "-Xlinker", "--global-base=524288",
-        "-Xlinker", "-z",
-        "-Xlinker", "stack-size=524288",
-    ])
+    .unsafeFlags(
+        [
+            "-Xlinker", "--stack-first",
+            "-Xlinker", "--global-base=524288",
+            "-Xlinker", "-z",
+            "-Xlinker", "stack-size=524288",
+        ],
+        .when(platforms: [.wasi])
+    )
 ]
 
 let package = Package(

Plugins/BridgeJS/Sources/BridgeJSCore/ClosureCodegen.swift

@@ -16,7 +16,7 @@ public struct ClosureCodegen {
 
     func collectClosureSignatures(from type: BridgeType, into signatures: inout Set<ClosureSignature>) {
         switch type {
-        case .closure(let signature):
+        case .closure(let signature, _):
             signatures.insert(signature)
             for paramType in signature.parameters {
                 collectClosureSignatures(from: paramType, into: &signatures)
@@ -32,15 +32,14 @@ public struct ClosureCodegen {
     func renderClosureHelpers(_ signature: ClosureSignature) throws -> [DeclSyntax] {
         let mangledName = signature.mangleName
         let helperName = "_BJS_Closure_\(mangledName)"
-        let boxClassName = "_BJS_ClosureBox_\(mangledName)"
 
         let closureParams = signature.parameters.enumerated().map { _, type in
             "\(type.swiftType)"
         }.joined(separator: ", ")
 
         let swiftEffects = (signature.isAsync ? " async" : "") + (signature.isThrows ? " throws" : "")
         let swiftReturnType = signature.returnType.swiftType
-        let closureType = "(\(closureParams))\(swiftEffects) -> \(swiftReturnType)"
+        let swiftClosureType = "(\(closureParams))\(swiftEffects) -> \(swiftReturnType)"
 
         let externName = "invoke_js_callback_\(signature.moduleName)_\(mangledName)"
 
@@ -69,13 +68,15 @@ public struct ClosureCodegen {
         // Generate extern declaration using CallJSEmission
         let externDecl = builder.renderImportDecl()
 
-        let boxClassDecl: DeclSyntax = """
-            private final class \(raw: boxClassName): _BridgedSwiftClosureBox {
-                let closure: \(raw: closureType)
-                init(_ closure: @escaping \(raw: closureType)) {
-                    self.closure = closure
-                }
+        let makeClosureExternDecl: DeclSyntax = """
+            #if arch(wasm32)
+            @_extern(wasm, module: "bjs", name: "make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)")
+            fileprivate func make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)(_ boxPtr: UnsafeMutableRawPointer, _ file: UnsafePointer<UInt8>, _ line: UInt32) -> Int32
+            #else
+            fileprivate func make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)(_ boxPtr: UnsafeMutableRawPointer, _ file: UnsafePointer<UInt8>, _ line: UInt32) -> Int32 {
+                fatalError("Only available on WebAssembly")
             }
+            #endif
             """
 
         let helperEnumDecl = EnumDeclSyntax(
@@ -84,33 +85,6 @@ public struct ClosureCodegen {
             },
             name: .identifier(helperName),
             memberBlockBuilder: {
-                DeclSyntax(
-                    FunctionDeclSyntax(
-                        modifiers: DeclModifierListSyntax {
-                            DeclModifierSyntax(name: .keyword(.static))
-                        },
-                        name: .identifier("bridgeJSLower"),
-                        signature: FunctionSignatureSyntax(
-                            parameterClause: FunctionParameterClauseSyntax {
-                                FunctionParameterSyntax(
-                                    firstName: .wildcardToken(),
-                                    secondName: .identifier("closure"),
-                                    colon: .colonToken(),
-                                    type: TypeSyntax("@escaping \(raw: closureType)")
-                                )
-                            },
-                            returnClause: ReturnClauseSyntax(
-                                arrow: .arrowToken(),
-                                type: IdentifierTypeSyntax(name: .identifier("UnsafeMutableRawPointer"))
-                            )
-                        ),
-                        body: CodeBlockSyntax {
-                            "let box = \(raw: boxClassName)(closure)"
-                            "return Unmanaged.passRetained(box).toOpaque()"
-                        }
-                    )
-                )
-
                 DeclSyntax(
                     FunctionDeclSyntax(
                         modifiers: DeclModifierListSyntax {
@@ -128,7 +102,7 @@ public struct ClosureCodegen {
                             },
                             returnClause: ReturnClauseSyntax(
                                 arrow: .arrowToken(),
-                                type: IdentifierTypeSyntax(name: .identifier(closureType))
+                                type: IdentifierTypeSyntax(name: .identifier(swiftClosureType))
                             )
                         ),
                         body: CodeBlockSyntax {
@@ -178,11 +152,32 @@ public struct ClosureCodegen {
                 )
             }
         )
-        return [externDecl, boxClassDecl, DeclSyntax(helperEnumDecl)]
+        let typedClosureExtension: DeclSyntax = """
+            extension JSTypedClosure where Signature == \(raw: swiftClosureType) {
+                init(fileID: StaticString = #fileID, line: UInt32 = #line, _ body: @escaping \(raw: swiftClosureType)) {
+                    self.init(
+                        makeClosure: make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName),
+                        body: body,
+                        fileID: fileID,
+                        line: line
+                    )
+                }
+            }
+            """
+
+        return [
+            externDecl, makeClosureExternDecl, DeclSyntax(helperEnumDecl), typedClosureExtension,
+        ]
     }
 
     func renderClosureInvokeHandler(_ signature: ClosureSignature) throws -> DeclSyntax {
-        let boxClassName = "_BJS_ClosureBox_\(signature.mangleName)"
+        let closureParams = signature.parameters.enumerated().map { _, type in
+            "\(type.swiftType)"
+        }.joined(separator: ", ")
+        let swiftEffects = (signature.isAsync ? " async" : "") + (signature.isThrows ? " throws" : "")
+        let swiftReturnType = signature.returnType.swiftType
+        let swiftClosureType = "(\(closureParams))\(swiftEffects) -> \(swiftReturnType)"
+        let boxType = "_BridgeJSTypedClosureBox<\(swiftClosureType)>"
         let abiName = "invoke_swift_closure_\(signature.moduleName)_\(signature.mangleName)"
 
         // Build ABI parameters directly with WasmCoreType (no string conversion needed)
@@ -205,7 +200,7 @@ public struct ClosureCodegen {
             liftedParams.append("\(paramType.swiftType).bridgeJSLiftParameter(\(argNames.joined(separator: ", ")))")
         }
 
-        let closureCallExpr = ExprSyntax("box.closure(\(raw: liftedParams.joined(separator: ", ")))")
+        let closureCallExpr = ExprSyntax("closure(\(raw: liftedParams.joined(separator: ", ")))")
 
         // Determine return type
         let abiReturnWasmType: WasmCoreType?
@@ -217,6 +212,8 @@ public struct ClosureCodegen {
             abiReturnWasmType = nil
         }
 
+        let throwReturn = abiReturnWasmType?.swiftReturnPlaceholderStmt ?? "return"
+
         // Build signature using SwiftSignatureBuilder
         let funcSignature = SwiftSignatureBuilder.buildABIFunctionSignature(
             abiParameters: abiParams,
@@ -225,7 +222,7 @@ public struct ClosureCodegen {
 
         // Build body
         let body = CodeBlockItemListSyntax {
-            "let box = Unmanaged<\(raw: boxClassName)>.fromOpaque(boxPtr).takeUnretainedValue()"
+            "let closure = Unmanaged<\(raw: boxType)>.fromOpaque(boxPtr).takeUnretainedValue().closure"
             if signature.returnType == .void {
                 closureCallExpr
             } else {
@@ -315,7 +312,7 @@ public struct ClosureCodegen {
                     for setter in type.setters {
                         collectClosureSignatures(from: setter.type, into: &closureSignatures)
                     }
-                    for method in type.methods {
+                    for method in type.methods + type.staticMethods {
                         collectClosureSignatures(from: method.parameters, into: &closureSignatures)
                         collectClosureSignatures(from: method.returnType, into: &closureSignatures)
                     }

Plugins/BridgeJS/Sources/BridgeJSCore/ExportSwift.swift

@@ -124,7 +124,7 @@ public class ExportSwift {
             let liftingExpr: ExprSyntax
 
             switch param.type {
-            case .closure(let signature):
+            case .closure(let signature, _):
                 typeNameForIntrinsic = param.type.swiftType
                 liftingExpr = ExprSyntax("_BJS_Closure_\(raw: signature.mangleName).bridgeJSLift(\(raw: param.name))")
             case .swiftStruct(let structName):
@@ -364,8 +364,8 @@ public class ExportSwift {
             }
 
             switch returnType {
-            case .closure(let signature):
-                append("return _BJS_Closure_\(raw: signature.mangleName).bridgeJSLower(ret)")
+            case .closure(_, useJSTypedClosure: false):
+                append("return JSTypedClosure(ret).bridgeJSLowerReturn()")
             case .array, .nullable(.array, _):
                 let stackCodegen = StackCodegen()
                 for stmt in stackCodegen.lowerStatements(for: returnType, accessor: "ret", varPrefix: "ret") {
@@ -423,14 +423,7 @@ public class ExportSwift {
         }
 
         private func returnPlaceholderStmt() -> String {
-            switch abiReturnType {
-            case .i32: return "return 0"
-            case .i64: return "return 0"
-            case .f32: return "return 0.0"
-            case .f64: return "return 0.0"
-            case .pointer: return "return UnsafeMutableRawPointer(bitPattern: -1).unsafelyUnwrapped"
-            case .none: return "return"
-            }
+            return abiReturnType?.swiftReturnPlaceholderStmt ?? "return"
         }
     }
 
@@ -1749,13 +1742,19 @@ extension BridgeType {
         case .associatedValueEnum(let name): return name
         case .swiftStruct(let name): return name
         case .namespaceEnum(let name): return name
-        case .closure(let signature):
+        case .closure(let signature, let useJSTypedClosure):
             let paramTypes = signature.parameters.map { $0.swiftType }.joined(separator: ", ")
             let effectsStr = (signature.isAsync ? " async" : "") + (signature.isThrows ? " throws" : "")
-            return "(\(paramTypes))\(effectsStr) -> \(signature.returnType.swiftType)"
+            let closureType = "(\(paramTypes))\(effectsStr) -> \(signature.returnType.swiftType)"
+            return useJSTypedClosure ? "JSTypedClosure<\(closureType)>" : closureType
         }
     }
 
+    var isClosureType: Bool {
+        if case .closure = self { return true }
+        return false
+    }
+
     struct LiftingIntrinsicInfo: Sendable {
         let parameters: [(name: String, type: WasmCoreType)]
 
@@ -1853,7 +1852,7 @@ extension BridgeType {
         case .namespaceEnum:
             throw BridgeJSCoreError("Namespace enums are not supported to pass as parameters")
         case .closure:
-            return .swiftHeapObject
+            return .jsObject
         case .array, .dictionary:
             return .array
         }

Plugins/BridgeJS/Sources/BridgeJSCore/ImportTS.swift

@@ -85,6 +85,8 @@ public struct ImportTS {
         var destructuredVarNames: [String] = []
         // Stack-lowered parameters should be evaluated in reverse order to match LIFO stacks
         var stackLoweringStmts: [CodeBlockItemSyntax] = []
+        // Values to extend lifetime during call
+        var valuesToExtendLifetimeDuringCall: [ExprSyntax] = []
 
         init(moduleName: String, abiName: String, context: BridgeContext = .importTS) {
             self.moduleName = moduleName
@@ -95,15 +97,17 @@ public struct ImportTS {
         func lowerParameter(param: Parameter) throws {
             let loweringInfo = try param.type.loweringParameterInfo(context: context)
 
-            let initializerExpr: ExprSyntax
             switch param.type {
-            case .closure(let signature):
-                initializerExpr = ExprSyntax(
-                    "_BJS_Closure_\(raw: signature.mangleName).bridgeJSLower(\(raw: param.name))"
-                )
+            case .closure(let signature, useJSTypedClosure: false):
+                let jsTypedClosureType = BridgeType.closure(signature, useJSTypedClosure: true).swiftType
+                body.append("let \(raw: param.name) = \(raw: jsTypedClosureType)(\(raw: param.name))")
+                // The just created JSObject is not owned by the caller unlike those passed in parameters,
+                // so we need to extend its lifetime during the call to ensure the JSObject.id is valid.
+                valuesToExtendLifetimeDuringCall.append("\(raw: param.name)")
             default:
-                initializerExpr = ExprSyntax("\(raw: param.name).bridgeJSLowerParameter()")
+                break
             }
+            let initializerExpr = ExprSyntax("\(raw: param.name).bridgeJSLowerParameter()")
 
             if loweringInfo.loweredParameters.isEmpty {
                 let stmt = CodeBlockItemSyntax(
@@ -193,7 +197,7 @@ public struct ImportTS {
             let liftingInfo = try returnType.liftingReturnInfo(context: context)
             body.append(contentsOf: stackLoweringStmts)
 
-            let callExpr = FunctionCallExprSyntax(
+            var callExpr = FunctionCallExprSyntax(
                 calledExpression: ExprSyntax("\(raw: abiName)"),
                 leftParen: .leftParenToken(),
                 arguments: LabeledExprListSyntax {
@@ -204,12 +208,33 @@ public struct ImportTS {
                 rightParen: .rightParenToken()
             )
 
-            if returnType == .void {
-                body.append(CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr)))))
-            } else if returnType.usesSideChannelForOptionalReturn() {
-                // Side channel returns don't need "let ret ="
-                body.append(CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr)))))
-            } else if liftingInfo.valueToLift == nil {
+            if !valuesToExtendLifetimeDuringCall.isEmpty {
+                callExpr = FunctionCallExprSyntax(
+                    calledExpression: ExprSyntax("withExtendedLifetime"),
+                    leftParen: .leftParenToken(),
+                    arguments: LabeledExprListSyntax {
+                        LabeledExprSyntax(
+                            expression: TupleExprSyntax(
+                                elements: LabeledExprListSyntax {
+                                    for value in valuesToExtendLifetimeDuringCall {
+                                        LabeledExprSyntax(expression: value)
+                                    }
+                                }
+                            )
+                        )
+                    },
+                    rightParen: .rightParenToken(),
+                    trailingClosure: ClosureExprSyntax(
+                        leftBrace: .leftBraceToken(),
+                        statements: CodeBlockItemListSyntax {
+                            CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr))))
+                        },
+                        rightBrace: .rightBraceToken()
+                    )
+                )
+            }
+
+            if returnType == .void || returnType.usesSideChannelForOptionalReturn() || liftingInfo.valueToLift == nil {
                 body.append(CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr)))))
             } else {
                 body.append("let ret = \(raw: callExpr)")
@@ -249,7 +274,7 @@ public struct ImportTS {
                 abiReturnType = liftingInfo.valueToLift
                 let liftExpr: ExprSyntax
                 switch returnType {
-                case .closure(let signature):
+                case .closure(let signature, _):
                     liftExpr = ExprSyntax("_BJS_Closure_\(raw: signature.mangleName).bridgeJSLift(ret)")
                 default:
                     if liftingInfo.valueToLift != nil {
@@ -722,11 +747,9 @@ struct SwiftSignatureBuilder {
     }
 
     /// Builds a parameter type syntax from a BridgeType.
-    ///
-    /// Swift closure parameters must be `@escaping` because they are boxed and can be invoked from JavaScript.
     static func buildParameterTypeSyntax(from type: BridgeType) -> TypeSyntax {
         switch type {
-        case .closure:
+        case .closure(_, useJSTypedClosure: false):
             return TypeSyntax("@escaping \(raw: type.swiftType)")
         default:
             return buildTypeSyntax(from: type)
@@ -930,8 +953,8 @@ extension BridgeType {
         case .jsValue: return .jsValue
         case .void: return .void
         case .closure:
-            // Swift closure is boxed and passed to JS as a pointer.
-            return LoweringParameterInfo(loweredParameters: [("pointer", .pointer)])
+            // Swift closure is passed to JS as a JS function reference.
+            return LoweringParameterInfo(loweredParameters: [("funcRef", .i32)])
         case .unsafePointer:
             return LoweringParameterInfo(loweredParameters: [("pointer", .pointer)])
         case .swiftHeapObject(let className):