Skip to content

Commit 0919be8

Browse files
kateinoigakukunkateinoigakukun
committed
BridgeJS: Add JSTypedClosure API
The new API allows managing JS closures converted from Swift closures from Swift side. It allows us to get the underlying `JSObject` and manage its lifetime manually from Swift. ```swift @js func makeIntToInt() throws(JSException) -> JSTypedClosure<(Int) -> Int> { return JSTypedClosure { x in return x + 1 } } @JSFunction func takeIntToInt(_ transform: JSTypedClosure<(Int) -> Int>) throws(JSException) let closure = JSTypedClosure<(Int) -> Int> { x in return x * 2 } defer { closure.release() } try takeIntToInt(closure) ``` Likewise to `JSClosure`, API users are responsible for "manually" releasing the closure when it's no longer needed by calling `release()`. After releasing, the closure becomes unusable and calling it will throw a JS exception (note that we will not segfault even if the closure is called after releasing). ```swift let closure = JSTypedClosure<(Int) -> Int> { x in return x * 2 } closure.release() try closure(10) // "JSException: Attempted to call a released JSTypedClosure created at <file>:<line>" ``` As a belt-and-suspenders measure, the underlying JS closure is also registered with a `FinalizationRegistry` to ensure that the Swift closure box is released when the JS closure is garbage collected, in case the API user forgets to call `release()`. However, relying on this mechanism is **not recommended** because the timing of garbage collection is non-deterministic and it's not guaranteed that it will happen in a timely manner. ---- On the top of the new API, this commit also fixes memory leak issues of closures exported to JS.
1 parent 444393f commit 0919be8

26 files changed

Lines changed: 4105 additions & 1665 deletions

File tree

‎Package.swift‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,15 @@ let useLegacyResourceBundling =
99
Context.environment["JAVASCRIPTKIT_USE_LEGACY_RESOURCE_BUNDLING"].flatMap(Bool.init) ?? false
1010

1111
let testingLinkerFlags: [LinkerSetting] = [
12-
.unsafeFlags([
13-
"-Xlinker", "--stack-first",
14-
"-Xlinker", "--global-base=524288",
15-
"-Xlinker", "-z",
16-
"-Xlinker", "stack-size=524288",
17-
])
12+
.unsafeFlags(
13+
[
14+
"-Xlinker", "--stack-first",
15+
"-Xlinker", "--global-base=524288",
16+
"-Xlinker", "-z",
17+
"-Xlinker", "stack-size=524288",
18+
],
19+
.when(platforms: [.wasi])
20+
)
1821
]
1922

