BridgeJS: Correctly emit @JS methods in extensions

Plugins/BridgeJS/Sources/BridgeJSCore/SwiftToSkeleton.swift

@@ -43,12 +43,14 @@ public final class SwiftToSkeleton {
         var perSourceErrors: [(inputFilePath: String, errors: [DiagnosticError])] = []
         var importedFiles: [ImportedFileSkeleton] = []
         var exported = ExportedSkeleton(functions: [], classes: [], enums: [], exposeToGlobal: exposeToGlobal)
+        var exportCollectors: [ExportSwiftAPICollector] = []
 
         for (sourceFile, inputFilePath) in sourceFiles {
             progress.print("Processing \(inputFilePath)")
 
             let exportCollector = ExportSwiftAPICollector(parent: self)
             exportCollector.walk(sourceFile)
+            exportCollectors.append(exportCollector)
 
             let typeNameCollector = ImportSwiftMacrosJSImportTypeNameCollector(viewMode: .sourceAccurate)
             typeNameCollector.walk(sourceFile)
@@ -74,7 +76,15 @@ public final class SwiftToSkeleton {
             if !importedFile.isEmpty {
                 importedFiles.append(importedFile)
             }
-            exportCollector.finalize(&exported)
+        }
+
+        // Resolve extensions against all collectors. This needs to happen at this point so we can resolve both same file and cross file extensions.
+        for source in exportCollectors {
+            source.resolveDeferredExtensions(against: exportCollectors)
+        }
+
+        for collector in exportCollectors {
+            collector.finalize(&exported)
         }
 
         if !perSourceErrors.isEmpty {
@@ -486,6 +496,8 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
     var exportedStructNames: [String] = []
     var exportedStructByName: [String: ExportedStruct] = [:]
     var errors: [DiagnosticError] = []
+    /// Extensions collected during the walk, to be resolved after all files have been walked
+    var deferredExtensions: [ExtensionDeclSyntax] = []
 
     func finalize(_ result: inout ExportedSkeleton) {
         result.functions.append(contentsOf: exportedFunctions)
@@ -1388,6 +1400,64 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
         }
     }
 
+    override func visit(_ node: ExtensionDeclSyntax) -> SyntaxVisitorContinueKind {
+        // Defer until all type declarations in the module have been collected.
+        deferredExtensions.append(node)
+        return .skipChildren
+    }
+
+    func resolveDeferredExtensions(against collectors: [ExportSwiftAPICollector]) {
+        for ext in deferredExtensions {
+            var resolved = false
+            for collector in collectors {
+                if collector.resolveExtension(ext) {
+                    resolved = true
+                    break
+                }
+            }
+            if !resolved {
+                diagnose(
+                    node: ext.extendedType,
+                    message: "Unsupported type '\(ext.extendedType.trimmedDescription)'.",
+                    hint: "You can only extend `@JS` annotated types defined in the same module"
+                )
+            }
+        }
+    }
+
+    /// Walks extension members under the matching type’s state, returning whether the type was found.
+    ///
+    /// Note: The lookup scans dictionaries keyed by `makeKey(name:namespace:)`, matching only by
+    /// plain name. If two types share a name but differ by namespace, `.first(where:)` picks
+    /// whichever comes first. This is acceptable today since namespace collisions are unlikely,
+    /// but may need refinement if namespace-qualified extension resolution is added.
+    func resolveExtension(_ ext: ExtensionDeclSyntax) -> Bool {
+        let name = ext.extendedType.trimmedDescription
+        let state: State
+        if let entry = exportedClassByName.first(where: { $0.value.name == name }) {
+            state = .classBody(name: name, key: entry.key)
+        } else if let entry = exportedStructByName.first(where: { $0.value.name == name }) {
+            state = .structBody(name: name, key: entry.key)
+        } else if let entry = exportedEnumByName.first(where: { $0.value.name == name }) {
+            state = .enumBody(name: name, key: entry.key)
+        } else if exportedProtocolByName.values.contains(where: { $0.name == name }) {
+            diagnose(
+                node: ext.extendedType,
+                message: "Protocol extensions are not supported by BridgeJS.",
+                hint: "You cannot extend `@JS` protocol '\(name)' with additional members"
+            )
+            return true
+        } else {
+            return false
+        }
+        stateStack.push(state: state)
+        for member in ext.memberBlock.members {
+            walk(member)
+        }
+        stateStack.pop()
+        return true
+    }
+
     override func visit(_ node: EnumDeclSyntax) -> SyntaxVisitorContinueKind {
         guard let jsAttribute = node.attributes.firstJSAttribute else {
             return .skipChildren

Plugins/BridgeJS/Tests/BridgeJSToolTests/BridgeJSCodegenTests.swift

@@ -167,6 +167,23 @@ import Testing
         try snapshotCodegen(skeleton: skeleton, name: "CrossFileFunctionTypes.ReverseOrder")
     }
 
+    @Test
+    func codegenCrossFileExtension() throws {
+        let swiftAPI = SwiftToSkeleton(progress: .silent, moduleName: "TestModule", exposeToGlobal: false)
+        let classURL = Self.multifileInputsDirectory.appendingPathComponent("CrossFileExtensionClass.swift")
+        swiftAPI.addSourceFile(
+            Parser.parse(source: try String(contentsOf: classURL, encoding: .utf8)),
+            inputFilePath: "CrossFileExtensionClass.swift"
+        )
+        let extensionURL = Self.multifileInputsDirectory.appendingPathComponent("CrossFileExtension.swift")
+        swiftAPI.addSourceFile(
+            Parser.parse(source: try String(contentsOf: extensionURL, encoding: .utf8)),
+            inputFilePath: "CrossFileExtension.swift"
+        )
+        let skeleton = try swiftAPI.finalize()
+        try snapshotCodegen(skeleton: skeleton, name: "CrossFileExtension")
+    }
+
     @Test
     func codegenSkipsEmptySkeletons() throws {
         let swiftAPI = SwiftToSkeleton(progress: .silent, moduleName: "TestModule", exposeToGlobal: false)

Plugins/BridgeJS/Tests/BridgeJSToolTests/Inputs/MacroSwift/Multifile/CrossFileExtension.swift

@@ -0,0 +1,5 @@
+extension Greeter {
+    @JS func greetFormally() -> String {
+        return "Good day, " + self.name + "."
+    }
+}

Plugins/BridgeJS/Tests/BridgeJSToolTests/Inputs/MacroSwift/Multifile/CrossFileExtensionClass.swift

@@ -0,0 +1,4 @@
+@JS class Greeter {
+    @JS init(name: String) {}
+    @JS func greet() -> String { return "" }
+}

Plugins/BridgeJS/Tests/BridgeJSToolTests/Inputs/MacroSwift/StaticFunctions.swift

@@ -38,3 +38,19 @@ enum APIResult {
         }
     }
 }
+
+extension MathUtils {
+    @JS static func divide(a: Int, b: Int) -> Int {
+        return a / b
+    }
+
+    @JS static var pi: Double { 3.14159 }
+}
+
+extension Calculator {
+    @JS static func cube(value: Int) -> Int {
+        return value * value * value
+    }
+
+    @JS static var version: String { "1.0" }
+}

Plugins/BridgeJS/Tests/BridgeJSToolTests/Inputs/MacroSwift/SwiftClass.swift

@@ -12,6 +12,20 @@
     }
 }
 
