Improve the performance of ProtocolConformanceReferenceBuilder through parallelism

Author: amorde

Sources/Indexer/SourceFileCollector.swift

@@ -4,6 +4,7 @@ import Logger
 import SourceGraph
 import SwiftIndexStore
 import SystemPackage
+import Shared
 
 public struct SourceFileCollector {
     private let indexStorePaths: Set<FilePath>

Sources/Indexer/JobPool.swift renamed to Sources/Shared/JobPool.swift

@@ -1,10 +1,13 @@
 import Foundation
-import Shared
 
-struct JobPool<Job> {
+public struct JobPool<Job> {
     let jobs: [Job]
 
-    func forEach(_ block: @escaping (Job) throws -> Void) throws {
+    public init(jobs: [Job]) {
+        self.jobs = jobs
+    }
+
+    public func forEach(_ block: @escaping (Job) throws -> Void) throws {
         var error: Error?
 
         DispatchQueue.concurrentPerform(iterations: jobs.count) { idx in
@@ -23,7 +26,8 @@ struct JobPool<Job> {
         }
     }
 
-    func flatMap<Result>(_ block: @escaping (Job) throws -> [Result]) throws -> [Result] {
+    /// Throwing variant
+    public func flatMap<Result>(_ block: @escaping (Job) throws -> [Result]) throws -> [Result] {
         var error: Error?
         var results: [Result] = []
         let lock = UnfairLock()
@@ -49,4 +53,21 @@ struct JobPool<Job> {
 
         return results
     }
+
+    /// Non-throwing variant
+    public func flatMap<Result>(_ block: @escaping (Job) -> [Result]) -> [Result] {
+        var results: [Result] = []
+        let lock = UnfairLock()
+
+        DispatchQueue.concurrentPerform(iterations: jobs.count) { idx in
+            let job = jobs[idx]
+            let result = block(job)
+
+            lock.perform {
+                results.append(contentsOf: result)
+            }
+        }
+
+        return results
+    }
 }

Sources/SourceGraph/Mutators/ProtocolConformanceReferenceBuilder.swift

@@ -19,10 +19,10 @@ final class ProtocolConformanceReferenceBuilder: SourceGraphMutator {
     // MARK: - Private
 
     private func referenceConformingDeclarationsImplementedInSuperclass() -> Set<Reference> {
-        var newReferences: Set<Reference> = []
         let protocols = graph.declarations(ofKind: .protocol)
 
-        for proto in protocols {
+        let newReferences = Set(JobPool(jobs: Array(protocols)).flatMap { [graph] proto in
+            var result: [Reference] = []
             // Find all classes that implement this protocol.
             let conformingClasses = graph.references(to: proto)
                 .reduce(into: Set<Declaration>()) { result, ref in
@@ -66,14 +66,16 @@ final class ProtocolConformanceReferenceBuilder: SourceGraphMutator {
                                 )
                                 reference.name = declInSuperclass.name
                                 reference.parent = unimplementedProtoDecl
-                                graph.add(reference, from: unimplementedProtoDecl)
-                                newReferences.insert(reference)
+                                result.append(reference)
                             }
                         }
                     }
                 }
             }
-        }
+            return result
+        })
+        // Perform mutations on the graph based on the calculated references
+        newReferences.forEach { graph.add($0, from: $0.parent!) }
 
         return newReferences
     }