use SymbolKit's overload groups implementation to group overloads (#846)

rdar://116550528

Co-authored-by: David Rönnqvist <[email protected]>

Author: QuietMisdreavus

Package.resolved

@@ -1156,6 +1156,8 @@ public class DocumentationContext: DocumentationContextDataProviderDelegate {
 
                     // For inherited symbols we remove the source docs (if inheriting docs is disabled) before creating their documentation nodes.
                     for (_, relationships) in unifiedSymbolGraph.relationshipsByLanguage {
+                        var overloadGroups = [String: [String]]()
+
                         for relationship in relationships {
                             // Check for an origin key.
                             if let sourceOrigin = relationship[mixin: SymbolGraph.Relationship.SourceOrigin.self],
@@ -1169,8 +1171,13 @@ public class DocumentationContext: DocumentationContextDataProviderDelegate {
                                     localCache: documentationCache,
                                     moduleName: moduleName
                                 )
+                            } else if relationship.kind == .overloadOf {
+                                // An 'overloadOf' relationship points from symbol -> group
+                                overloadGroups[relationship.target, default: []].append(relationship.source)
                             }
                         }
+
+                        addOverloadGroupReferences(overloadGroups: overloadGroups)
                     }
 
                     if let rootURL = symbolGraphLoader.mainModuleURL(forModule: moduleName), let rootModule = unifiedSymbolGraph.moduleData[rootURL] {
@@ -1285,8 +1292,6 @@ public class DocumentationContext: DocumentationContextDataProviderDelegate {
             }
             emitWarningsForSymbolsMatchedInMultipleDocumentationExtensions(with: symbolsWithMultipleDocumentationExtensionMatches)
             symbolsWithMultipleDocumentationExtensionMatches.removeAll()
-
-            try groupOverloadedSymbols(with: linkResolver.localResolver)
 
             // Create inherited API collections
             try GeneratedDocumentationTopics.createInheritedSymbolsAPICollections(
@@ -2330,34 +2335,78 @@ public class DocumentationContext: DocumentationContextDataProviderDelegate {
         return automaticallyCuratedSymbols
     }
 
-    /// Handles overloaded symbols by grouping them together into one page.
-    private func groupOverloadedSymbols(with linkResolver: PathHierarchyBasedLinkResolver) throws {
+    private func addOverloadGroupReferences(overloadGroups: [String: [String]]) {
         guard FeatureFlags.current.isExperimentalOverloadedSymbolPresentationEnabled else {
             return
         }
 
-        try linkResolver.traverseOverloadedSymbols { overloadedSymbolReferences in
+        for (overloadGroupID, overloadSymbolIDs) in overloadGroups {
+            guard overloadSymbolIDs.count > 1 else {
+                assertionFailure("Overload group \(overloadGroupID) contained \(overloadSymbolIDs.count) symbols, but should have more than one symbol to be valid.")
+                continue
+            }
+            guard let overloadGroupNode = documentationCache[overloadGroupID] else {
+                preconditionFailure("Overload group \(overloadGroupID) doesn't have a local entity")
+            }
 
-            // Tell each symbol what other symbols overload it.
-            for (index, symbolReference) in overloadedSymbolReferences.indexed() {
-                let documentationNode = try entity(with: symbolReference)
+            var overloadSymbolNodes = overloadSymbolIDs.map {
+                guard let node = documentationCache[$0] else {
+                    preconditionFailure("Overloaded symbol \($0) doesn't have a local entity")
+                }
+                return node
+            }
+            if overloadSymbolNodes.allSatisfy({ node in
+                (node.semantic as? Symbol)?.overloadsVariants.firstValue != nil
+            }) {
+                // If SymbolKit saved its sort of the symbols, use that ordering here
+                overloadSymbolNodes.sort(by: { lhs, rhs in
+                    let lhsIndex = (lhs.semantic as! Symbol).overloadsVariants.firstValue!.displayIndex
+                    let rhsIndex = (rhs.semantic as! Symbol).overloadsVariants.firstValue!.displayIndex
+                    return lhsIndex < rhsIndex
+                })
+            } else {
+                assertionFailure("""
+                    Overload group \(overloadGroupNode.reference.absoluteString.singleQuoted) was not properly initialized with overload data from SymbolKit.
+                    Symbols without overload data: \(Array(overloadSymbolNodes.filter({ ($0.semantic as? Symbol)?.overloadsVariants.firstValue == nil }).map(\.reference.absoluteString.singleQuoted)))
+                    """)
+                return
+            }
 
+            func addOverloadReferences(
+                to documentationNode: DocumentationNode,
+                at index: Int,
+                overloadSymbolReferences: [DocumentationNode]
+            ) {
                 guard let symbol = documentationNode.semantic as? Symbol else {
                     preconditionFailure("""
-                    Only symbols can be overloads. Found non-symbol overload for \(symbolReference.absoluteString.singleQuoted).
-                    Non-symbols should already have been filtered out in `PathHierarchyBasedLinkResolver.traverseOverloadedSymbols(_:)`.
+                    Only symbols can be overloads. Found non-symbol overload for \(documentationNode.reference.absoluteString.singleQuoted).
+                    Non-symbols should already have been filtered out by SymbolKit.
                     """)
                 }
-                guard symbolReference.sourceLanguage == .swift else {
-                    assertionFailure("Overload groups is only supported for Swift symbols.")
-                    continue
+                guard documentationNode.reference.sourceLanguage == .swift else {
+                    assertionFailure("""
+                    Overload groups are only supported for Swift symbols.
+                    The symbol at \(documentationNode.reference.absoluteString.singleQuoted) is listed as \(documentationNode.reference.sourceLanguage.name).
+                    """)
+                    return
                 }
 
-                var otherOverloadedSymbolReferences = overloadedSymbolReferences
+                var otherOverloadedSymbolReferences = overloadSymbolReferences.map(\.reference)
                 otherOverloadedSymbolReferences.remove(at: index)
-                let overloads = Symbol.Overloads(references: otherOverloadedSymbolReferences, displayIndex: index)
+                let displayIndex = symbol.overloadsVariants.firstValue?.displayIndex ?? index
+                let overloads = Symbol.Overloads(references: otherOverloadedSymbolReferences, displayIndex: displayIndex)
                 symbol.overloadsVariants = .init(swiftVariant: overloads)
             }
+
+            // The overload group node itself is a clone of the first symbol, so the code above can
+            // swap out the first element in the overload references to create the alternate
+            // declarations section properly. However, it is also a distinct symbol, node, and page,
+            // so the first overload itself should also be handled separately in the loop below.
+            addOverloadReferences(to: overloadGroupNode, at: 0, overloadSymbolReferences: overloadSymbolNodes)
+
+            for (index, node) in overloadSymbolNodes.indexed() {
+                addOverloadReferences(to: node, at: index, overloadSymbolReferences: overloadSymbolNodes)
+            }
         }
     }
 

Sources/SwiftDocC/Infrastructure/DocumentationContext.swift

@@ -136,7 +136,14 @@ struct PathHierarchy {
                     allNodes[id, default: []].append(node)
                 }
             }
-            
+
+            for relationship in graph.relationships where relationship.kind == .overloadOf {
+                // An 'overloadOf' relationship points from symbol -> group. We want to disfavor the
+                // individual overload symbols in favor of resolving links to their overload group
+                // symbol.
+                nodes[relationship.source]?.specialBehaviors.insert(.disfavorInLinkCollision)
+            }
+
             // If there are multiple symbol graphs (for example for different source languages or platforms) then the nodes may have already been added to the hierarchy.
             var topLevelCandidates = nodes.filter { _, node in node.parent == nil }
             for relationship in graph.relationships where relationship.kind.formsHierarchy {
@@ -192,7 +199,7 @@ struct PathHierarchy {
                 }
                 topLevelCandidates.removeValue(forKey: relationship.source)
             }
-            
+
             // The hierarchy doesn't contain any non-symbol nodes yet. It's OK to unwrap the `symbol` property.
             for topLevelNode in topLevelCandidates.values where topLevelNode.symbol!.pathComponents.count == 1 {
                 moduleNode.add(symbolChild: topLevelNode)
@@ -520,23 +527,6 @@ extension PathHierarchy {
         }
         return Array(result) + modules.map { $0.identifier }
     }
-
-    func traverseOverloadedSymbolGroups(observe: (_ overloadedSymbols: [ResolvedIdentifier]) throws -> Void) rethrows {
-        for node in lookup.values where node.symbol != nil {
-            for disambiguation in node.children.values {
-                var overloadGroups = [SymbolGraph.Symbol.KindIdentifier: [ResolvedIdentifier]]()
-                for element in disambiguation.storage {
-                    guard let kind = element.node.symbol?.kind.identifier, kind.isOverloadableKind else {
-                        continue
-                    }
-                    overloadGroups[kind, default: []].append(element.node.identifier)
-                }
-                for overloads in overloadGroups.values where overloads.count > 1 {
-                    try observe(overloads)
-                }
-            }
-        }
-    }
 }
 
 // MARK: Removing nodes

Sources/SwiftDocC/Infrastructure/Extensions/KindIdentifier+Overloads.swift

@@ -56,7 +56,7 @@ final class PathHierarchyBasedLinkResolver {
             observe(reference, parentReference)
         }
     }
-    
+
     /// Returns the direct descendants of the given page that match the given source language filter.
     ///
     /// A descendant is included if it has a language representation in at least one of the languages in the given language filter or if the language filter is empty.
@@ -92,14 +92,7 @@ final class PathHierarchyBasedLinkResolver {
         }
         return results
     }
-    
-    /// Traverse all symbols of the same kind that have collisions.
-    func traverseOverloadedSymbols(_ observe: (_ overloadedSymbols: [ResolvedTopicReference]) throws -> Void) rethrows {
-        try pathHierarchy.traverseOverloadedSymbolGroups() { overloadedSymbols in
-            try observe(overloadedSymbols.map { resolvedReferenceMap[$0]! })
-        }
-    }
-    
+
     /// Returns a list of all the top level symbols.
     func topLevelSymbols() -> [ResolvedTopicReference] {
         return pathHierarchy.topLevelSymbols().map { resolvedReferenceMap[$0]! }

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

@@ -78,6 +78,10 @@ struct SymbolGraphLoader {
 
                 configureSymbolGraph?(&symbolGraph)
 
+                if FeatureFlags.current.isExperimentalOverloadedSymbolPresentationEnabled {
+                    symbolGraph.createOverloadGroupSymbols()
+                }
+
                 let (moduleName, isMainSymbolGraph) = Self.moduleNameFor(symbolGraph, at: symbolGraphURL)
                 // If the bundle provides availability defaults add symbol availability data.
                 self.addDefaultAvailability(to: &symbolGraph, moduleName: moduleName)

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

@@ -223,7 +223,17 @@ public struct DocumentationNode {
             }
             return nil
         }
-        
+
+        let overloadVariants = DocumentationDataVariants(
+            symbolData: unifiedSymbol.mixins,
+            platformName: platformName
+        ) { mixins -> Symbol.Overloads? in
+            guard let overloadData = mixins[SymbolGraph.Symbol.OverloadData.mixinKey] as? SymbolGraph.Symbol.OverloadData else {
+                return nil
+            }
+            return .init(references: [], displayIndex: overloadData.overloadGroupIndex)
+        }
+
         var languages = Set([reference.sourceLanguage])
         var operatingSystemName = platformName.map({ Set([$0]) }) ?? []
 
@@ -300,7 +310,8 @@ public struct DocumentationNode {
             httpParametersSectionVariants: .empty,
             httpResponsesSectionVariants: .empty,
             redirectsVariants: .empty,
-            crossImportOverlayModule: moduleData.bystanders.map({ (moduleData.name, $0) })
+            crossImportOverlayModule: moduleData.bystanders.map({ (moduleData.name, $0) }),
+            overloadsVariants: overloadVariants
         )
 
         try! semanticSymbol.mergeDeclarations(unifiedSymbol: unifiedSymbol)

Sources/SwiftDocC/Infrastructure/Symbol Graph/SymbolGraphLoader.swift

@@ -3938,6 +3938,7 @@ let expected = """
 
             let symbolReference = try XCTUnwrap(context.documentationCache.reference(symbolID: "s:12Minimal_docs4TestV"))
 
+
             // Resolve from various locations in the bundle
             for parent in [bundle.rootReference, bundle.documentationRootReference, bundle.tutorialsRootReference, symbolReference] {
                 switch context.resolve(unresolved, in: parent) {
@@ -4356,14 +4357,31 @@ let expected = """
                 ))
             ])
         ])
