Skip to content

BridgeJS: Add opt-in TypedArray bridging for numeric arrays #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
krodak wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

BridgeJS: Add opt-in TypedArray bridging for numeric arrays #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 13 additions & 13 deletions Benchmarks/Sources/Generated/BridgeJS.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1829,9 +1829,9 @@ public func _bjs_ArrayRoundtrip_init() -> UnsafeMutableRawPointer {

@_expose(wasm, "bjs_ArrayRoundtrip_takeIntArray")
@_cdecl("bjs_ArrayRoundtrip_takeIntArray")
public func _bjs_ArrayRoundtrip_takeIntArray(_ _self: UnsafeMutableRawPointer) -> Void {
public func _bjs_ArrayRoundtrip_takeIntArray(_ _self: UnsafeMutableRawPointer, _ valuesSourceId: Int32, _ valuesCount: Int32) -> Void {
#if arch(wasm32)
ArrayRoundtrip.bridgeJSLiftParameter(_self).takeIntArray(_: [Int].bridgeJSStackPop())
ArrayRoundtrip.bridgeJSLiftParameter(_self).takeIntArray(_: _bridgeJS_typedArrayLiftParameter(valuesSourceId, valuesCount) as [Int])
#else
fatalError("Only available on WebAssembly")
#endif
Expand All @@ -1842,18 +1842,18 @@ public func _bjs_ArrayRoundtrip_takeIntArray(_ _self: UnsafeMutableRawPointer) -
public func _bjs_ArrayRoundtrip_makeIntArray(_ _self: UnsafeMutableRawPointer) -> Void {
#if arch(wasm32)
let ret = ArrayRoundtrip.bridgeJSLiftParameter(_self).makeIntArray()
ret.bridgeJSStackPush()
_bridgeJS_typedArrayPush(ret)
#else
fatalError("Only available on WebAssembly")
#endif
}

@_expose(wasm, "bjs_ArrayRoundtrip_roundtripIntArray")
@_cdecl("bjs_ArrayRoundtrip_roundtripIntArray")
public func _bjs_ArrayRoundtrip_roundtripIntArray(_ _self: UnsafeMutableRawPointer) -> Void {
public func _bjs_ArrayRoundtrip_roundtripIntArray(_ _self: UnsafeMutableRawPointer, _ valuesSourceId: Int32, _ valuesCount: Int32) -> Void {
#if arch(wasm32)
let ret = ArrayRoundtrip.bridgeJSLiftParameter(_self).roundtripIntArray(_: [Int].bridgeJSStackPop())
ret.bridgeJSStackPush()
let ret = ArrayRoundtrip.bridgeJSLiftParameter(_self).roundtripIntArray(_: _bridgeJS_typedArrayLiftParameter(valuesSourceId, valuesCount) as [Int])
_bridgeJS_typedArrayPush(ret)
#else
fatalError("Only available on WebAssembly")
#endif
Expand All @@ -1864,17 +1864,17 @@ public func _bjs_ArrayRoundtrip_roundtripIntArray(_ _self: UnsafeMutableRawPoint
public func _bjs_ArrayRoundtrip_makeIntArrayLarge(_ _self: UnsafeMutableRawPointer) -> Void {
#if arch(wasm32)
let ret = ArrayRoundtrip.bridgeJSLiftParameter(_self).makeIntArrayLarge()
ret.bridgeJSStackPush()
_bridgeJS_typedArrayPush(ret)
#else
fatalError("Only available on WebAssembly")
#endif
}

@_expose(wasm, "bjs_ArrayRoundtrip_takeDoubleArray")
@_cdecl("bjs_ArrayRoundtrip_takeDoubleArray")
public func _bjs_ArrayRoundtrip_takeDoubleArray(_ _self: UnsafeMutableRawPointer) -> Void {
public func _bjs_ArrayRoundtrip_takeDoubleArray(_ _self: UnsafeMutableRawPointer, _ valuesSourceId: Int32, _ valuesCount: Int32) -> Void {
#if arch(wasm32)
ArrayRoundtrip.bridgeJSLiftParameter(_self).takeDoubleArray(_: [Double].bridgeJSStackPop())
ArrayRoundtrip.bridgeJSLiftParameter(_self).takeDoubleArray(_: _bridgeJS_typedArrayLiftParameter(valuesSourceId, valuesCount) as [Double])
#else
fatalError("Only available on WebAssembly")
#endif
Expand All @@ -1885,18 +1885,18 @@ public func _bjs_ArrayRoundtrip_takeDoubleArray(_ _self: UnsafeMutableRawPointer
public func _bjs_ArrayRoundtrip_makeDoubleArray(_ _self: UnsafeMutableRawPointer) -> Void {
#if arch(wasm32)
let ret = ArrayRoundtrip.bridgeJSLiftParameter(_self).makeDoubleArray()
ret.bridgeJSStackPush()
_bridgeJS_typedArrayPush(ret)
#else
fatalError("Only available on WebAssembly")
#endif
}

@_expose(wasm, "bjs_ArrayRoundtrip_roundtripDoubleArray")
@_cdecl("bjs_ArrayRoundtrip_roundtripDoubleArray")
public func _bjs_ArrayRoundtrip_roundtripDoubleArray(_ _self: UnsafeMutableRawPointer) -> Void {
public func _bjs_ArrayRoundtrip_roundtripDoubleArray(_ _self: UnsafeMutableRawPointer, _ valuesSourceId: Int32, _ valuesCount: Int32) -> Void {
#if arch(wasm32)
let ret = ArrayRoundtrip.bridgeJSLiftParameter(_self).roundtripDoubleArray(_: [Double].bridgeJSStackPop())
ret.bridgeJSStackPush()
let ret = ArrayRoundtrip.bridgeJSLiftParameter(_self).roundtripDoubleArray(_: _bridgeJS_typedArrayLiftParameter(valuesSourceId, valuesCount) as [Double])
_bridgeJS_typedArrayPush(ret)
#else
fatalError("Only available on WebAssembly")
#endif
Expand Down
13 changes: 13 additions & 0 deletions Plugins/BridgeJS/Sources/BridgeJSCore/ClosureCodegen.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,19 @@ public struct ClosureCodegen {

for (index, paramType) in signature.parameters.enumerated() {
let paramName = "param\(index)"

// Numeric arrays use bulk TypedArray transfer with (sourceId, count) WASM params
if case .array(let elementType) = paramType, elementType.isNumericScalar {
let sourceIdName = "\(paramName)SourceId"
let countName = "\(paramName)Count"
abiParams.append((sourceIdName, .i32))
abiParams.append((countName, .i32))
liftedParams.append(
"_bridgeJS_typedArrayLiftParameter(\(sourceIdName), \(countName)) as [\(elementType.swiftType)]"
)
continue
}

let liftInfo = try paramType.liftParameterInfo()

for (argName, wasmType) in liftInfo.parameters {
Expand Down
14 changes: 14 additions & 0 deletions Plugins/BridgeJS/Sources/BridgeJSCore/ExportSwift.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,18 @@ public class ExportSwift {
case .swiftStruct(let structName):
typeNameForIntrinsic = structName
liftingExpr = ExprSyntax("\(raw: structName).bridgeJSLiftParameter()")
case .array(let elementType) where elementType.isNumericScalar:
// Numeric arrays use bulk TypedArray transfer with (sourceId, count) WASM params
let elementSwiftType = elementType.swiftType
typeNameForIntrinsic = param.type.swiftType
liftingExpr = ExprSyntax(
"_bridgeJS_typedArrayLiftParameter(\(raw: param.name)SourceId, \(raw: param.name)Count) as [\(raw: elementSwiftType)]"
)
// Override the ABI signatures: two i32 params instead of stack-based
liftedParameterExprs.append(liftingExpr)
abiParameterSignatures.append(("\(param.name)SourceId", .i32))
abiParameterSignatures.append(("\(param.name)Count", .i32))
return
case .array:
typeNameForIntrinsic = param.type.swiftType
liftingExpr = StackCodegen().liftExpression(for: param.type)
Expand Down Expand Up @@ -301,6 +313,8 @@ public class ExportSwift {
switch returnType {
case .closure(_, useJSTypedClosure: false):
append("return JSTypedClosure(ret).bridgeJSLowerReturn()")
case .array(let elementType) where elementType.isNumericScalar:
append("_bridgeJS_typedArrayPush(ret)")
case .array, .nullable(.array, _):
let stackCodegen = StackCodegen()
for stmt in stackCodegen.lowerStatements(for: returnType, accessor: "ret", varPrefix: "ret") {
Expand Down
10 changes: 10 additions & 0 deletions Plugins/BridgeJS/Sources/BridgeJSCore/ImportTS.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,10 @@ public struct ImportTS {
// 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(param.name)
case .array(let elementType) where elementType.isNumericScalar:
// Numeric arrays use bulk TypedArray transfer instead of element-by-element
stackLoweringStmts.insert("_bridgeJS_typedArrayPush(\(param.name))", at: 0)
return
default:
break
}
Expand Down Expand Up @@ -302,6 +306,12 @@ public struct ImportTS {
switch returnType {
case .closure(let signature, _):
liftExpr = "_BJS_Closure_\(signature.mangleName).bridgeJSLift(ret)"
case .array(let elementType) where elementType.isNumericScalar:
// Numeric arrays: JS pushed (sourceId, count) onto i32 stack
let swiftElementType = elementType.swiftType
body.write("let _count = _swift_js_pop_i32()")
body.write("let _sourceId = _swift_js_pop_i32()")
liftExpr = "_bridgeJS_typedArrayLiftParameter(_sourceId, _count) as [\(swiftElementType)]"
default:
if liftingInfo.valueToLift != nil {
liftExpr = "\(returnType.swiftType).bridgeJSLiftReturn(ret)"
Expand Down
62 changes: 62 additions & 0 deletions Plugins/BridgeJS/Sources/BridgeJSLink/BridgeJSLink.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,8 @@ public struct BridgeJSLink {
"let \(JSGlueVariableScope.reservedF32Stack) = [];",
"let \(JSGlueVariableScope.reservedF64Stack) = [];",
"let \(JSGlueVariableScope.reservedPointerStack) = [];",
"let \(JSGlueVariableScope.reservedTaStack) = [];",
"const \(JSGlueVariableScope.reservedTypedArrayConstructors) = [Int8Array, Uint8Array, Int16Array, Uint16Array, Int32Array, Uint32Array, BigInt64Array, BigUint64Array, Float32Array, Float64Array];",
"const \(JSGlueVariableScope.reservedEnumHelpers) = {};",
"const \(JSGlueVariableScope.reservedStructHelpers) = {};",
"",
Expand Down Expand Up @@ -487,6 +489,66 @@ public struct BridgeJSLink {
printer.write("return \(JSGlueVariableScope.reservedI64Stack).pop();")
}
printer.write("}")
printer.write("bjs[\"swift_js_push_typed_array\"] = function(ptr, count, kind) {")
printer.indent {
printer.write(
"const Constructor = \(JSGlueVariableScope.reservedTypedArrayConstructors)[kind];"
)
printer.write(
"const elemSize = Constructor.BYTES_PER_ELEMENT;"
)
printer.write(
"const totalBytes = count * elemSize;"
)
printer.write(
"const copy = new Uint8Array(totalBytes);"
)
printer.write(
"copy.set(new Uint8Array(\(JSGlueVariableScope.reservedMemory).buffer, ptr, totalBytes));"
)
printer.write(
"\(JSGlueVariableScope.reservedTaStack).push(new Constructor(copy.buffer));"
)
}
printer.write("}")
printer.write("bjs[\"swift_js_init_typed_array_memory\"] = function(sourceId, destPtr, count, kind) {")
printer.indent {
printer.write(
"const source = \(JSGlueVariableScope.reservedSwift).\(JSGlueVariableScope.reservedMemory).getObject(sourceId);"
)
printer.write(
"\(JSGlueVariableScope.reservedSwift).\(JSGlueVariableScope.reservedMemory).release(sourceId);"
)
printer.write(
"const Constructor = \(JSGlueVariableScope.reservedTypedArrayConstructors)[kind];"
)
printer.write(
"const elemSize = Constructor.BYTES_PER_ELEMENT;"
)
printer.write(
"if (destPtr % elemSize === 0) {"
)
printer.indent {
printer.write(
"const dest = new Constructor(\(JSGlueVariableScope.reservedMemory).buffer, destPtr, count);"
)
printer.write("dest.set(source);")
}
printer.write(
"} else {"
)
printer.indent {
printer.write(
"const dest = new Uint8Array(\(JSGlueVariableScope.reservedMemory).buffer, destPtr, count * elemSize);"
)
printer.write(
"const srcBytes = new Uint8Array(source.buffer, source.byteOffset, source.byteLength);"
)
printer.write("dest.set(srcBytes);")
}
printer.write("}")
}
printer.write("}")
if !allStructs.isEmpty {
for structDef in allStructs {
printer.write("bjs[\"swift_js_struct_lower_\(structDef.abiName)\"] = function(objectId) {")
Expand Down
78 changes: 78 additions & 0 deletions Plugins/BridgeJS/Sources/BridgeJSLink/JSGlueGen.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ final class JSGlueVariableScope {
static let reservedStructHelpers = "structHelpers"
static let reservedSwiftClosureRegistry = "swiftClosureRegistry"
static let reservedMakeSwiftClosure = "makeClosure"
static let reservedTaStack = "taStack"
static let reservedTypedArrayConstructors = "typedArrayConstructors"

private let intrinsicRegistry: JSIntrinsicRegistry

Expand Down Expand Up @@ -63,6 +65,8 @@ final class JSGlueVariableScope {
reservedStructHelpers,
reservedSwiftClosureRegistry,
reservedMakeSwiftClosure,
reservedTaStack,
reservedTypedArrayConstructors,
]

init(intrinsicRegistry: JSIntrinsicRegistry) {
Expand Down Expand Up @@ -1317,6 +1321,8 @@ struct IntrinsicJSFragment: Sendable {
)
case .namespaceEnum(let string):
throw BridgeJSLinkError(message: "Namespace enums are not supported to be passed as parameters: \(string)")
case .array(let elementType) where elementType.isNumericScalar:
return numericArrayLower(elementType: elementType)
case .array(let elementType):
return try arrayLower(elementType: elementType)
case .dictionary(let valueType):
Expand Down Expand Up @@ -1374,6 +1380,8 @@ struct IntrinsicJSFragment: Sendable {
throw BridgeJSLinkError(
message: "Namespace enums are not supported to be returned from functions: \(string)"
)
case .array(let elementType) where elementType.isNumericScalar:
return numericArrayLift(elementType: elementType)
case .array(let elementType):
return try arrayLift(elementType: elementType)
case .dictionary(let valueType):
Expand Down Expand Up @@ -1472,6 +1480,8 @@ struct IntrinsicJSFragment: Sendable {
message:
"Namespace enums are not supported to be passed as parameters to imported JS functions: \(string)"
)
case .array(let elementType) where elementType.isNumericScalar:
return numericArrayLift(elementType: elementType)
case .array(let elementType):
return try arrayLift(elementType: elementType)
case .dictionary(let valueType):
Expand Down Expand Up @@ -1534,6 +1544,8 @@ struct IntrinsicJSFragment: Sendable {
throw BridgeJSLinkError(
message: "Namespace enums are not supported to be returned from imported JS functions: \(string)"
)
case .array(let elementType) where elementType.isNumericScalar:
return numericArrayLowerReturn(elementType: elementType)
case .array(let elementType):
return try arrayLower(elementType: elementType)
case .dictionary(let valueType):
Expand Down Expand Up @@ -1828,6 +1840,72 @@ struct IntrinsicJSFragment: Sendable {

// MARK: - Array Helpers

/// Lowers a numeric array from JS to Swift via bulk TypedArray transfer.
/// Converts the JS Array to a TypedArray, retains it, and passes (sourceId, count) as WASM params.
static func numericArrayLower(elementType: BridgeType) -> IntrinsicJSFragment {
let kind = elementType.typedArrayKind!
return IntrinsicJSFragment(
parameters: ["arr"],
printCode: { arguments, context in
let (scope, printer) = (context.scope, context.printer)
let arr = arguments[0]
let typedArrVar = scope.variable("typedArr")
let idVar = scope.variable("typedArrId")
printer.write(
"const \(typedArrVar) = new \(JSGlueVariableScope.reservedTypedArrayConstructors)[\(kind)](\(arr));"
)
printer.write(
"const \(idVar) = \(JSGlueVariableScope.reservedSwift).memory.retain(\(typedArrVar));"
)
return [idVar, "\(typedArrVar).length"]
}
)
}

/// Lowers a numeric array from JS to Swift for return values (stack-based).
/// Pushes (sourceId, count) onto the i32 stack.
static func numericArrayLowerReturn(elementType: BridgeType) -> IntrinsicJSFragment {
let kind = elementType.typedArrayKind!
return IntrinsicJSFragment(
parameters: ["arr"],
printCode: { arguments, context in
let (scope, printer) = (context.scope, context.printer)
let arr = arguments[0]
let typedArrVar = scope.variable("typedArr")
let idVar = scope.variable("typedArrId")
printer.write(
"const \(typedArrVar) = new \(JSGlueVariableScope.reservedTypedArrayConstructors)[\(kind)](\(arr));"
)
printer.write(
"const \(idVar) = \(JSGlueVariableScope.reservedSwift).memory.retain(\(typedArrVar));"
)
scope.emitPushI32Parameter(idVar, printer: printer)
scope.emitPushI32Parameter("\(typedArrVar).length", printer: printer)
return []
}
)
}

/// Lifts a numeric array from Swift to JS via bulk TypedArray transfer.
/// Swift side pushed the typed array view onto taStack; JS pops and converts to plain Array.
static func numericArrayLift(elementType: BridgeType) -> IntrinsicJSFragment {
return IntrinsicJSFragment(
parameters: [],
printCode: { arguments, context in
let (scope, printer) = (context.scope, context.printer)
let srcVar = scope.variable("taSrc")
let resultVar = scope.variable("arrayResult")
let iVar = scope.variable("i")
printer.write("const \(srcVar) = \(JSGlueVariableScope.reservedTaStack).pop();")
printer.write("const \(resultVar) = new Array(\(srcVar).length);")
printer.write(
"for (let \(iVar) = 0; \(iVar) < \(srcVar).length; \(iVar)++) \(resultVar)[\(iVar)] = \(srcVar)[\(iVar)];"
)
return [resultVar]
}
)
}

/// Lowers an array from JS to Swift by iterating elements and pushing to stacks
static func arrayLower(elementType: BridgeType) throws -> IntrinsicJSFragment {
return IntrinsicJSFragment(
Expand Down
Loading
Loading