Fix topologicalSort performance on ResolvedPackage

MaxDesiatov · MaxDesiatov · commit 693b4aea7bb3 · 2023-12-14T15:37:22.000Z
diff --git a/Sources/Build/BuildPlan/BuildPlan+Test.swift b/Sources/Build/BuildPlan/BuildPlan+Test.swift
@@ -85,6 +85,7 @@ extension BuildPlan {
                     testDiscoverySrc: Sources(paths: discoveryPaths, root: discoveryDerivedDir)
                 )
                 let discoveryResolvedTarget = ResolvedTarget(
+                    packageIdentity: testProduct.packageIdentity,
                     underlying: discoveryTarget,
                     dependencies: testProduct.targets.map { .target($0, conditions: []) },
                     defaultLocalization: testProduct.defaultLocalization,
@@ -124,6 +125,7 @@ extension BuildPlan {
                     testEntryPointSources: entryPointSources
                 )
                 let entryPointResolvedTarget = ResolvedTarget(
+                    packageIdentity: testProduct.packageIdentity,
                     underlying: entryPointTarget,
                     dependencies: testProduct.targets.map { .target($0, conditions: []) } + resolvedTargetDependencies,
                     defaultLocalization: testProduct.defaultLocalization,
@@ -168,6 +170,7 @@ extension BuildPlan {
                             testEntryPointSources: entryPointResolvedTarget.underlying.sources
                         )
                         let entryPointResolvedTarget = ResolvedTarget(
+                            packageIdentity: testProduct.packageIdentity,
                             underlying: entryPointTarget,
                             dependencies: entryPointResolvedTarget.dependencies + resolvedTargetDependencies,
                             defaultLocalization: testProduct.defaultLocalization,
diff --git a/Sources/PackageGraph/PackageGraph+Loading.swift b/Sources/PackageGraph/PackageGraph+Loading.swift
@@ -373,6 +373,7 @@ private func createResolvedPackages(
         // Create target builders for each target in the package.
         let targetBuilders = package.targets.map {
             ResolvedTargetBuilder(
+                packageIdentity: package.identity,
                 target: $0,
                 observabilityScope: packageObservabilityScope,
                 platformVersionProvider: platformVersionProvider
@@ -861,6 +862,7 @@ private final class ResolvedProductBuilder: ResolvedBuilder<ResolvedProduct> {
 
     override func constructImpl() throws -> ResolvedProduct {
         return ResolvedProduct(
+            packageIdentity: packageBuilder.package.identity,
             product: product,
             targets: try targets.map{ try $0.construct() }
         )
@@ -882,6 +884,8 @@ private final class ResolvedTargetBuilder: ResolvedBuilder<ResolvedTarget> {
     /// The target reference.
     let target: Target
 
+    let packageIdentity: PackageIdentity
+
     /// The target dependencies of this target.
     var dependencies: [Dependency] = []
 
@@ -895,11 +899,13 @@ private final class ResolvedTargetBuilder: ResolvedBuilder<ResolvedTarget> {
     let platformVersionProvider: PlatformVersionProvider
 
     init(
+        packageIdentity: PackageIdentity,
         target: Target,
         observabilityScope: ObservabilityScope,
         platformVersionProvider: PlatformVersionProvider
     ) {
         self.target = target
+        self.packageIdentity = packageIdentity
         self.observabilityScope = observabilityScope
         self.platformVersionProvider = platformVersionProvider
     }
@@ -929,6 +935,7 @@ private final class ResolvedTargetBuilder: ResolvedBuilder<ResolvedTarget> {
         }
 
         return ResolvedTarget(
+            packageIdentity: self.packageIdentity,
             underlying: self.target,
             dependencies: dependencies,
             defaultLocalization: self.defaultLocalization,
diff --git a/Sources/PackageGraph/PackageGraph.swift b/Sources/PackageGraph/PackageGraph.swift
@@ -13,7 +13,7 @@
 import PackageLoading
 import PackageModel
 
-import func TSCBasic.topologicalSort
+import struct TSCBasic.OrderedSet
 
 enum PackageGraphError: Swift.Error {
     /// Indicates a non-root package with no targets.
@@ -276,3 +276,75 @@ extension PackageGraphError: CustomStringConvertible {
         }
     }
 }
+
+enum GraphError: Error {
+    /// A cycle was detected in the input.
+    case unexpectedCycle
+}
+
+/// Perform a topological sort of a graph.
+///
+/// This function is optimized for use cases where cycles are unexpected, and
+/// does not attempt to retain information on the exact nodes in the cycle.
+///
+/// - Parameters:
+///   - nodes: The list of input nodes to sort.
+///   - successors: A closure for fetching the successors of a particular node.
+///
+/// - Returns: A list of the transitive closure of nodes reachable from the
+/// inputs, ordered such that every node in the list follows all of its
+/// predecessors.
+///
+/// - Throws: GraphError.unexpectedCycle
+///
+/// - Complexity: O(v + e) where (v, e) are the number of vertices and edges
+/// reachable from the input nodes via the relation.
+private func topologicalSort<T: Identifiable>(
+    _ nodes: [T], successors: (T) throws -> [T]
+) throws -> [T] {
+    // Implements a topological sort via recursion and reverse postorder DFS.
+    func visit(
+        _ node: T,
+        _ stack: inout OrderedSet<T.ID>,
+        _ visited: inout Set<T.ID>,
+        _ result: inout [T.ID],
+        _ successors: (T) throws -> [T]
+    ) throws {
+        // Mark this node as visited -- we are done if it already was.
+        if !visited.insert(node.id).inserted {
+            return
+        }
+
+        // Otherwise, visit each adjacent node.
+        for succ in try successors(node) {
+            guard stack.append(succ.id) else {
+                // If the successor is already in this current stack, we have found a cycle.
+                //
+                // FIXME: We could easily include information on the cycle we found here.
+                throw GraphError.unexpectedCycle
+            }
+            try visit(succ, &stack, &visited, &result, successors)
+            let popped = stack.removeLast()
+            assert(popped == succ.id)
+        }
+
+        // Add to the result.
+        result.append(node.id)
+    }
+
+    let idsToNodes = Dictionary(uniqueKeysWithValues: nodes.map { ($0.id, $0) })
+
+    // FIXME: This should use a stack not recursion.
+    var visited = Set<T.ID>()
+    var result = [T.ID]()
+    var stack = OrderedSet<T.ID>()
+    for node in nodes {
+        precondition(stack.isEmpty)
+        stack.append(node.id)
+        try visit(node, &stack, &visited, &result, successors)
+        let popped = stack.removeLast()
+        assert(popped == node.id)
+    }
+
+    return result.reversed().compactMap { idsToNodes[$0] }
+}
diff --git a/Sources/PackageGraph/Resolution/ResolvedPackage.swift b/Sources/PackageGraph/Resolution/ResolvedPackage.swift
@@ -87,3 +87,9 @@ extension ResolvedPackage: CustomStringConvertible {
         return "<ResolvedPackage: \(self.identity)>"
     }
 }
+
+extension ResolvedPackage: Identifiable {
+    public var id: PackageIdentity {
+        self.underlying.identity
+    }
+}
diff --git a/Sources/PackageGraph/Resolution/ResolvedProduct.swift b/Sources/PackageGraph/Resolution/ResolvedProduct.swift
@@ -24,6 +24,8 @@ public struct ResolvedProduct: Hashable {
         self.underlying.type
     }
 
+    public let packageIdentity: PackageIdentity
+
     /// The underlying product.
     public let underlying: Product
 
@@ -61,8 +63,9 @@ public struct ResolvedProduct: Hashable {
         }
     }
 
-    public init(product: Product, targets: [ResolvedTarget]) {
+    public init(packageIdentity: PackageIdentity, product: Product, targets: [ResolvedTarget]) {
         assert(product.targets.count == targets.count && product.targets.map(\.name) == targets.map(\.name))
+        self.packageIdentity = packageIdentity
         self.underlying = product
         self.targets = targets
 
@@ -85,6 +88,7 @@ public struct ResolvedProduct: Hashable {
                 testEntryPointPath: testEntryPointPath
             )
             return ResolvedTarget(
+                packageIdentity: packageIdentity,
                 underlying: swiftTarget,
                 dependencies: targets.map { .target($0, conditions: []) },
                 defaultLocalization: defaultLocalization ?? .none, // safe since this is a derived product
diff --git a/Sources/PackageGraph/Resolution/ResolvedTarget.swift b/Sources/PackageGraph/Resolution/ResolvedTarget.swift
@@ -129,6 +129,8 @@ public struct ResolvedTarget: Hashable {
         return underlying.sources
     }
 
+    let packageIdentity: PackageIdentity
+
     /// The underlying target represented in this resolved target.
     public let underlying: Target
 
@@ -148,12 +150,14 @@ public struct ResolvedTarget: Hashable {
 
     /// Create a resolved target instance.
     public init(
+        packageIdentity: PackageIdentity,
         underlying: Target,
         dependencies: [ResolvedTarget.Dependency],
         defaultLocalization: String? = nil,
         supportedPlatforms: [SupportedPlatform],
         platformVersionProvider: PlatformVersionProvider
     ) {
+        self.packageIdentity = packageIdentity
         self.underlying = underlying
         self.dependencies = dependencies
         self.defaultLocalization = defaultLocalization
@@ -190,3 +194,25 @@ extension ResolvedTarget.Dependency: CustomStringConvertible {
         return str
     }
 }
+
+extension ResolvedTarget.Dependency: Identifiable {
+    public struct ID: Hashable {
+        enum Kind: Hashable {
+            case target
+            case product
+        }
+
+        let kind: Kind
+        let packageIdentity: PackageIdentity
+        let name: String
+    }
+
+    public var id: ID {
+        switch self {
+        case .target(let target, _):
+            return .init(kind: .target, packageIdentity: target.packageIdentity, name: target.name)
+        case .product(let product, _):
+            return .init(kind: .product, packageIdentity: product.packageIdentity, name: product.name)
+        }
+    }
+}
diff --git a/Tests/PackageGraphPerformanceTests/PackageGraphPerfTests.swift b/Tests/PackageGraphPerformanceTests/PackageGraphPerfTests.swift
@@ -126,10 +126,12 @@ final class PackageGraphPerfTests: XCTestCasePerf {
                     )
                 ],
                 targets: [
-                    try .init(name: "Target\(sequenceNumber)",
-                              dependencies: dependencySequence.map {
-                                  .product(name: "Target\($0)", package: "Package\($0)")
-                              })
+                    try .init(
+                        name: "Target\(sequenceNumber)",
+                        dependencies: dependencySequence.map {
+                            .product(name: "Target\($0)", package: "Package\($0)")
+                        }
+                    )
                 ]
             )
         }
diff --git a/Tests/PackageGraphTests/TargetTests.swift b/Tests/PackageGraphTests/TargetTests.swift
@@ -13,11 +13,12 @@
 import XCTest
 
 import PackageGraph
-import PackageModel
+@testable import PackageModel
 
 private extension ResolvedTarget {
-    init(name: String, deps: ResolvedTarget...) {
+    init(packageIdentity: String, name: String, deps: ResolvedTarget...) {
         self.init(
+            packageIdentity: .init(packageIdentity),
             underlying: SwiftTarget(
                 name: name,
                 type: .library,
@@ -48,9 +49,9 @@ class TargetDependencyTests: XCTestCase {
 
     func test1() throws {
         testTargets {
-            let t1 = ResolvedTarget(name: "t1")
-            let t2 = ResolvedTarget(name: "t2", deps: t1)
-            let t3 = ResolvedTarget(name: "t3", deps: t2)
+            let t1 = ResolvedTarget(packageIdentity: "pkg", name: "t1")
+            let t2 = ResolvedTarget(packageIdentity: "pkg", name: "t2", deps: t1)
+            let t3 = ResolvedTarget(packageIdentity: "pkg", name: "t3", deps: t2)
 
             XCTAssertEqual(try t3.recursiveTargetDependencies(), [t2, t1])
             XCTAssertEqual(try t2.recursiveTargetDependencies(), [t1])
@@ -59,10 +60,10 @@ class TargetDependencyTests: XCTestCase {
 
     func test2() throws {
         testTargets {
-            let t1 = ResolvedTarget(name: "t1")
-            let t2 = ResolvedTarget(name: "t2", deps: t1)
-            let t3 = ResolvedTarget(name: "t3", deps: t2, t1)
-            let t4 = ResolvedTarget(name: "t4", deps: t2, t3, t1)
+            let t1 = ResolvedTarget(packageIdentity: "pkg", name: "t1")
+            let t2 = ResolvedTarget(packageIdentity: "pkg", name: "t2", deps: t1)
+            let t3 = ResolvedTarget(packageIdentity: "pkg", name: "t3", deps: t2, t1)
+            let t4 = ResolvedTarget(packageIdentity: "pkg", name: "t4", deps: t2, t3, t1)
 
             XCTAssertEqual(try t4.recursiveTargetDependencies(), [t3, t2, t1])
             XCTAssertEqual(try t3.recursiveTargetDependencies(), [t2, t1])
@@ -72,10 +73,10 @@ class TargetDependencyTests: XCTestCase {
 
     func test3() throws {
         testTargets {
-            let t1 = ResolvedTarget(name: "t1")
-            let t2 = ResolvedTarget(name: "t2", deps: t1)
-            let t3 = ResolvedTarget(name: "t3", deps: t2, t1)
-            let t4 = ResolvedTarget(name: "t4", deps: t1, t2, t3)
+            let t1 = ResolvedTarget(packageIdentity: "pkg", name: "t1")
+            let t2 = ResolvedTarget(packageIdentity: "pkg", name: "t2", deps: t1)
+            let t3 = ResolvedTarget(packageIdentity: "pkg", name: "t3", deps: t2, t1)
+            let t4 = ResolvedTarget(packageIdentity: "pkg", name: "t4", deps: t1, t2, t3)
 
             XCTAssertEqual(try t4.recursiveTargetDependencies(), [t3, t2, t1])
             XCTAssertEqual(try t3.recursiveTargetDependencies(), [t2, t1])
@@ -85,10 +86,10 @@ class TargetDependencyTests: XCTestCase {
 
     func test4() throws {
         testTargets {
-            let t1 = ResolvedTarget(name: "t1")
-            let t2 = ResolvedTarget(name: "t2", deps: t1)
-            let t3 = ResolvedTarget(name: "t3", deps: t2)
-            let t4 = ResolvedTarget(name: "t4", deps: t3)
+            let t1 = ResolvedTarget(packageIdentity: "pkg", name: "t1")
+            let t2 = ResolvedTarget(packageIdentity: "pkg", name: "t2", deps: t1)
+            let t3 = ResolvedTarget(packageIdentity: "pkg", name: "t3", deps: t2)
+            let t4 = ResolvedTarget(packageIdentity: "pkg", name: "t4", deps: t3)
 
             XCTAssertEqual(try t4.recursiveTargetDependencies(), [t3, t2, t1])
             XCTAssertEqual(try t3.recursiveTargetDependencies(), [t2, t1])
@@ -98,12 +99,12 @@ class TargetDependencyTests: XCTestCase {
 
     func test5() throws {
         testTargets {
-            let t1 = ResolvedTarget(name: "t1")
-            let t2 = ResolvedTarget(name: "t2", deps: t1)
-            let t3 = ResolvedTarget(name: "t3", deps: t2)
-            let t4 = ResolvedTarget(name: "t4", deps: t3)
-            let t5 = ResolvedTarget(name: "t5", deps: t2)
-            let t6 = ResolvedTarget(name: "t6", deps: t5, t4)
+            let t1 = ResolvedTarget(packageIdentity: "pkg", name: "t1")
+            let t2 = ResolvedTarget(packageIdentity: "pkg", name: "t2", deps: t1)
+            let t3 = ResolvedTarget(packageIdentity: "pkg", name: "t3", deps: t2)
+            let t4 = ResolvedTarget(packageIdentity: "pkg", name: "t4", deps: t3)
+            let t5 = ResolvedTarget(packageIdentity: "pkg", name: "t5", deps: t2)
+            let t6 = ResolvedTarget(packageIdentity: "pkg", name: "t6", deps: t5, t4)
 
             // precise order is not important, but it is important that the following are true
             let t6rd = try t6.recursiveTargetDependencies()
@@ -122,12 +123,12 @@ class TargetDependencyTests: XCTestCase {
 
     func test6() throws {
         testTargets {
-            let t1 = ResolvedTarget(name: "t1")
-            let t2 = ResolvedTarget(name: "t2", deps: t1)
-            let t3 = ResolvedTarget(name: "t3", deps: t2)
-            let t4 = ResolvedTarget(name: "t4", deps: t3)
-            let t5 = ResolvedTarget(name: "t5", deps: t2)
-            let t6 = ResolvedTarget(name: "t6", deps: t4, t5) // same as above, but these two swapped
+            let t1 = ResolvedTarget(packageIdentity: "pkg", name: "t1")
+            let t2 = ResolvedTarget(packageIdentity: "pkg", name: "t2", deps: t1)
+            let t3 = ResolvedTarget(packageIdentity: "pkg", name: "t3", deps: t2)
+            let t4 = ResolvedTarget(packageIdentity: "pkg", name: "t4", deps: t3)
+            let t5 = ResolvedTarget(packageIdentity: "pkg", name: "t5", deps: t2)
+            let t6 = ResolvedTarget(packageIdentity: "pkg", name: "t6", deps: t4, t5) // same as above, but these two swapped
 
             // precise order is not important, but it is important that the following are true
             let t6rd = try t6.recursiveTargetDependencies()

Original file line number	Diff line number	Diff line change
`@@ -87,3 +87,9 @@ extension ResolvedPackage: CustomStringConvertible {`
`87`	`87`	`return "<ResolvedPackage: \(self.identity)>"`
`88`	`88`	`}`
`89`	`89`	`}`
	`90`	`+`
	`91`	`+extension ResolvedPackage: Identifiable {`
	`92`	`+ public var id: PackageIdentity {`
	`93`	`+ self.underlying.identity`
	`94`	`+ }`
	`95`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,8 @@ public struct ResolvedProduct: Hashable {`
`24`	`24`	`self.underlying.type`
`25`	`25`	`}`
`26`	`26`
	`27`	`+ public let packageIdentity: PackageIdentity`
	`28`	`+`
`27`	`29`	`/// The underlying product.`
`28`	`30`	`public let underlying: Product`
`29`	`31`
`@@ -61,8 +63,9 @@ public struct ResolvedProduct: Hashable {`
`61`	`63`	`}`
`62`	`64`	`}`
`63`	`65`
`64`		`- public init(product: Product, targets: [ResolvedTarget]) {`
	`66`	`+ public init(packageIdentity: PackageIdentity, product: Product, targets: [ResolvedTarget]) {`
`65`	`67`	`assert(product.targets.count == targets.count && product.targets.map(\.name) == targets.map(\.name))`
	`68`	`+ self.packageIdentity = packageIdentity`
`66`	`69`	`self.underlying = product`
`67`	`70`	`self.targets = targets`
`68`	`71`
`@@ -85,6 +88,7 @@ public struct ResolvedProduct: Hashable {`
`85`	`88`	`testEntryPointPath: testEntryPointPath`
`86`	`89`	`)`
`87`	`90`	`return ResolvedTarget(`
	`91`	`+ packageIdentity: packageIdentity,`
`88`	`92`	`underlying: swiftTarget,`
`89`	`93`	`dependencies: targets.map { .target($0, conditions: []) },`
`90`	`94`	`defaultLocalization: defaultLocalization ?? .none, // safe since this is a derived product`
Original file line number	Diff line number	Diff line change
`@@ -126,10 +126,12 @@ final class PackageGraphPerfTests: XCTestCasePerf {`
`126`	`126`	`)`
`127`	`127`	`],`
`128`	`128`	`targets: [`
`129`		`- try .init(name: "Target\(sequenceNumber)",`
`130`		`- dependencies: dependencySequence.map {`
`131`		`- .product(name: "Target\($0)", package: "Package\($0)")`
`132`		`- })`
	`129`	`+ try .init(`
	`130`	`+ name: "Target\(sequenceNumber)",`
	`131`	`+ dependencies: dependencySequence.map {`
	`132`	`+ .product(name: "Target\($0)", package: "Package\($0)")`
	`133`	`+ }`
	`134`	`+ )`
`133`	`135`	`]`
`134`	`136`	`)`
`135`	`137`	`}`