+extension Greeter {
+    @JS func greetEnthusiastically() -> String {
+        return "Hey, " + self.name + "!!!"
+    }
+
+    @JS var nameCount: Int { name.count }
+
+    @JS static func greetAnonymously() -> String {
+        return "Hello."
+    }
+
+    @JS static var defaultGreeting: String { "Hello, world!" }
+}
+
 @JS func takeGreeter(greeter: Greeter) {
     print(greeter.greet())
 }

Plugins/BridgeJS/Tests/BridgeJSToolTests/Inputs/MacroSwift/SwiftStruct.swift

@@ -60,3 +60,26 @@
 }
 
 @JS func roundtripContainer(_ container: Container) -> Container
+
+@JS struct Vector2D {
+    var dx: Double
+    var dy: Double
+}
+
+extension Vector2D {
+    @JS func magnitude() -> Double {
+        return (dx * dx + dy * dy).squareRoot()
+    }
+
+    @JS func scaled(by factor: Double) -> Vector2D {
+        return Vector2D(dx: dx * factor, dy: dy * factor)
+    }
+}
+
+extension DataPoint {
+    @JS static func origin() -> DataPoint {
+        return DataPoint(x: 0, y: 0, label: "origin", optCount: nil, optFlag: nil)
+    }
+
+    @JS static var dimensions: Int { 2 }
+}

Plugins/BridgeJS/Tests/BridgeJSToolTests/__Snapshots__/BridgeJSCodegenTests/CrossFileExtension.json

