BridgeJS: Allow extensions to contain types

Author: wfltaylor

Plugins/BridgeJS/Sources/BridgeJSCore/SwiftToSkeleton.swift

@@ -525,12 +525,12 @@ public final class SwiftToSkeleton {
 
         if let typeDecl = typeDeclResolver.resolve(type) {
             if typeDecl.is(ProtocolDeclSyntax.self) {
-                let swiftCallName = SwiftToSkeleton.computeSwiftCallName(for: typeDecl, itemName: typeDecl.name.text)
+                let swiftCallName = computeSwiftCallName(for: typeDecl, itemName: typeDecl.name.text)
                 return .swiftProtocol(swiftCallName)
             }
 
             if let enumDecl = typeDecl.as(EnumDeclSyntax.self) {
-                let swiftCallName = SwiftToSkeleton.computeSwiftCallName(for: enumDecl, itemName: enumDecl.name.text)
+                let swiftCallName = computeSwiftCallName(for: enumDecl, itemName: enumDecl.name.text)
                 if let jsAttribute = enumDecl.attributes.firstJSAttribute,
                     let aliasTarget = extractAliasTarget(from: jsAttribute)
                 {
@@ -569,7 +569,7 @@ public final class SwiftToSkeleton {
             }
 
             if let structDecl = typeDecl.as(StructDeclSyntax.self) {
-                let swiftCallName = SwiftToSkeleton.computeSwiftCallName(
+                let swiftCallName = computeSwiftCallName(
                     for: structDecl,
                     itemName: structDecl.name.text
                 )
@@ -587,7 +587,7 @@ public final class SwiftToSkeleton {
             guard typeDecl.is(ClassDeclSyntax.self) || typeDecl.is(ActorDeclSyntax.self) else {
                 return nil
             }
-            let swiftCallName = SwiftToSkeleton.computeSwiftCallName(for: typeDecl, itemName: typeDecl.name.text)
+            let swiftCallName = computeSwiftCallName(for: typeDecl, itemName: typeDecl.name.text)
 
             // A type annotated with @JSClass is a JavaScript object wrapper (imported),
             // even if it is declared as a Swift class.
@@ -627,7 +627,7 @@ public final class SwiftToSkeleton {
     private func resolveExternal(for type: TypeSyntax, errors: inout [DiagnosticError]) -> BridgeType? {
         guard
             !externalModuleIndex.isEmpty,
-            var components = typeDeclResolver.qualifiedComponents(from: type)
+            var components = type.qualifiedComponents
         else {
             return nil
         }
@@ -766,27 +766,50 @@ public final class SwiftToSkeleton {
         return nil
     }
 
-    /// Computes the full Swift call name by walking up the AST hierarchy to find all parent enums
+    /// This currently doesn’t work correctly for extensions on types defined in other modules,
+    /// which is fine for now since we don’t support extending @JS types from other modules.
+    /// This will need updating when we do.
+    fileprivate func enclosingDeclarations(of node: some SyntaxProtocol) -> [Syntax] {
+        var declarations: [Syntax] = []
+        var visitedExtendedTypes: Set<SyntaxIdentifier> = []
+        var currentNode: Syntax? = Syntax(node).parent
+
+        while let parent = currentNode {
+            if let extensionDecl = parent.as(ExtensionDeclSyntax.self) {
+                if let extendedDecl = typeDeclResolver.resolve(extensionDecl.extendedType),
+                    visitedExtendedTypes.insert(extendedDecl.id).inserted
+                {
+                    declarations.append(Syntax(extendedDecl))
+                    currentNode = Syntax(extendedDecl).parent
+                } else {
+                    currentNode = parent.parent
+                }
+            } else {
+                declarations.append(parent)
+                currentNode = parent.parent
+            }
+        }
+        return declarations
+    }
+
     /// This generates the qualified name needed for Swift code generation (e.g., "Networking.API.HTTPServer")
-    fileprivate static func computeSwiftCallName(for node: some SyntaxProtocol, itemName: String) -> String {
+    fileprivate func computeSwiftCallName(for node: some SyntaxProtocol, itemName: String) -> String {
         var swiftPath: [String] = []
-        var currentNode: Syntax? = node.parent
 
-        while let parent = currentNode {
-            if let enumDecl = parent.as(EnumDeclSyntax.self),
+        for declaration in enclosingDeclarations(of: node) {
+            if let enumDecl = declaration.as(EnumDeclSyntax.self),
                 enumDecl.attributes.hasJSAttribute()
             {
                 swiftPath.insert(enumDecl.name.text, at: 0)
-            } else if let structDecl = parent.as(StructDeclSyntax.self),
+            } else if let structDecl = declaration.as(StructDeclSyntax.self),
                 structDecl.attributes.hasJSAttribute()
             {
                 swiftPath.insert(structDecl.name.text, at: 0)
-            } else if let classDecl = parent.as(ClassDeclSyntax.self),
+            } else if let classDecl = declaration.as(ClassDeclSyntax.self),
                 classDecl.attributes.hasJSAttribute()
             {
                 swiftPath.insert(classDecl.name.text, at: 0)
             }
-            currentNode = parent.parent
         }
 
         if swiftPath.isEmpty {
@@ -1861,7 +1884,7 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
             resolvedNamespace: namespaceResult.namespace,
             parentTypeNamespace: computeParentTypeNamespace(for: node)
         )
-        let swiftCallName = SwiftToSkeleton.computeSwiftCallName(for: node, itemName: name)
+        let swiftCallName = parent.computeSwiftCallName(for: node, itemName: name)
         let explicitAccessControl = computeExplicitAtLeastInternalAccessControl(
             for: node,
             message: "Class visibility must be at least internal"
@@ -1921,25 +1944,23 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
     }
 
     /// 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
+        guard let extendedDecl = parent.typeDeclResolver.resolve(ext.extendedType) else {
+            return false
+        }
+        let swiftCallName = parent.computeSwiftCallName(for: extendedDecl, itemName: extendedDecl.name.text)
         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 }) {
+        if let entry = exportedClassByName.first(where: { $0.value.swiftCallName == swiftCallName }) {
+            state = .classBody(name: entry.value.name, key: entry.key)
+        } else if let entry = exportedStructByName.first(where: { $0.value.swiftCallName == swiftCallName }) {
+            state = .structBody(name: entry.value.name, key: entry.key)
+        } else if let entry = exportedEnumByName.first(where: { $0.value.swiftCallName == swiftCallName }) {
+            state = .enumBody(name: entry.value.name, key: entry.key)
+        } else if exportedProtocolByName.values.contains(where: { $0.name == swiftCallName }) {
             diagnose(
                 node: ext.extendedType,
                 message: "Protocol extensions are not supported by BridgeJS.",
-                hint: "You cannot extend `@JS` protocol '\(name)' with additional members"
+                hint: "You cannot extend `@JS` protocol '\(swiftCallName)' with additional members"
             )
             return true
         } else {
@@ -1958,7 +1979,7 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
         jsAttribute: AttributeSyntax,
         aliasTarget: TypeSyntax
     ) {
-        let swiftCallName = SwiftToSkeleton.computeSwiftCallName(for: node, itemName: node.name.text)
+        let swiftCallName = parent.computeSwiftCallName(for: node, itemName: node.name.text)
         if extractNamespace(from: jsAttribute) != nil {
             errors.append(
                 DiagnosticError(
@@ -2023,7 +2044,7 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
             parentTypeNamespace: computeParentTypeNamespace(for: node)
         )
         let emitStyle = extractEnumStyle(from: jsAttribute) ?? .const
-        let swiftCallName = SwiftToSkeleton.computeSwiftCallName(for: node, itemName: name)
+        let swiftCallName = parent.computeSwiftCallName(for: node, itemName: name)
         let explicitAccessControl = computeExplicitAtLeastInternalAccessControl(
             for: node,
             message: "Enum visibility must be at least internal"
@@ -2209,7 +2230,7 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
             resolvedNamespace: namespaceResult.namespace,
             parentTypeNamespace: computeParentTypeNamespace(for: node)
         )
-        let swiftCallName = SwiftToSkeleton.computeSwiftCallName(for: node, itemName: name)
+        let swiftCallName = parent.computeSwiftCallName(for: node, itemName: name)
         let explicitAccessControl = computeExplicitAtLeastInternalAccessControl(
             for: node,
             message: "Struct visibility must be at least internal"
@@ -2524,10 +2545,9 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
     /// Method allows for explicit namespace for top level enum, it will be used as base namespace and will concat enum name
     private func computeNamespace(for node: some SyntaxProtocol) -> [String]? {
         var namespace: [String] = []
-        var currentNode: Syntax? = node.parent
 
-        while let parent = currentNode {
-            if let enumDecl = parent.as(EnumDeclSyntax.self),
+        for declaration in parent.enclosingDeclarations(of: node) {
+            if let enumDecl = declaration.as(EnumDeclSyntax.self),
                 enumDecl.attributes.hasJSAttribute()
             {
                 let isNamespaceEnum = !enumDecl.memberBlock.members.contains { member in
@@ -2544,27 +2564,24 @@ private final class ExportSwiftAPICollector: SyntaxAnyVisitor {
                     }
                 }
             }
-            currentNode = parent.parent
         }
 
         return namespace.isEmpty ? nil : namespace
     }
 
     private func computeParentTypeNamespace(for node: some SyntaxProtocol) -> [String]? {
         var path: [String] = []
-        var currentNode: Syntax? = node.parent
 
-        while let parent = currentNode {
-            if let structDecl = parent.as(StructDeclSyntax.self),
+        for declaration in parent.enclosingDeclarations(of: node) {
+            if let structDecl = declaration.as(StructDeclSyntax.self),
                 structDecl.attributes.hasJSAttribute()
             {
                 path.insert(structDecl.name.text, at: 0)
-            } else if let classDecl = parent.as(ClassDeclSyntax.self),
+            } else if let classDecl = declaration.as(ClassDeclSyntax.self),
                 classDecl.attributes.hasJSAttribute()
             {
                 path.insert(classDecl.name.text, at: 0)
             }
-            currentNode = parent.parent
         }
 
         return path.isEmpty ? nil : path

Plugins/BridgeJS/Sources/BridgeJSCore/TypeDeclResolver.swift

@@ -16,8 +16,7 @@ class TypeDeclResolver {
 
     private class TypeDeclCollector: SyntaxVisitor {
         let resolver: TypeDeclResolver
-        var scope: [TypeDecl] = []
-        var rootTypeDecls: [TypeDecl] = []
+        var scope: [String] = []
 
         init(resolver: TypeDeclResolver) {
             self.resolver = resolver
@@ -26,17 +25,14 @@ class TypeDeclResolver {
 
         func visitNominalDecl(_ node: TypeDecl) -> SyntaxVisitorContinueKind {
             let name = node.name.text
-            let qualifiedName = scope.map(\.name.text) + [name]
+            let qualifiedName = scope + [name]
             resolver.typeDeclByQualifiedName[qualifiedName] = node
-            scope.append(node)
+            scope.append(name)
             return .visitChildren
         }
 
         func visitPostNominalDecl() {
-            let type = scope.removeLast()
-            if scope.isEmpty {
-                rootTypeDecls.append(type)
-            }
+            scope.removeLast()
         }
 
         override func visit(_ node: StructDeclSyntax) -> SyntaxVisitorContinueKind {
@@ -72,10 +68,21 @@ class TypeDeclResolver {
 
         override func visit(_ node: TypeAliasDeclSyntax) -> SyntaxVisitorContinueKind {
             let name = node.name.text
-            let qualifiedName = scope.map(\.name.text) + [name]
+            let qualifiedName = scope + [name]
             resolver.typeAliasByQualifiedName[qualifiedName] = node
             return .skipChildren
         }
+
+        override func visit(_ node: ExtensionDeclSyntax) -> SyntaxVisitorContinueKind {
+            guard let components = node.memberScopeComponents else {
+                return .skipChildren
+            }
+            scope.append(contentsOf: components)
+            return .visitChildren
+        }
+        override func visitPost(_ node: ExtensionDeclSyntax) {
+            scope.removeLast(node.memberScopeComponents?.count ?? 0)
+        }
     }
 
     /// Collects type declarations from a parsed Swift source file
@@ -91,6 +98,10 @@ class TypeDeclResolver {
         while let parent = context.parent {
             if let parent = parent.asProtocol(NamedDeclSyntax.self), parent.isProtocol(DeclGroupSyntax.self) {
                 innerToOuter.append(parent.name.text)
+            } else if let extensionDecl = parent.as(ExtensionDeclSyntax.self),
+                let components = extensionDecl.memberScopeComponents
+            {
+                innerToOuter.append(contentsOf: components.reversed())
             }
             context = parent
         }
@@ -106,7 +117,7 @@ class TypeDeclResolver {
         /// Search for the type declaration from the innermost scope to the outermost scope
         for i in (0...scope.count).reversed() {
             let qualifiedName = Array(scope[0..<i] + [name])
-            if typeDeclByQualifiedName[qualifiedName] != nil {
+            if typeDeclByQualifiedName[qualifiedName] != nil || typeAliasByQualifiedName[qualifiedName] != nil {
                 return qualifiedName
             }
         }
@@ -132,15 +143,15 @@ class TypeDeclResolver {
     ///
     /// Resolution strategy:
     /// 1. If the node is IdentifierTypeSyntax, call `lookupType(for:)` which attempts scope-aware qualification via `tryQualify`.
-    /// 2. Otherwise, attempt to build a fully qualified name with `qualifiedComponents(from:)` and look it up with `lookupType(fullyQualified:)`.
+    /// 2. Otherwise, attempt to build a fully qualified name with `qualifiedComponents` and look it up with `lookupType(fullyQualified:)`.
     ///
     /// - Parameter type: The SwiftSyntax node representing a type appearance in source code.
     /// - Returns: The nominal declaration (enum/class/actor/struct) if found, otherwise nil.
     func resolve(_ type: TypeSyntax) -> TypeDecl? {
         if let id = type.as(IdentifierTypeSyntax.self) {
             return lookupType(for: id)
         }
-        if let components = qualifiedComponents(from: type) {
+        if let components = type.qualifiedComponents {
             return lookupType(fullyQualified: components)
         }
         return nil
@@ -155,20 +166,29 @@ class TypeDeclResolver {
             let qualifiedName = tryQualify(type: id)
             return typeAliasByQualifiedName[qualifiedName]
         }
-        if let components = qualifiedComponents(from: type) {
+        if let components = type.qualifiedComponents {
             return typeAliasByQualifiedName[components]
         }
         return nil
     }
 
-    func qualifiedComponents(from type: TypeSyntax) -> QualifiedName? {
-        if let m = type.as(MemberTypeSyntax.self) {
-            guard let base = qualifiedComponents(from: TypeSyntax(m.baseType)) else { return nil }
+}
+
+extension TypeSyntax {
+    var qualifiedComponents: TypeDeclResolver.QualifiedName? {
+        if let m = self.as(MemberTypeSyntax.self) {
+            guard let base = TypeSyntax(m.baseType).qualifiedComponents else { return nil }
             return base + [m.name.text]
-        } else if let id = type.as(IdentifierTypeSyntax.self) {
+        } else if let id = self.as(IdentifierTypeSyntax.self) {
             return [id.name.text]
         } else {
             return nil
         }
     }
 }
+
+extension ExtensionDeclSyntax {
+    var memberScopeComponents: TypeDeclResolver.QualifiedName? {
+        extendedType.qualifiedComponents
+    }
+}

Plugins/BridgeJS/Tests/BridgeJSToolTests/BridgeJSCodegenTests.swift

@@ -336,6 +336,28 @@ import Testing
         try snapshotCodegen(skeleton: skeleton, name: "CrossFileExtension")
     }
 
+    @Test
+    func codegenCrossFileNestedTypeExtension() throws {
+        let swiftAPI = SwiftToSkeleton(
+            progress: .silent,
+            moduleName: "TestModule",
+            exposeToGlobal: false,
+            externalModuleIndex: .empty
+        )
+        let classURL = Self.multifileInputsDirectory.appendingPathComponent("CrossFileNestedTypeClass.swift")
+        swiftAPI.addSourceFile(
+            Parser.parse(source: try String(contentsOf: classURL, encoding: .utf8)),
+            inputFilePath: "CrossFileNestedTypeClass.swift"
+        )
+        let extensionURL = Self.multifileInputsDirectory.appendingPathComponent("CrossFileNestedTypeExtension.swift")
+        swiftAPI.addSourceFile(
+            Parser.parse(source: try String(contentsOf: extensionURL, encoding: .utf8)),
+            inputFilePath: "CrossFileNestedTypeExtension.swift"
+        )
+        let skeleton = try swiftAPI.finalize()
+        try snapshotCodegen(skeleton: skeleton, name: "CrossFileNestedTypeExtension")
+    }
+
     @Test
     func codegenSkipsEmptySkeletons() throws {
         let swiftAPI = SwiftToSkeleton(