-        let (_, _, context) = try loadBundle(from: tempURL)
-        
+        let (_, bundle, context) = try loadBundle(from: tempURL)
+        let moduleReference = ResolvedTopicReference(bundleIdentifier: bundle.identifier, path: "/documentation/ModuleName", sourceLanguage: .swift)
+
         for kindID in overloadableKindIDs {
             var seenIndices = Set<Int>()
-            // Find the 4 symbols of this specific kind
-            let overloadedReferences = try symbols.filter { $0.kind.identifier == kindID }
+            // Find the 4 symbols of this specific kind. SymbolKit will have assigned a display
+            // index based on their sorted USRs, so sort them ahead of time based on that
+            let overloadedReferences = try symbols.filter { $0.kind.identifier == kindID }.sorted(by: \.identifier.precise)
                 .map { try XCTUnwrap(context.documentationCache.reference(symbolID: $0.identifier.precise)) }
-            
+
+            let overloadGroupNode: DocumentationNode
+            let overloadGroupSymbol: Symbol
+            let overloadGroupReferences: Symbol.Overloads
+            switch context.resolve(.unresolved(.init(topicURL: .init(symbolPath: "SymbolName-\(kindID.identifier)"))), in: moduleReference, fromSymbolLink: true) {
+            case let .failure(_, errorMessage):
+                XCTFail("Could not resolve overload group page for \(kindID.identifier). Error message: \(errorMessage)")
+                continue
+            case let .success(overloadGroupReference):
+                overloadGroupNode = try context.entity(with: overloadGroupReference)
+                overloadGroupSymbol = try XCTUnwrap(overloadGroupNode.semantic as? Symbol)
+                overloadGroupReferences = try XCTUnwrap(overloadGroupSymbol.overloadsVariants.firstValue)
+
+                XCTAssertEqual(overloadGroupReferences.displayIndex, 0)
+            }
+
             // Check that each symbol lists the other 3 overloads
             for (index, reference) in overloadedReferences.indexed() {
                 let overloadedDocumentationNode = try XCTUnwrap(context.documentationCache[reference])
@@ -4378,9 +4396,18 @@ let expected = """
                 }
 
                 // Each symbol needs to tell the renderer where it belongs in the array of overloaded declarations.
-                let displayIndex = try XCTUnwrap(overloads.displayIndex)
-                XCTAssertFalse(seenIndices.contains(displayIndex))
-                seenIndices.insert(displayIndex)
+                XCTAssertFalse(seenIndices.contains(overloads.displayIndex))
+                XCTAssertEqual(overloads.displayIndex, index)
+                seenIndices.insert(overloads.displayIndex)
+
+                if overloads.displayIndex == 0 {
+                    // The first declaration in the display list should be the same declaration as
+                    // the overload group page
+                    XCTAssertEqual(overloadedSymbol.declaration.first?.value.declarationFragments, overloadGroupSymbol.declaration.first?.value.declarationFragments)
+                } else {
+                    // Otherwise, this reference should also be referenced by the overload group
+                    XCTAssert(overloadGroupReferences.references.contains(reference))
+                }
             }
             // Check that all the overloads was encountered
             for index in overloadedReferences.indices {
@@ -4425,13 +4452,17 @@ let expected = """
             }
         }
     }
-    
+
     // A test helper that creates a symbol with a given identifier and kind.
-    private func makeSymbol(identifier: String, kind: SymbolGraph.Symbol.KindIdentifier) -> SymbolGraph.Symbol {
+    private func makeSymbol(
+        name: String = "SymbolName",
+        identifier: String,
+        kind: SymbolGraph.Symbol.KindIdentifier
+    ) -> SymbolGraph.Symbol {
         return SymbolGraph.Symbol(
             identifier: .init(precise: identifier, interfaceLanguage: SourceLanguage.swift.id),
-            names: .init(title: "SymbolName", navigator: nil, subHeading: nil, prose: nil),
-            pathComponents: ["SymbolName"],
+            names: .init(title: name, navigator: nil, subHeading: nil, prose: nil),
+            pathComponents: [name],
             docComment: nil,
             accessLevel: .public,
             kind: .init(parsedIdentifier: kind, displayName: "Kind Display Name"),