2023
let package = Package(

‎Package@swift-6.2.swift‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,15 @@ let tracingTrait = Trait(
1616
)
1717

1818
let testingLinkerFlags: [LinkerSetting] = [
19-
.unsafeFlags([
20-
"-Xlinker", "--stack-first",
21-
"-Xlinker", "--global-base=524288",
22-
"-Xlinker", "-z",
23-
"-Xlinker", "stack-size=524288",
24-
])
19+
.unsafeFlags(
20+
[
21+
"-Xlinker", "--stack-first",
22+
"-Xlinker", "--global-base=524288",
23+
"-Xlinker", "-z",
24+
"-Xlinker", "stack-size=524288",
25+
],
26+
.when(platforms: [.wasi])
27+
)
2528
]
2629

2730
let package = Package(

‎Plugins/BridgeJS/Sources/BridgeJSCore/ClosureCodegen.swift‎

Lines changed: 39 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public struct ClosureCodegen {
1616

1717
func collectClosureSignatures(from type: BridgeType, into signatures: inout Set<ClosureSignature>) {
1818
switch type {
19-
case .closure(let signature):
19+
case .closure(let signature, _):
2020
signatures.insert(signature)
2121
for paramType in signature.parameters {
2222
collectClosureSignatures(from: paramType, into: &signatures)
@@ -32,15 +32,14 @@ public struct ClosureCodegen {
3232
func renderClosureHelpers(_ signature: ClosureSignature) throws -> [DeclSyntax] {
3333
let mangledName = signature.mangleName
3434
let helperName = "_BJS_Closure_\(mangledName)"
35-
let boxClassName = "_BJS_ClosureBox_\(mangledName)"
3635

3736
let closureParams = signature.parameters.enumerated().map { _, type in
3837
"\(type.swiftType)"
3938
}.joined(separator: ", ")
4039

4140
let swiftEffects = (signature.isAsync ? " async" : "") + (signature.isThrows ? " throws" : "")
4241
let swiftReturnType = signature.returnType.swiftType
43-
let closureType = "(\(closureParams))\(swiftEffects) -> \(swiftReturnType)"
42+
let swiftClosureType = "(\(closureParams))\(swiftEffects) -> \(swiftReturnType)"
4443

4544
let externName = "invoke_js_callback_\(signature.moduleName)_\(mangledName)"
4645

@@ -69,13 +68,15 @@ public struct ClosureCodegen {
6968
// Generate extern declaration using CallJSEmission
7069
let externDecl = builder.renderImportDecl()
7170

72-
let boxClassDecl: DeclSyntax = """
73-
private final class \(raw: boxClassName): _BridgedSwiftClosureBox {
74-
let closure: \(raw: closureType)
75-
init(_ closure: @escaping \(raw: closureType)) {
76-
self.closure = closure
77-
}
71+
let makeClosureExternDecl: DeclSyntax = """
72+
#if arch(wasm32)
73+
@_extern(wasm, module: "bjs", name: "make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)")
74+
fileprivate func make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)(_ boxPtr: UnsafeMutableRawPointer, _ file: UnsafePointer<UInt8>, _ line: UInt32) -> Int32
75+
#else
76+
fileprivate func make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName)(_ boxPtr: UnsafeMutableRawPointer, _ file: UnsafePointer<UInt8>, _ line: UInt32) -> Int32 {
77+
fatalError("Only available on WebAssembly")
7878
}
79+
#endif
7980
"""
8081

8182
let helperEnumDecl = EnumDeclSyntax(
@@ -84,33 +85,6 @@ public struct ClosureCodegen {
8485
},
8586
name: .identifier(helperName),
8687
memberBlockBuilder: {
87-
DeclSyntax(
88-
FunctionDeclSyntax(
89-
modifiers: DeclModifierListSyntax {
90-
DeclModifierSyntax(name: .keyword(.static))
91-
},
92-
name: .identifier("bridgeJSLower"),
93-
signature: FunctionSignatureSyntax(
94-
parameterClause: FunctionParameterClauseSyntax {
95-
FunctionParameterSyntax(
96-
firstName: .wildcardToken(),
97-
secondName: .identifier("closure"),
98-
colon: .colonToken(),
99-
type: TypeSyntax("@escaping \(raw: closureType)")
100-
)
101-
},
102-
returnClause: ReturnClauseSyntax(
103-
arrow: .arrowToken(),
104-
type: IdentifierTypeSyntax(name: .identifier("UnsafeMutableRawPointer"))
105-
)
106-
),
107-
body: CodeBlockSyntax {
108-
"let box = \(raw: boxClassName)(closure)"
109-
"return Unmanaged.passRetained(box).toOpaque()"
110-
}
111-
)
112-
)
113-
11488
DeclSyntax(
11589
FunctionDeclSyntax(
11690
modifiers: DeclModifierListSyntax {
@@ -128,7 +102,7 @@ public struct ClosureCodegen {
128102
},
129103
returnClause: ReturnClauseSyntax(
130104
arrow: .arrowToken(),
131-
type: IdentifierTypeSyntax(name: .identifier(closureType))
105+
type: IdentifierTypeSyntax(name: .identifier(swiftClosureType))
132106
)
133107
),
134108
body: CodeBlockSyntax {
@@ -178,11 +152,32 @@ public struct ClosureCodegen {
178152
)
179153
}
180154
)
181-
return [externDecl, boxClassDecl, DeclSyntax(helperEnumDecl)]
155+
let typedClosureExtension: DeclSyntax = """
156+
extension JSTypedClosure where Signature == \(raw: swiftClosureType) {
157+
init(fileID: StaticString = #fileID, line: UInt32 = #line, _ body: @escaping \(raw: swiftClosureType)) {
158+
self.init(
159+
makeClosure: make_swift_closure_\(raw: signature.moduleName)_\(raw: signature.mangleName),
160+
body: body,
161+
fileID: fileID,
162+
line: line
163+
)
164+
}
165+
}
166+
"""
167+
168+
return [
169+
externDecl, makeClosureExternDecl, DeclSyntax(helperEnumDecl), typedClosureExtension,
170+
]
182171
}
183172

184173
func renderClosureInvokeHandler(_ signature: ClosureSignature) throws -> DeclSyntax {
185-
let boxClassName = "_BJS_ClosureBox_\(signature.mangleName)"
174+
let closureParams = signature.parameters.enumerated().map { _, type in
175+
"\(type.swiftType)"
176+
}.joined(separator: ", ")
177+
let swiftEffects = (signature.isAsync ? " async" : "") + (signature.isThrows ? " throws" : "")
178+
let swiftReturnType = signature.returnType.swiftType
179+
let swiftClosureType = "(\(closureParams))\(swiftEffects) -> \(swiftReturnType)"
180+
let boxType = "_BridgeJSTypedClosureBox<\(swiftClosureType)>"
186181
let abiName = "invoke_swift_closure_\(signature.moduleName)_\(signature.mangleName)"
187182

188183
// Build ABI parameters directly with WasmCoreType (no string conversion needed)
@@ -205,7 +200,7 @@ public struct ClosureCodegen {
205200
liftedParams.append("\(paramType.swiftType).bridgeJSLiftParameter(\(argNames.joined(separator: ", ")))")
206201
}
207202

208-
let closureCallExpr = ExprSyntax("box.closure(\(raw: liftedParams.joined(separator: ", ")))")
203+
let closureCallExpr = ExprSyntax("closure(\(raw: liftedParams.joined(separator: ", ")))")
209204

210205
// Determine return type
211206
let abiReturnWasmType: WasmCoreType?
@@ -217,6 +212,8 @@ public struct ClosureCodegen {
217212
abiReturnWasmType = nil
218213
}
219214

215+
let throwReturn = abiReturnWasmType?.swiftReturnPlaceholderStmt ?? "return"
216+
220217
// Build signature using SwiftSignatureBuilder
221218
let funcSignature = SwiftSignatureBuilder.buildABIFunctionSignature(
222219
abiParameters: abiParams,
@@ -225,7 +222,7 @@ public struct ClosureCodegen {
225222

226223
// Build body
227224
let body = CodeBlockItemListSyntax {
228-
"let box = Unmanaged<\(raw: boxClassName)>.fromOpaque(boxPtr).takeUnretainedValue()"
225+
"let closure = Unmanaged<\(raw: boxType)>.fromOpaque(boxPtr).takeUnretainedValue().closure"
229226
if signature.returnType == .void {
230227
closureCallExpr
231228
} else {
@@ -315,7 +312,7 @@ public struct ClosureCodegen {
315312
for setter in type.setters {
316313
collectClosureSignatures(from: setter.type, into: &closureSignatures)
317314
}
318-
for method in type.methods {
315+
for method in type.methods + type.staticMethods {
319316
collectClosureSignatures(from: method.parameters, into: &closureSignatures)
320317
collectClosureSignatures(from: method.returnType, into: &closureSignatures)
321318
}

‎Plugins/BridgeJS/Sources/BridgeJSCore/ExportSwift.swift‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ public class ExportSwift {
124124
let liftingExpr: ExprSyntax
125125

126126
switch param.type {
127-
case .closure(let signature):
127+
case .closure(let signature, _):
128128
typeNameForIntrinsic = param.type.swiftType
129129
liftingExpr = ExprSyntax("_BJS_Closure_\(raw: signature.mangleName).bridgeJSLift(\(raw: param.name))")
130130
case .swiftStruct(let structName):
@@ -364,8 +364,8 @@ public class ExportSwift {
364364
}
365365

366366
switch returnType {
367-
case .closure(let signature):
368-
append("return _BJS_Closure_\(raw: signature.mangleName).bridgeJSLower(ret)")
367+
case .closure(_, useJSTypedClosure: false):
368+
append("return JSTypedClosure(ret).bridgeJSLowerReturn()")
369369
case .array, .nullable(.array, _):
370370
let stackCodegen = StackCodegen()
371371
for stmt in stackCodegen.lowerStatements(for: returnType, accessor: "ret", varPrefix: "ret") {
@@ -423,14 +423,7 @@ public class ExportSwift {
423423
}
424424

425425
private func returnPlaceholderStmt() -> String {
426-
switch abiReturnType {
427-
case .i32: return "return 0"
428-
case .i64: return "return 0"
429-
case .f32: return "return 0.0"
430-
case .f64: return "return 0.0"
431-
case .pointer: return "return UnsafeMutableRawPointer(bitPattern: -1).unsafelyUnwrapped"
432-
case .none: return "return"
433-
}
426+
return abiReturnType?.swiftReturnPlaceholderStmt ?? "return"
434427
}
435428
}
436429

@@ -1749,13 +1742,19 @@ extension BridgeType {
17491742
case .associatedValueEnum(let name): return name
17501743
case .swiftStruct(let name): return name
17511744
case .namespaceEnum(let name): return name
1752-
case .closure(let signature):
1745+
case .closure(let signature, let useJSTypedClosure):
17531746
let paramTypes = signature.parameters.map { $0.swiftType }.joined(separator: ", ")
17541747
let effectsStr = (signature.isAsync ? " async" : "") + (signature.isThrows ? " throws" : "")
1755-
return "(\(paramTypes))\(effectsStr) -> \(signature.returnType.swiftType)"
1748+
let closureType = "(\(paramTypes))\(effectsStr) -> \(signature.returnType.swiftType)"
1749+
return useJSTypedClosure ? "JSTypedClosure<\(closureType)>" : closureType
17561750
}
17571751
}
17581752

1753+
var isClosureType: Bool {
1754+
if case .closure = self { return true }
1755+
return false
1756+
}
1757+
17591758
struct LiftingIntrinsicInfo: Sendable {
17601759
let parameters: [(name: String, type: WasmCoreType)]
17611760

@@ -1853,7 +1852,7 @@ extension BridgeType {
18531852
case .namespaceEnum:
18541853
throw BridgeJSCoreError("Namespace enums are not supported to pass as parameters")
18551854
case .closure:
1856-
return .swiftHeapObject
1855+
return .jsObject
18571856
case .array, .dictionary:
18581857
return .array
18591858
}

‎Plugins/BridgeJS/Sources/BridgeJSCore/ImportTS.swift‎

Lines changed: 42 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ public struct ImportTS {
8585
var destructuredVarNames: [String] = []
8686
// Stack-lowered parameters should be evaluated in reverse order to match LIFO stacks
8787
var stackLoweringStmts: [CodeBlockItemSyntax] = []
88+
// Values to extend lifetime during call
89+
var valuesToExtendLifetimeDuringCall: [ExprSyntax] = []
8890

8991
init(moduleName: String, abiName: String, context: BridgeContext = .importTS) {
9092
self.moduleName = moduleName
@@ -95,15 +97,17 @@ public struct ImportTS {
9597
func lowerParameter(param: Parameter) throws {
9698
let loweringInfo = try param.type.loweringParameterInfo(context: context)
9799

98-
let initializerExpr: ExprSyntax
99100
switch param.type {
100-
case .closure(let signature):
101-
initializerExpr = ExprSyntax(
102-
"_BJS_Closure_\(raw: signature.mangleName).bridgeJSLower(\(raw: param.name))"
103-
)
101+
case .closure(let signature, useJSTypedClosure: false):
102+
let jsTypedClosureType = BridgeType.closure(signature, useJSTypedClosure: true).swiftType
103+
body.append("let \(raw: param.name) = \(raw: jsTypedClosureType)(\(raw: param.name))")
104+
// The just created JSObject is not owned by the caller unlike those passed in parameters,
105+
// so we need to extend its lifetime during the call to ensure the JSObject.id is valid.
106+
valuesToExtendLifetimeDuringCall.append("\(raw: param.name)")
104107
default:
105-
initializerExpr = ExprSyntax("\(raw: param.name).bridgeJSLowerParameter()")
108+
break
106109
}
110+
let initializerExpr = ExprSyntax("\(raw: param.name).bridgeJSLowerParameter()")
107111

108112
if loweringInfo.loweredParameters.isEmpty {
109113
let stmt = CodeBlockItemSyntax(
@@ -193,7 +197,7 @@ public struct ImportTS {
193197
let liftingInfo = try returnType.liftingReturnInfo(context: context)
194198
body.append(contentsOf: stackLoweringStmts)
195199

196-
let callExpr = FunctionCallExprSyntax(
200+
var callExpr = FunctionCallExprSyntax(
197201
calledExpression: ExprSyntax("\(raw: abiName)"),
198202
leftParen: .leftParenToken(),
199203
arguments: LabeledExprListSyntax {
@@ -204,12 +208,33 @@ public struct ImportTS {
204208
rightParen: .rightParenToken()
205209
)
206210

207-
if returnType == .void {
208-
body.append(CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr)))))
209-
} else if returnType.usesSideChannelForOptionalReturn() {
210-
// Side channel returns don't need "let ret ="
211-
body.append(CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr)))))
212-
} else if liftingInfo.valueToLift == nil {
211+
if !valuesToExtendLifetimeDuringCall.isEmpty {
212+
callExpr = FunctionCallExprSyntax(
213+
calledExpression: ExprSyntax("withExtendedLifetime"),
214+
leftParen: .leftParenToken(),
215+
arguments: LabeledExprListSyntax {
216+
LabeledExprSyntax(
217+
expression: TupleExprSyntax(
218+
elements: LabeledExprListSyntax {
219+
for value in valuesToExtendLifetimeDuringCall {
220+
LabeledExprSyntax(expression: value)
221+
}
222+
}
223+
)
224+
)
225+
},
226+
rightParen: .rightParenToken(),
227+
trailingClosure: ClosureExprSyntax(
228+
leftBrace: .leftBraceToken(),
229+
statements: CodeBlockItemListSyntax {
230+
CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr))))
231+
},
232+
rightBrace: .rightBraceToken()
233+
)
234+
)
235+
}
236+
237+
if returnType == .void || returnType.usesSideChannelForOptionalReturn() || liftingInfo.valueToLift == nil {
213238
body.append(CodeBlockItemSyntax(item: .stmt(StmtSyntax(ExpressionStmtSyntax(expression: callExpr)))))
214239
} else {
215240
body.append("let ret = \(raw: callExpr)")
@@ -249,7 +274,7 @@ public struct ImportTS {
249274
abiReturnType = liftingInfo.valueToLift
250275
let liftExpr: ExprSyntax
251276
switch returnType {
252-
case .closure(let signature):
277+
case .closure(let signature, _):
253278
liftExpr = ExprSyntax("_BJS_Closure_\(raw: signature.mangleName).bridgeJSLift(ret)")
254279
default:
255280
if liftingInfo.valueToLift != nil {
@@ -722,11 +747,9 @@ struct SwiftSignatureBuilder {
722747
}
723748

724749
/// Builds a parameter type syntax from a BridgeType.
725-
///
726-
/// Swift closure parameters must be `@escaping` because they are boxed and can be invoked from JavaScript.
727750
static func buildParameterTypeSyntax(from type: BridgeType) -> TypeSyntax {
728751
switch type {
729-
case .closure:
752+
case .closure(_, useJSTypedClosure: false):
730753
return TypeSyntax("@escaping \(raw: type.swiftType)")
731754
default:
732755
return buildTypeSyntax(from: type)
@@ -930,8 +953,8 @@ extension BridgeType {
930953
case .jsValue: return .jsValue
931954
case .void: return .void
932955
case .closure:
933-
// Swift closure is boxed and passed to JS as a pointer.
934-
return LoweringParameterInfo(loweredParameters: [("pointer", .pointer)])
956+
// Swift closure is passed to JS as a JS function reference.
957+
return LoweringParameterInfo(loweredParameters: [("funcRef", .i32)])
935958
case .unsafePointer:
936959
return LoweringParameterInfo(loweredParameters: [("pointer", .pointer)])
937960
case .swiftHeapObject(let className):

0 commit comments

Comments
 (0)