[BuildPlan] Add a way to traverse the graph and compute destinations for all modules

xedin · xedin · commit 13b861c07a26 · 2024-08-08T22:03:34.000-07:00
This is going to be used by sourcekit-lsp to build a graph/dictionary
of all the targets and depths they appear during package loading and
replaces the need for topological sort and `buildTriple`.
diff --git a/Sources/Basics/Graph/GraphAlgorithms.swift b/Sources/Basics/Graph/GraphAlgorithms.swift
@@ -86,3 +86,50 @@ public func depthFirstSearch<T: Hashable>(
         }
     }
 }
+
+private struct TraversalNode<T: Hashable>: Hashable {
+    let parent: T?
+    let curr: T
+    let depth: Int
+}
+
+/// Implements a pre-order depth-first search that traverses the whole graph and
+/// doesn't distinguish between unique and duplicate nodes. The method expects
+/// the graph to be acyclic but doesn't check that.
+///
+/// - Parameters:
+///   - nodes: The list of input nodes to sort.
+///   - successors: A closure for fetching the successors of a particular node.
+///   - onNext: A callback to indicate the node currently being processed
+///             including its parent (if any) and its depth.
+///
+/// - Complexity: O(v + e) where (v, e) are the number of vertices and edges
+/// reachable from the input nodes via the relation.
+public func depthFirstSearch<T: Hashable>(
+    _ nodes: [T],
+    successors: (T) throws -> [T],
+    onNext: (T, _ parent: T?, _ depth: Int) throws -> Void
+) rethrows {
+    var stack = OrderedSet<TraversalNode<T>>()
+
+    for node in nodes {
+        precondition(stack.isEmpty)
+        stack.append(TraversalNode(parent: nil, curr: node, depth: 0))
+
+        while !stack.isEmpty {
+            let node = stack.removeLast()
+
+            try onNext(node.curr, node.parent, node.depth)
+
+            for succ in try successors(node.curr) {
+                stack.append(
+                    TraversalNode(
+                        parent: node.curr,
+                        curr: succ,
+                        depth: node.depth + 1
+                    )
+                )
+            }
+        }
+    }
+}
diff --git a/Sources/Build/BuildPlan/BuildPlan.swift b/Sources/Build/BuildPlan/BuildPlan.swift
@@ -904,6 +904,66 @@ extension BuildPlan {
 extension BuildPlan {
     fileprivate typealias Destination = BuildParameters.Destination
 
+    fileprivate enum TraversalNode: Hashable {
+        case package(ResolvedPackage)
+        case product(ResolvedProduct, BuildParameters.Destination)
+        case module(ResolvedModule, BuildParameters.Destination)
+
+        var destination: BuildParameters.Destination {
+            switch self {
+            case .package:
+                .target
+            case .product(_, let destination):
+                destination
+            case .module(_, let destination):
+                destination
+            }
+        }
+
+        static func `for`(
+            product: ResolvedProduct,
+            context destination: BuildParameters.Destination
+        ) -> Self {
+            switch product.type {
+            case .macro, .plugin:
+                .product(product, .host)
+            case .test:
+                .product(product, product.modules.contains(where: self.hasMacroDependency) ? .host : destination)
+            default:
+                .product(product, destination)
+            }
+        }
+
+        static func `for`(
+            module: ResolvedModule,
+            context destination: BuildParameters.Destination
+        ) -> Self {
+            switch module.type {
+            case .macro, .plugin:
+                // Macros and plugins are ways built for host
+                .module(module, .host)
+            case .test:
+                .module(module, self.hasMacroDependency(module: module) ? .host : destination)
+            default:
+                // By default assume the destination of the context.
+                // This means that i.e. test products that reference macros
+                // would force all of their successors to be `host`
+                .module(module, destination)
+            }
+        }
+
+        static func hasMacroDependency(module: ResolvedModule) -> Bool {
+            module.dependencies.contains(where: {
+                switch $0 {
+                case .product(let productDependency, _):
+                    productDependency.type == .macro
+                case .module(let moduleDependency, _):
+                    moduleDependency.type == .macro
+                }
+            })
+        }
+    }
+
     /// Traverse the modules graph and find a destination for every product and module.
     /// All non-macro/plugin products and modules have `target` destination with one
     /// notable exception - test products/modules with direct macro dependency.
@@ -912,57 +972,8 @@ extension BuildPlan {
         onProduct: (ResolvedProduct, Destination) throws -> Void,
         onModule: (ResolvedModule, Destination) async throws -> Void
     ) async rethrows {
-        enum Node: Hashable {
-            case package(ResolvedPackage)
-            case product(ResolvedProduct, Destination)
-            case module(ResolvedModule, Destination)
-
-            static func `for`(
-                product: ResolvedProduct,
-                context destination: Destination
-            ) -> Node {
-                switch product.type {
-                case .macro, .plugin:
-                    .product(product, .host)
-                case .test:
-                    .product(product, product.modules.contains(where: self.hasMacroDependency) ? .host : destination)
-                default:
-                    .product(product, destination)
-                }
-            }
-
-            static func `for`(
-                module: ResolvedModule,
-                context destination: Destination
-            ) -> Node {
-                switch module.type {
-                case .macro, .plugin:
-                    // Macros and plugins are ways built for host
-                    .module(module, .host)
-                case .test:
-                    .module(module, self.hasMacroDependency(module: module) ? .host : destination)
-                default:
-                    // By default assume the destination of the context.
-                    // This means that i.e. test products that reference macros
-                    // would force all of their successors to be `host`
-                    .module(module, destination)
-                }
-            }
-
-            static func hasMacroDependency(module: ResolvedModule) -> Bool {
-                module.dependencies.contains(where: {
-                    switch $0 {
-                    case .product(let productDependency, _):
-                        productDependency.type == .macro
-                    case .module(let moduleDependency, _):
-                        moduleDependency.type == .macro
-                    }
-                })
-            }
-        }
-
-        func successors(for package: ResolvedPackage) -> [Node] {
-            var successors: [Node] = []
+        func successors(for package: ResolvedPackage) -> [TraversalNode] {
+            var successors: [TraversalNode] = []
             for product in package.products {
                 if case .test = product.underlying.type,
                    !graph.rootPackages.contains(id: package.id)
@@ -989,7 +1000,7 @@ extension BuildPlan {
         func successors(
             for product: ResolvedProduct,
             destination: Destination
-        ) -> [Node] {
+        ) -> [TraversalNode] {
             guard destination == .host else {
                 return []
             }
@@ -1002,12 +1013,12 @@ extension BuildPlan {
         func successors(
             for module: ResolvedModule,
             destination: Destination
-        ) -> [Node] {
+        ) -> [TraversalNode] {
             guard destination == .host else {
                 return []
             }
 
-            return module.dependencies.reduce(into: [Node]()) { partial, dependency in
+            return module.dependencies.reduce(into: [TraversalNode]()) { partial, dependency in
                 switch dependency {
                 case .product(let product, conditions: _):
                     partial.append(.for(product: product, context: destination))
@@ -1017,7 +1028,7 @@ extension BuildPlan {
             }
         }
 
-        try await depthFirstSearch(graph.packages.map { Node.package($0) }) { node in
+        try await depthFirstSearch(graph.packages.map { TraversalNode.package($0) }) { node in
             switch node {
             case .package(let package):
                 successors(for: package)
@@ -1040,6 +1051,71 @@ extension BuildPlan {
             // No de-duplication is necessary we only want unique nodes.
         }
     }
+
+    /// Traverses the modules graph, computes destination of every module reference and
+    /// provides the data to the caller by means of `onModule` callback. The products
+    /// are completely transparent to this method and are represented by their module dependencies.
+    package func traverseModules(
+        _ onModule: (
+            (ResolvedModule, BuildParameters.Destination),
+            _ parent: (ResolvedModule, BuildParameters.Destination)?,
+            _ depth: Int
+        ) -> Void
+    ) {
+        func successors(for package: ResolvedPackage) -> [TraversalNode] {
+            package.modules.compactMap {
+                if case .test = $0.underlying.type,
+                   !self.graph.rootPackages.contains(id: package.id)
+                {
+                    return nil
+                }
+                return .for(module: $0, context: .target)
+            }
+        }
+
+        func successors(
+            for module: ResolvedModule,
+            destination: Destination
+        ) -> [TraversalNode] {
+            module.dependencies.reduce(into: [TraversalNode]()) { partial, dependency in
+                switch dependency {
+                case .product(let product, conditions: _):
+                    let parent: TraversalNode = .for(product: product, context: destination)
+                    for module in product.modules {
+                        partial.append(.for(module: module, context: parent.destination))
+                    }
+                case .module(let module, _):
+                    partial.append(.for(module: module, context: destination))
+                }
+            }
+        }
+
+        depthFirstSearch(self.graph.packages.map { TraversalNode.package($0) }) {
+            switch $0 {
+            case .package(let package):
+                successors(for: package)
+            case .module(let module, let destination):
+                successors(for: module, destination: destination)
+            case .product:
+                []
+            }
+        } onNext: { current, parent, depth in
+            let parentModule: (ResolvedModule, BuildParameters.Destination)? = switch parent {
+            case .package, .product, nil:
+                nil
+            case .module(let module, let destination):
+                (module, destination)
+            }
+
+            switch current {
+            case .package, .product:
+                break
+
+            case .module(let module, let destination):
+                onModule((module, destination), parentModule, depth)
+            }
+        }
+    }
 }
 
 extension Basics.Diagnostic {
diff --git a/Sources/PackageGraph/ModulesGraph.swift b/Sources/PackageGraph/ModulesGraph.swift
@@ -107,14 +107,6 @@ public struct ModulesGraph {
     /// Returns all the modules in the graph, regardless if they are reachable from the root modules or not.
     public private(set) var allModules: IdentifiableSet<ResolvedModule>
 
-    /// Returns all modules within the graph in topological order, starting with low-level modules (that have no
-    /// dependencies).
-    package var allModulesInTopologicalOrder: [ResolvedModule] {
-        get throws {
-            try topologicalSort(Array(allModules)) { $0.dependencies.compactMap { $0.module } }.reversed()
-        }
-    }
-
     /// Returns all the products in the graph, regardless if they are reachable from the root modules or not.
     public private(set) var allProducts: IdentifiableSet<ResolvedProduct>
 
diff --git a/Sources/SourceKitLSPAPI/BuildDescription.swift b/Sources/SourceKitLSPAPI/BuildDescription.swift
@@ -142,11 +142,22 @@ public struct BuildDescription {
         }
     }
 
-    /// Returns all targets within the module graph in topological order, starting with low-level targets (that have no
-    /// dependencies).
-    public func allTargetsInTopologicalOrder(in modulesGraph: ModulesGraph) throws -> [BuildTarget] {
-        try modulesGraph.allModulesInTopologicalOrder.compactMap {
-            getBuildTarget(for: $0, in: modulesGraph)
+    public func traverseModules(
+        callback: (any BuildTarget, _ parent: (any BuildTarget)?, _ depth: Int) -> Void
+    ) {
+        // TODO: One the `targetMap` is switched over to use `IdentifiableSet<ModuleBuildDescription>`
+        // we can introduce `BuildPlan.description(ResolvedModule, BuildParameters.Destination)`
+        // and start using that here.
+        self.buildPlan.traverseModules { module, parent, depth in
+            let parentDescription: (any BuildTarget)? = if let parent {
+                getBuildTarget(for: parent.0, in: self.buildPlan.graph)
+            } else {
+                nil
+            }
+
+            if let description = getBuildTarget(for: module.0, in: self.buildPlan.graph) {
+                callback(description, parentDescription, depth)
+            }
         }
     }
 
diff --git a/Tests/BuildTests/BuildPlanTests.swift b/Tests/BuildTests/BuildPlanTests.swift
diff --git a/Tests/SourceKitLSPAPITests/SourceKitLSPAPITests.swift b/Tests/SourceKitLSPAPITests/SourceKitLSPAPITests.swift
