Fix crash when a nested symbol is extended and the a local symbol shadow extended module name (#1202)

* Use `inContextOf` relationships to form hierarchy

* Fixed inversed filters in debug assertion messages

* Avoid computing disambiguation during path hierarchy initialization

rdar://149838723

Author: d-ronnqvist

Sources/SwiftDocC/Infrastructure/Link Resolution/PathHierarchy+DisambiguatedPaths.swift

@@ -193,6 +193,8 @@ extension PathHierarchy.DisambiguationContainer {
         includeLanguage: Bool = false,
         allowAdvancedDisambiguation: Bool = true
     ) -> [(value: PathHierarchy.Node, disambiguation: Disambiguation)] {
+        assert(elements.allSatisfy({ $0.node.identifier != nil }), "All nodes should have been assigned an identifier before their disambiguation can be computed.")
+        
         var collisions = _disambiguatedValues(for: elements, includeLanguage: includeLanguage, allowAdvancedDisambiguation: allowAdvancedDisambiguation)
 
         // If all but one of the collisions are disfavored, remove the disambiguation for the only favored element.

Sources/SwiftDocC/Infrastructure/Link Resolution/PathHierarchy+Find.swift

@@ -447,19 +447,26 @@ extension PathHierarchy.DisambiguationContainer {
         case lookupCollision([(node: PathHierarchy.Node, disambiguation: String)])
     }
 
-    /// Attempts to find a value in the disambiguation tree based on partial disambiguation information.
+    /// Attempts to find the only element in the disambiguation container without using any disambiguation information.
+    ///
+    /// - Returns: The only element in the container or `nil` if the container has more than one element.
+    func singleMatch() -> PathHierarchy.Node? {
+        if storage.count <= 1 {
+            return storage.first?.node
+        } else {
+            return storage.singleMatch({ !$0.node.isDisfavoredInLinkCollisions })?.node
+        }
+    }
+    
+    /// Attempts to find a value in the disambiguation container based on partial disambiguation information.
     ///
     /// There are 3 possible results:
     ///  - No match is found; indicated by a `nil` return value.
     ///  - Exactly one match is found; indicated by a non-nil return value.
     ///  - More than one match is found; indicated by a raised error listing the matches and their missing disambiguation.
     func find(_ disambiguation: PathHierarchy.PathComponent.Disambiguation?) throws -> PathHierarchy.Node? {
-        if disambiguation == nil {
-            if storage.count <= 1 {
-                return storage.first?.node
-            } else if let favoredMatch = storage.singleMatch({ !$0.node.isDisfavoredInLinkCollisions }) {
-                return favoredMatch.node
-            }
+        if disambiguation == nil, let match = singleMatch() {
+            return match
         }
 
         switch disambiguation {

Sources/SwiftDocC/Infrastructure/Link Resolution/PathHierarchy.swift

@@ -232,7 +232,7 @@ struct PathHierarchy {
                         return original
                     }
                 }(node.symbol!)[...].dropLast()
-                while !components.isEmpty, let child = try? parent.children[components.first!]?.find(nil) {
+                while !components.isEmpty, let child = parent.children[components.first!]?.singleMatch() {
                     parent = child
                     components = components.dropFirst()
                 }
@@ -311,7 +311,7 @@ struct PathHierarchy {
         assert(
             allNodes.allSatisfy({ $0.value[0].parent != nil || roots[$0.key] != nil }), """
             Every node should either have a parent node or be a root node. \
-            This wasn't true for \(allNodes.filter({ $0.value[0].parent != nil || roots[$0.key] != nil }).map(\.key).sorted())
+            This wasn't true for \(allNodes.filter({ $0.value[0].parent == nil && roots[$0.key] == nil }).map(\.key).sorted())
             """
         )
 
@@ -320,7 +320,7 @@ struct PathHierarchy {
                 Array(sequence(first: node, next: \.parent)).last!.symbol!.kind.identifier == .module })
             }), """
             Every node should reach a root node by following its parents up. \
-            This wasn't true for \(allNodes.filter({ $0.value.allSatisfy({ Array(sequence(first: $0, next: \.parent)).last!.symbol!.kind.identifier == .module }) }).map(\.key).sorted())
+            This wasn't true for \(allNodes.filter({ $0.value.allSatisfy({ Array(sequence(first: $0, next: \.parent)).last!.symbol!.kind.identifier != .module }) }).map(\.key).sorted())
             """
         )
 
@@ -366,7 +366,7 @@ struct PathHierarchy {
         assert(
             lookup.allSatisfy({ $0.value.parent != nil || roots[$0.value.name] != nil }), """
             Every node should either have a parent node or be a root node. \
-            This wasn't true for \(allNodes.filter({ $0.value[0].parent != nil || roots[$0.key] != nil }).map(\.key).sorted())
+            This wasn't true for \(allNodes.filter({ $0.value[0].parent == nil && roots[$0.key] == nil }).map(\.key).sorted())
             """
         )
 
@@ -389,7 +389,7 @@ struct PathHierarchy {
         assert(
             lookup.values.allSatisfy({ $0.parent?.identifier == nil || lookup[$0.parent!.identifier] != nil }), """
             Every node's findable parent should exist in the lookup. \
-            This wasn't true for \(lookup.values.filter({ $0.parent?.identifier == nil || lookup[$0.parent!.identifier] != nil }).map(\.symbol!.identifier.precise).sorted())
+            This wasn't true for \(lookup.values.filter({ $0.parent?.identifier != nil && lookup[$0.parent!.identifier] == nil }).map(\.symbol!.identifier.precise).sorted())
             """
         )
 
@@ -799,7 +799,7 @@ private extension SymbolGraph.Relationship.Kind {
     /// Whether or not this relationship kind forms a hierarchical relationship between the source and the target.
     var formsHierarchy: Bool {
         switch self {
-        case .memberOf, .optionalMemberOf, .requirementOf, .optionalRequirementOf, .extensionTo, .declaredIn:
+        case .memberOf, .optionalMemberOf, .requirementOf, .optionalRequirementOf, .extensionTo, .inContextOf, .declaredIn:
             return true
         default:
             return false

Tests/SwiftDocCTests/Infrastructure/PathHierarchyTests.swift

@@ -2191,6 +2191,78 @@ class PathHierarchyTests: XCTestCase {
         try assertFindsPath("Inner/InnerClass/something()", in: tree, asSymbolID: "s:5Inner0A5ClassC5OuterE9somethingyyF")
     }
 
+    func testExtensionSymbolsWithSameNameAsExtendedModule() throws {
+        // ---- ExtendedModule
+        // public struct SomeStruct {
+        //     public struct SomeNestedStruct {}
+        // }
+        //
+        // ---- ModuleName
+        // public import ExtendedModule
+        //
+        // // Shadow the ExtendedModule module with a local type
+        // public enum ExtendedModule {}
+        //
+        // // Extend the nested type from the extended module
+        // public extension SomeStruct.SomeNestedStruct {
+        //     func doSomething() {}
+        // }
+        
+        let extensionMixin = SymbolGraph.Symbol.Swift.Extension(extendedModule: "ExtendedModule", typeKind: .struct, constraints: [])
+        
+        let extensionSymbolID      = "s:e:s:14ExtendedModule10SomeStructV0c6NestedD0V0B4NameE11doSomethingyyF"
+        let extendedMethodSymbolID =     "s:14ExtendedModule10SomeStructV0c6NestedD0V0B4NameE11doSomethingyyF"
+        
+        let catalog = Folder(name: "CatalogName.docc", content: [
+            JSONFile(name: "ModuleName.symbols.json", content: makeSymbolGraph(
+                moduleName: "ModuleName",
+                symbols: [
+                    makeSymbol(id: "s:10ModuleName08ExtendedA0O", kind: .enum, pathComponents: ["ExtendedModule"])
+                ])
+            ),
+            
+            JSONFile(name: "[email protected]", content: makeSymbolGraph(
+                moduleName: "ModuleName",
+                symbols: [
+                    // The 'SomeNestedStruct' extension
+                    makeSymbol(id: extensionSymbolID, kind: .extension, pathComponents: ["SomeStruct", "SomeNestedStruct"], otherMixins: [extensionMixin]),
+                    // The 'doSomething()' method added in the extension
+                    makeSymbol(id: extendedMethodSymbolID, kind: .method, pathComponents: ["SomeStruct", "SomeNestedStruct", "doSomething()"], otherMixins: [extensionMixin]),
+                ],
+                relationships: [
+                    // 'doSomething()' is a member of the extension
+                    .init(source: extendedMethodSymbolID, target: extensionSymbolID, kind: .memberOf, targetFallback: "ExtendedModule.SomeStruct.SomeNestedStruct"),
+                    // The extension extends the external 'SomeNestedStruct' symbol
+                    .init(source: extensionSymbolID, target: "s:14ExtendedModule10SomeStructV0c6NestedD0V", kind: .extensionTo, targetFallback: "ExtendedModule.SomeStruct.SomeNestedStruct"),
+                ])
+            ),
+        ])
+        
+        let (_, context) = try loadBundle(catalog: catalog)
+        let tree = context.linkResolver.localResolver.pathHierarchy
+
+        let paths = tree.caseInsensitiveDisambiguatedPaths()
+        XCTAssertEqual(paths[extendedMethodSymbolID], "/ModuleName/ExtendedModule/SomeStruct/SomeNestedStruct/doSomething()")
+        
+        try assertPathCollision("ModuleName/ExtendedModule", in: tree, collisions: [
+            ("s:m:s:e:\(extensionSymbolID)", "-module.extension"),
+            ("s:10ModuleName08ExtendedA0O", "-enum"),
+        ])
+        // If the first path component is ambiguous, it should have the same error as if that was a later path component.
+        try assertPathCollision("ExtendedModule", in: tree, collisions: [
+            ("s:m:s:e:\(extensionSymbolID)", "-module.extension"),
+            ("s:10ModuleName08ExtendedA0O", "-enum"),
+        ])
+        
+        try assertFindsPath("ExtendedModule-enum", in: tree, asSymbolID: "s:10ModuleName08ExtendedA0O")
+        try assertFindsPath("ExtendedModule-module.extension", in: tree, asSymbolID: "s:m:s:e:\(extensionSymbolID)")
+        
+        // The "Inner" struct doesn't have "InnerStruct" or "InnerClass" descendants so the path is not ambiguous.
+        try assertFindsPath("ExtendedModule/SomeStruct", in: tree, asSymbolID: "s:e:\(extensionSymbolID)")
+        try assertFindsPath("ExtendedModule/SomeStruct/SomeNestedStruct", in: tree, asSymbolID: extensionSymbolID)
+        try assertFindsPath("ExtendedModule/SomeStruct/SomeNestedStruct/doSomething()", in: tree, asSymbolID: extendedMethodSymbolID)
+    }
+    
     func testContinuesSearchingIfNonSymbolMatchesSymbolLink() throws {
         let exampleDocumentation = Folder(name: "CatalogName.docc", content: [
             JSONFile(name: "ModuleName.symbols.json", content: makeSymbolGraph(moduleName: "ModuleName", symbols: [