BridgeJS: Add JSTypedClosure API

Author: kateinoigakukun

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

@@ -32,15 +32,16 @@ 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 typedClosureType = "JSTypedClosure<\(swiftClosureType)>"
+        let boxType = "_BridgeJSTypedClosureBox<\(swiftClosureType)>"
 
         let externName = "invoke_js_callback_\(signature.moduleName)_\(mangledName)"
 
@@ -65,17 +66,25 @@ public struct ClosureCodegen {
 
         // Get the body code
         let bodyCode = builder.getBody()
+        let wasmBody = SwiftCodePattern.buildWasmConditionalCompilation(wasmBody: bodyCode.statements)
+            .formatted(using: BasicFormat()).description
+        let closureParamsList = signature.parameters.enumerated().map { "param\($0.offset)" }.joined(separator: ", ")
+        let wrapperParamClause = closureParamsList.isEmpty ? "" : "\(closureParamsList) "
+        let wrapperParamIn = closureParamsList.isEmpty ? "" : "in "
+        let wrapperCallArguments = closureParamsList.isEmpty ? "" : "\(closureParamsList)"
 
         // 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) -> Int32
+            #else
+            fileprivate func make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)(_ boxPtr: UnsafeMutableRawPointer) -> Int32 {
+                fatalError("Only available on WebAssembly")
             }
+            #endif
             """
 
         let helperEnumDecl = EnumDeclSyntax(
@@ -84,33 +93,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,61 +110,60 @@ public struct ClosureCodegen {
                             },
                             returnClause: ReturnClauseSyntax(
                                 arrow: .arrowToken(),
-                                type: IdentifierTypeSyntax(name: .identifier(closureType))
+                                type: IdentifierTypeSyntax(name: .identifier(swiftClosureType))
                             )
                         ),
                         body: CodeBlockSyntax {
                             "let callback = JSObject.bridgeJSLiftParameter(callbackId)"
-                            ReturnStmtSyntax(
-                                expression: ClosureExprSyntax(
-                                    leftBrace: .leftBraceToken(),
-                                    signature: ClosureSignatureSyntax(
-                                        capture: ClosureCaptureClauseSyntax(
-                                            leftSquare: .leftSquareToken(),
-                                            items: ClosureCaptureListSyntax {
-                                                #if canImport(SwiftSyntax602)
-                                                ClosureCaptureSyntax(
-                                                    name: .identifier("", presence: .missing),
-                                                    initializer: InitializerClauseSyntax(
-                                                        equal: .equalToken(presence: .missing),
-                                                        nil,
-                                                        value: ExprSyntax("callback")
-                                                    ),
-                                                    trailingTrivia: nil
-                                                )
-                                                #else
-                                                ClosureCaptureSyntax(
-                                                    expression: ExprSyntax("callback")
-                                                )
-                                                #endif
-                                            },
-                                            rightSquare: .rightSquareToken()
-                                        ),
-                                        parameterClause: .simpleInput(
-                                            ClosureShorthandParameterListSyntax {
-                                                for (index, _) in signature.parameters.enumerated() {
-                                                    ClosureShorthandParameterSyntax(name: .identifier("param\(index)"))
-                                                }
-                                            }
-                                        ),
-                                        inKeyword: .keyword(.in)
-                                    ),
-                                    statements: CodeBlockItemListSyntax {
-                                        SwiftCodePattern.buildWasmConditionalCompilation(wasmBody: bodyCode.statements)
-                                    },
-                                    rightBrace: .rightBraceToken()
-                                )
-                            )
+                            """
+                            let callable: \(raw: swiftClosureType) = { [callback] \(raw: closureParamsList) in
+                            \(raw: wasmBody)
+                            }
+                            """
+                            "return callable"
                         }
                     )
                 )
             }
         )
-        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)) {
+                    let box = \(raw: boxType)(body, fileID: fileID, line: line)
+                    let pointer = Unmanaged.passRetained(box).toOpaque()
+                    let funcRefValue = make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)(pointer)
+                    let jsObject = JSObject(id: UInt32(bitPattern: funcRefValue))
+                    self.init(
+                        _boxPointer: pointer,
+                        jsObject: jsObject
+                    )
+                }
+
+                static func oneshot(fileID: StaticString = #fileID, line: UInt32 = #line, _ body: @escaping \(raw: swiftClosureType)) -> JSTypedClosure<Signature> {
+                    var typedClosure: JSTypedClosure<Signature>!
+                    let wrapper: Signature = { \(raw: wrapperParamClause)\(raw: wrapperParamIn)
+                        defer { typedClosure.release() }
+                        return body(\(raw: wrapperCallArguments))
+                    }
+                    typedClosure = JSTypedClosure(fileID: fileID, line: line, wrapper)
+                    return typedClosure
+                }
+            }
+            """
+
+        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 +186,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 +198,19 @@ public struct ClosureCodegen {
             abiReturnWasmType = nil
         }
 
