Make it easier to store other types of link disambiguating information in the path hierarchy (#828)

* Internally store disambiguated path hierarchy elements as small lists

* Continue to disfavor sparse nodes without disambiguation

* Expand on internal comments to describe more implementation details

Author: d-ronnqvist

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

@@ -112,12 +112,11 @@ extension PathHierarchy {
                         // When descending through placeholder nodes, we trust that the known disambiguation
                         // that they were created with is necessary.
                         var knownDisambiguation = ""
-                        let (kind, subtree) = tree.storage.first!
-                        if kind != "_" {
+                        let element = tree.storage.first!
+                        if let kind = element.kind {
                             knownDisambiguation += "-\(kind)"
                         }
-                        let hash = subtree.keys.first!
-                        if hash != "_" {
+                        if let hash = element.hash {
                             knownDisambiguation += "-\(hash)"
                         }
                         path = accumulatedPath + "/" + nameTransform(node.name) + knownDisambiguation
@@ -170,33 +169,45 @@ extension PathHierarchy {
 }
 
 extension PathHierarchy.DisambiguationContainer {
+    
+    static func disambiguatedValues<E: Sequence>(
+        for elements: E,
+        includeLanguage: Bool = false
+    ) -> [(value: PathHierarchy.Node, disambiguation: Disambiguation)] where E.Element == Element {
+        var collisions: [(value: PathHierarchy.Node, disambiguation: Disambiguation)] = []
+        
+        var remainingIDs = Set(elements.map(\.node.identifier))
+        
+        // Kind disambiguation is the most readable, so we start by checking if any element has a unique kind.
+        let groupedByKind = [String?: [Element]](grouping: elements, by: \.kind)
+        for (kind, elements) in groupedByKind where elements.count == 1 && kind != nil {
+            let element = elements.first!
+            if includeLanguage, let symbol = element.node.symbol {
+                collisions.append((value: element.node, disambiguation: .kind("\(SourceLanguage(id: symbol.identifier.interfaceLanguage).linkDisambiguationID).\(kind!)")))
+            } else {
+                collisions.append((value: element.node, disambiguation: .kind(kind!)))
+            }
+            remainingIDs.remove(element.node.identifier)
+        }
+        if remainingIDs.isEmpty {
+            return collisions
+        }
+        
+        for element in elements where remainingIDs.contains(element.node.identifier) {
+            collisions.append((value: element.node, disambiguation: element.hash.map { .hash($0) } ?? .none))
+        }
+        return collisions
+    }
+    
     /// Returns all values paired with their disambiguation suffixes.
     ///
     /// - Parameter includeLanguage: Whether or not the kind disambiguation information should include the language, for example: "swift".
     func disambiguatedValues(includeLanguage: Bool = false) -> [(value: PathHierarchy.Node, disambiguation: Disambiguation)] {
         if storage.count == 1 {
-            let tree = storage.values.first!
-            if tree.count == 1 {
-                return [(tree.values.first!, .none)]
-            }
+            return [(storage.first!.node, .none)]
         }
 
-        var collisions: [(value: PathHierarchy.Node, disambiguation: Disambiguation)] = []
-        for (kind, kindTree) in storage {
-            if kindTree.count == 1 {
-                // No other match has this kind
-                if includeLanguage, let symbol = kindTree.first!.value.symbol {
-                    collisions.append((value: kindTree.first!.value, disambiguation: .kind("\(SourceLanguage(id: symbol.identifier.interfaceLanguage).linkDisambiguationID).\(kind)")))
-                } else {
-                    collisions.append((value: kindTree.first!.value, disambiguation: .kind(kind)))
-                }
-                continue
-            }
-            for (usr, value) in kindTree {
-                collisions.append((value: value, disambiguation: .hash(usr)))
-            }
-        }
-        return collisions
+        return Self.disambiguatedValues(for: storage, includeLanguage: includeLanguage)
     }
 
     /// Returns all values paired with their disambiguation suffixes without needing to disambiguate between two different versions of the same symbol.
@@ -205,31 +216,29 @@ extension PathHierarchy.DisambiguationContainer {
     func disambiguatedValuesWithCollapsedUniqueSymbols(includeLanguage: Bool) -> [(value: PathHierarchy.Node, disambiguation: Disambiguation)] {
         typealias DisambiguationPair = (String, String)
 
-        var uniqueSymbolIDs = [String: [DisambiguationPair]]()
-        var nonSymbols = [DisambiguationPair]()
-        for (kind, kindTree) in storage {
-            for (hash, value) in kindTree {
-                guard let symbol = value.symbol else {
-                    nonSymbols.append((kind, hash))
-                    continue
-                }
-                if symbol.identifier.interfaceLanguage == "swift" {
-                    uniqueSymbolIDs[symbol.identifier.precise, default: []].insert((kind, hash), at: 0)
-                } else {
-                    uniqueSymbolIDs[symbol.identifier.precise, default: []].append((kind, hash))
-                }
+        var uniqueSymbolIDs = [String: [Element]]()
+        var nonSymbols = [Element]()
+        for element in storage {
+            guard let symbol = element.node.symbol else {
+                nonSymbols.append(element)
+                continue
+            }
+            if symbol.identifier.interfaceLanguage == "swift" {
+                uniqueSymbolIDs[symbol.identifier.precise, default: []].insert(element, at: 0)
+            } else {
+                uniqueSymbolIDs[symbol.identifier.precise, default: []].append(element)
             }
         }
 
-        var duplicateSymbols = [String: ArraySlice<DisambiguationPair>]()
+        var duplicateSymbols = [String: ArraySlice<Element>]()
 
-        var new = Self()
-        for (kind, hash) in nonSymbols {
-            new.add(kind, hash, storage[kind]![hash]!)
+        var new = PathHierarchy.DisambiguationContainer()
+        for element in nonSymbols {
+            new.add(element.node, kind: element.kind, hash: element.hash)
         }
         for (id, symbolDisambiguations) in uniqueSymbolIDs {
-            let (kind, hash) = symbolDisambiguations[0]
-            new.add(kind, hash, storage[kind]![hash]!)
+            let element = symbolDisambiguations.first!
+            new.add(element.node, kind: element.kind, hash: element.hash)
 
             if symbolDisambiguations.count > 1 {
                 duplicateSymbols[id] = symbolDisambiguations.dropFirst()
@@ -243,8 +252,8 @@ extension PathHierarchy.DisambiguationContainer {
 
         for (id, disambiguations) in duplicateSymbols {
             let primaryDisambiguation = disambiguated.first(where: { $0.value.symbol?.identifier.precise == id })!.disambiguation
-            for (kind, hash) in disambiguations {
-                disambiguated.append((storage[kind]![hash]!, primaryDisambiguation.updated(kind: kind, hash: hash)))
+            for element in disambiguations {
+                disambiguated.append((element.node, primaryDisambiguation.updated(kind: element.kind, hash: element.hash)))
             }
         }
 
@@ -280,14 +289,14 @@ extension PathHierarchy.DisambiguationContainer {
         }
 
         /// Creates a new disambiguation with a new kind or hash value.
-        func updated(kind: String, hash: String) -> Self {
+        func updated(kind: String?, hash: String?) -> Self {
             switch self {
             case .none:
                 return .none
             case .kind:
-                return .kind(kind)
+                return kind.map { .kind($0) } ?? self
             case .hash:
-                return .hash(hash)
+                return hash.map { .hash($0) } ?? self
             }
         }
     }

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

@@ -1,7 +1,7 @@
 /*
  This source file is part of the Swift.org open source project
 
- Copyright (c) 2023 Apple Inc. and the Swift project authors
+ Copyright (c) 2023-2024 Apple Inc. and the Swift project authors
  Licensed under Apache License v2.0 with Runtime Library Exception
 
  See https://swift.org/LICENSE.txt for license information
@@ -23,17 +23,17 @@ private extension PathHierarchy.Node {
         // Each node is printed as 3-layer hierarchy with the child names, their kind disambiguation, and their hash disambiguation.
         return DumpableNode(
             name: symbol.map { "{ \($0.identifier.precise) : \($0.identifier.interfaceLanguage).\($0.kind.identifier.identifier) }" } ?? "[ \(name) ]",
-            children:
-                children.sorted(by: \.key).map { (key, disambiguationTree) -> DumpableNode in
-                DumpableNode(
+            children: children.sorted(by: \.key).map { (key, disambiguationTree) -> DumpableNode in
+                let grouped = [String: [PathHierarchy.DisambiguationContainer.Element]](grouping: disambiguationTree.storage, by: { $0.kind ?? "_" })
+                return DumpableNode(
                     name: key,
-                    children: disambiguationTree.storage.sorted(by: \.key).map { (kind, kindTree) -> DumpableNode in
+                    children: grouped.sorted(by: \.key).map { (kind, kindTree) -> DumpableNode in
                         DumpableNode(
                             name: kind,
-                            children: kindTree.sorted(by: \.key).map { (usr, node) -> DumpableNode in
+                            children: kindTree.sorted(by: { lhs, rhs in (lhs.hash ?? "_") < (rhs.hash ?? "_") }).map { (element) -> DumpableNode in
                                 DumpableNode(
-                                    name: usr,
-                                    children: [node.dumpableNode()]
+                                    name: element.hash ?? "_",
+                                    children: [element.node.dumpableNode()]
                                 )
                             }
                         )

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

@@ -410,49 +410,33 @@ extension PathHierarchy.DisambiguationContainer {
     ///  - 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? {
-        var kind: String?
-        var hash: String?
-        switch disambiguation {
-        case .kindAndHash(kind: let maybeKind, hash: let maybeHash):
-            kind = maybeKind.map(String.init)
-            hash = maybeHash.map(String.init)
-        case nil:
-            kind = nil
-            hash = nil
+        if storage.count <= 1, disambiguation == nil {
+            return storage.first?.node
         }
 
-        if let kind = kind {
-            // Need to match the provided kind
-            guard let subtree = storage[kind] else { return nil }
-            if let hash = hash {
-                return subtree[hash]
-            } else if subtree.count == 1 {
-                return subtree.values.first
-            } else {
-                // Subtree contains more than one match.
-                throw Error.lookupCollision(subtree.map { ($0.value, $0.key) })
-            }
-        } else if storage.count == 1, let subtree = storage.values.first {
-            // Tree only contains one kind subtree
-            if let hash = hash {
-                return subtree[hash]
-            } else if subtree.count == 1 {
-                return subtree.values.first
-            } else {
-                // Subtree contains more than one match.
-                throw Error.lookupCollision(subtree.map { ($0.value, $0.key) })
-            }
-        } else if let hash = hash {
-            // Need to match the provided hash
-            let kinds = storage.filter { $0.value.keys.contains(hash) }
-            if kinds.isEmpty {
-                return nil
-            } else if kinds.count == 1 {
-                return kinds.first!.value[hash]
-            } else {
-                // Subtree contains more than one match
-                throw Error.lookupCollision(kinds.map { ($0.value[hash]!, $0.key) })
+        switch disambiguation {
+        case .kindAndHash(let kind, let hash):
+            switch (kind, hash) {
+            case (let kind?, let hash?):
+                return storage.first(where: { $0.kind == kind && $0.hash == hash })?.node
+            case (let kind?, nil):
+                let matches = storage.filter({ $0.kind == kind })
+                guard matches.count <= 1 else {
+                    // Suggest not only hash disambiguation, but also type signature disambiguation.
+                    throw Error.lookupCollision(Self.disambiguatedValues(for: matches).map { ($0.value, $0.disambiguation.value()) })
+                }
+                return matches.first?.node
+            case (nil, let hash?):
+                let matches = storage.filter({ $0.hash == hash })
+                guard matches.count <= 1 else {
+                    throw Error.lookupCollision(matches.map { ($0.node, $0.kind!) }) // An element wouldn't match if it didn't have kind disambiguation.
+                }
+                return matches.first?.node
+            case (nil, nil):
+                break
             }
+        case nil:
+            break
         }
         // Disambiguate by a mix of kinds and USRs
         throw Error.lookupCollision(self.disambiguatedValues().map { ($0.value, $0.disambiguation.value()) })
@@ -461,6 +445,12 @@ extension PathHierarchy.DisambiguationContainer {
 
 // MARK: Private helper extensions
 
+// Allow optional substrings to be compared to non-optional strings
+private func == <S1: StringProtocol, S2: StringProtocol>(lhs: S1?, rhs: S2) -> Bool {
+     guard let lhs = lhs else { return false }
+     return lhs == rhs
+ }
+
 private extension Sequence {
     /// Returns the only element of the sequence that satisfies the given predicate.
     /// - Parameters:

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

@@ -45,11 +45,9 @@ extension PathHierarchy.FileRepresentation {
                         isDisfavoredInCollision: node.isDisfavoredInCollision,
                         children: node.children.values.flatMap({ tree in
                             var disambiguations = [Node.Disambiguation]()
-                            for (kind, kindTree) in tree.storage {
-                                for (hash, childNode) in kindTree where childNode.identifier != nil { // nodes without identifiers can't be found in the tree
-                                    disambiguations.append(.init(kind: kind, hash: hash, nodeID: identifierMap[childNode.identifier]!))
+                            for element in tree.storage where element.node.identifier != nil { // nodes without identifiers can't be found in the tree
+                                disambiguations.append(.init(kind: element.kind, hash: element.hash, nodeID: identifierMap[element.node.identifier]!))
                                 }
-                            }
                             return disambiguations
                         }),
                         symbolID: node.symbol?.identifier
@@ -101,7 +99,7 @@ extension PathHierarchy {
         /// The container of tutorial overview pages.
         var tutorialOverviewContainer: Int
 
-        /// A node in the
+        /// A node in the hierarchy.
         struct Node: Codable {
             var name: String
             var isDisfavoredInCollision: Bool = false