@@ -0,0 +1,82 @@
+{
+  "exported" : {
+    "classes" : [
+      {
+        "constructor" : {
+          "abiName" : "bjs_Greeter_init",
+          "effects" : {
+            "isAsync" : false,
+            "isStatic" : false,
+            "isThrows" : false
+          },
+          "parameters" : [
+            {
+              "label" : "name",
+              "name" : "name",
+              "type" : {
+                "string" : {
+
+                }
+              }
+            }
+          ]
+        },
+        "methods" : [
+          {
+            "abiName" : "bjs_Greeter_greet",
+            "effects" : {
+              "isAsync" : false,
+              "isStatic" : false,
+              "isThrows" : false
+            },
+            "name" : "greet",
+            "parameters" : [
+
+            ],
+            "returnType" : {
+              "string" : {
+
+              }
+            }
+          },
+          {
+            "abiName" : "bjs_Greeter_greetFormally",
+            "effects" : {
+              "isAsync" : false,
+              "isStatic" : false,
+              "isThrows" : false
+            },
+            "name" : "greetFormally",
+            "parameters" : [
+
+            ],
+            "returnType" : {
+              "string" : {
+
+              }
+            }
+          }
+        ],
+        "name" : "Greeter",
+        "properties" : [
+
+        ],
+        "swiftCallName" : "Greeter"
+      }
+    ],
+    "enums" : [
+
+    ],
+    "exposeToGlobal" : false,
+    "functions" : [
+
+    ],
+    "protocols" : [
+
+    ],
+    "structs" : [
+
+    ]
+  },
+  "moduleName" : "TestModule"
+}

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

@@ -0,0 +1,60 @@
+@_expose(wasm, "bjs_Greeter_init")
+@_cdecl("bjs_Greeter_init")
+public func _bjs_Greeter_init(_ nameBytes: Int32, _ nameLength: Int32) -> UnsafeMutableRawPointer {
+    #if arch(wasm32)
+    let ret = Greeter(name: String.bridgeJSLiftParameter(nameBytes, nameLength))
+    return ret.bridgeJSLowerReturn()
+    #else
+    fatalError("Only available on WebAssembly")
+    #endif
+}
+
+@_expose(wasm, "bjs_Greeter_greet")
+@_cdecl("bjs_Greeter_greet")
+public func _bjs_Greeter_greet(_ _self: UnsafeMutableRawPointer) -> Void {
+    #if arch(wasm32)
+    let ret = Greeter.bridgeJSLiftParameter(_self).greet()
+    return ret.bridgeJSLowerReturn()
+    #else
+    fatalError("Only available on WebAssembly")
+    #endif
+}
+
+@_expose(wasm, "bjs_Greeter_greetFormally")
+@_cdecl("bjs_Greeter_greetFormally")
+public func _bjs_Greeter_greetFormally(_ _self: UnsafeMutableRawPointer) -> Void {
+    #if arch(wasm32)
+    let ret = Greeter.bridgeJSLiftParameter(_self).greetFormally()
+    return ret.bridgeJSLowerReturn()
+    #else
+    fatalError("Only available on WebAssembly")
+    #endif
+}
+
+@_expose(wasm, "bjs_Greeter_deinit")
+@_cdecl("bjs_Greeter_deinit")
+public func _bjs_Greeter_deinit(_ pointer: UnsafeMutableRawPointer) -> Void {
+    #if arch(wasm32)
+    Unmanaged<Greeter>.fromOpaque(pointer).release()
+    #else
+    fatalError("Only available on WebAssembly")
+    #endif
+}
+
+extension Greeter: ConvertibleToJSValue, _BridgedSwiftHeapObject {
+    var jsValue: JSValue {
+        return .object(JSObject(id: UInt32(bitPattern: _bjs_Greeter_wrap(Unmanaged.passRetained(self).toOpaque()))))
+    }
+}
+
+#if arch(wasm32)
+@_extern(wasm, module: "TestModule", name: "bjs_Greeter_wrap")
+fileprivate func _bjs_Greeter_wrap_extern(_ pointer: UnsafeMutableRawPointer) -> Int32
+#else
+fileprivate func _bjs_Greeter_wrap_extern(_ pointer: UnsafeMutableRawPointer) -> Int32 {
+    fatalError("Only available on WebAssembly")
+}
+#endif
+@inline(never) fileprivate func _bjs_Greeter_wrap(_ pointer: UnsafeMutableRawPointer) -> Int32 {
+    return _bjs_Greeter_wrap_extern(pointer)
+}