+        let throwReturn: String
+        if let abiReturnWasmType {
+            switch abiReturnWasmType {
+            case .i32: throwReturn = "return 0"
+            case .i64: throwReturn = "return 0"
+            case .f32: throwReturn = "return 0"
+            case .f64: throwReturn = "return 0"
+            case .pointer: throwReturn = "return UnsafeMutableRawPointer(bitPattern: 0)!"
+            }
+        } else {
+            throwReturn = "return"
+        }
+
         // Build signature using SwiftSignatureBuilder
         let funcSignature = SwiftSignatureBuilder.buildABIFunctionSignature(
             abiParameters: abiParams,
@@ -225,7 +219,11 @@ public struct ClosureCodegen {
 
         // Build body
         let body = CodeBlockItemListSyntax {
-            "let box = Unmanaged<\(raw: boxClassName)>.fromOpaque(boxPtr).takeUnretainedValue()"
+            "let box = Unmanaged<\(raw: boxType)>.fromOpaque(boxPtr).takeUnretainedValue()"
+            "guard let closure = box.closure else {"
+            "    box._bridgeJSThrowReleasedClosure()"
+            "    \(raw: throwReturn)"
+            "}"
             if signature.returnType == .void {
                 closureCallExpr
             } else {

Plugins/BridgeJS/Sources/BridgeJSCore/ExportSwift.swift

@@ -365,7 +365,7 @@ public class ExportSwift {
 
             switch returnType {
             case .closure(let signature):
-                append("return _BJS_Closure_\(raw: signature.mangleName).bridgeJSLower(ret)")
+                append("return JSTypedClosure(ret)._bridgeJSLowerReturn()")
             case .array, .nullable(.array, _):
                 let stackCodegen = StackCodegen()
                 for stmt in stackCodegen.lowerStatements(for: returnType, accessor: "ret", varPrefix: "ret") {
@@ -1628,6 +1628,35 @@ extension BridgeType {
         }
     }
 
+    /// Swift-side parameter type to use when sending a closure into JavaScript (ImportTS).
+    var swiftImportParameterType: String {
+        switch self {
+        case .closure(let signature):
+            let paramTypes = signature.parameters.map { $0.swiftType }.joined(separator: ", ")
+            let effectsStr = (signature.isAsync ? " async" : "") + (signature.isThrows ? " throws" : "")
+            return "JSTypedClosure<(\(paramTypes))\(effectsStr) -> \(signature.returnType.swiftType)>"
+        case .nullable(let wrapped, let kind):
+            if wrapped.isClosureType {
+                let wrappedType = wrapped.swiftImportParameterType
+                switch kind {
+                case .null:
+                    return "Optional<\(wrappedType)>"
+                case .undefined:
+                    return "JSUndefinedOr<\(wrappedType)>"
+                }
+            } else {
+                return swiftType
+            }
+        default:
+            return swiftType
+        }
+    }
+
+    var isClosureType: Bool {
+        if case .closure = self { return true }
+        return false
+    }
+
     struct LiftingIntrinsicInfo: Sendable {
         let parameters: [(name: String, type: WasmCoreType)]
 
@@ -1725,7 +1754,7 @@ extension BridgeType {
         case .namespaceEnum:
             throw BridgeJSCoreError("Namespace enums are not supported to pass as parameters")
         case .closure:
-            return .swiftHeapObject
+            return .jsObject
         case .array:
             return .array
         }

Plugins/BridgeJS/Sources/BridgeJSCore/ImportTS.swift

@@ -99,7 +99,7 @@ public struct ImportTS {
             switch param.type {
             case .closure(let signature):
                 initializerExpr = ExprSyntax(
-                    "_BJS_Closure_\(raw: signature.mangleName).bridgeJSLower(\(raw: param.name))"
+                    "\(raw: param.name)._bridgeJSLowerParameter()"
                 )
             default:
                 initializerExpr = ExprSyntax("\(raw: param.name).bridgeJSLowerParameter()")
@@ -722,12 +722,12 @@ 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:
-            return TypeSyntax("@escaping \(raw: type.swiftType)")
+            return TypeSyntax("\(raw: type.swiftImportParameterType)")
+        case .nullable(let wrapped, _) where wrapped.isClosureType:
+            return TypeSyntax("\(raw: type.swiftImportParameterType)")
         default:
             return buildTypeSyntax(from: type)
         }
@@ -930,8 +930,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):