traverse: add MixedDepthVisitor, use in cmake (#47750)

This visitor accepts the sub-dag of all nodes and unique edges that have
deptype X directly from given roots, or deptype Y transitively for any
of the roots.

3 files changed, 136 insertions, 19 deletions

diff --git a/lib/spack/spack/build_systems/cmake.py b/lib/spack/spack/build_systems/cmake.py
index dd1261c811..010d427633 100644
--- a/lib/spack/spack/build_systems/cmake.py
+++ b/lib/spack/spack/build_systems/cmake.py
@@ -9,7 +9,7 @@ import platform
 import re
 import sys
 from itertools import chain
-from typing import Any, List, Optional, Set, Tuple
+from typing import Any, List, Optional, Tuple
 
 import llnl.util.filesystem as fs
 from llnl.util.lang import stable_partition
@@ -21,6 +21,7 @@ import spack.package_base
 import spack.phase_callbacks
 import spack.spec
 import spack.util.prefix
+from spack import traverse
 from spack.directives import build_system, conflicts, depends_on, variant
 from spack.multimethod import when
 from spack.util.environment import filter_system_paths
@@ -166,15 +167,18 @@ def generator(*names: str, default: Optional[str] = None) -> None:
 def get_cmake_prefix_path(pkg: spack.package_base.PackageBase) -> List[str]:
     """Obtain the CMAKE_PREFIX_PATH entries for a package, based on the cmake_prefix_path package
     attribute of direct build/test and transitive link dependencies."""
-    # Add direct build/test deps
-    selected: Set[str] = {s.dag_hash() for s in pkg.spec.dependencies(deptype=dt.BUILD | dt.TEST)}
-    # Add transitive link deps
-    selected.update(s.dag_hash() for s in pkg.spec.traverse(root=False, deptype=dt.LINK))
-    # Separate out externals so they do not shadow Spack prefixes
-    externals, spack_built = stable_partition(
-        (s for s in pkg.spec.traverse(root=False, order="topo") if s.dag_hash() in selected),
-        lambda x: x.external,
+    edges = traverse.traverse_topo_edges_generator(
+        traverse.with_artificial_edges([pkg.spec]),
+        visitor=traverse.MixedDepthVisitor(
+            direct=dt.BUILD | dt.TEST, transitive=dt.LINK, key=traverse.by_dag_hash
+        ),
+        key=traverse.by_dag_hash,
+        root=False,
+        all_edges=False,  # cover all nodes, not all edges
     )
+    ordered_specs = [edge.spec for edge in edges]
+    # Separate out externals so they do not shadow Spack prefixes
+    externals, spack_built = stable_partition((s for s in ordered_specs), lambda x: x.external)
 
     return filter_system_paths(
         path for spec in chain(spack_built, externals) for path in spec.package.cmake_prefix_paths
diff --git a/lib/spack/spack/test/traverse.py b/lib/spack/spack/test/traverse.py
index 62ce24d366..79aae6eb55 100644
--- a/lib/spack/spack/test/traverse.py
+++ b/lib/spack/spack/test/traverse.py
@@ -20,9 +20,8 @@ def create_dag(nodes, edges):
     """
     specs = {name: Spec(name) for name in nodes}
     for parent, child, deptypes in edges:
-        specs[parent].add_dependency_edge(
-            specs[child], depflag=dt.canonicalize(deptypes), virtuals=()
-        )
+        depflag = deptypes if isinstance(deptypes, dt.DepFlag) else dt.canonicalize(deptypes)
+        specs[parent].add_dependency_edge(specs[child], depflag=depflag, virtuals=())
     return specs
 
 
@@ -454,3 +453,61 @@ def test_topo_is_bfs_for_trees(cover):
     assert list(traverse.traverse_nodes([binary_tree["A"]], order="topo", cover=cover)) == list(
         traverse.traverse_nodes([binary_tree["A"]], order="breadth", cover=cover)
     )
+
+
+@pytest.mark.parametrize("roots", [["A"], ["A", "B"], ["B", "A"], ["A", "B", "A"]])
+@pytest.mark.parametrize("order", ["breadth", "post", "pre"])
+@pytest.mark.parametrize("include_root", [True, False])
+def test_mixed_depth_visitor(roots, order, include_root):
+    """Test that the MixedDepthVisitor lists unique edges that are reachable either directly from
+    roots through build type edges, or transitively through link type edges. The tests ensures that
+    unique edges are listed exactly once."""
+    my_graph = create_dag(
+        nodes=["A", "B", "C", "D", "E", "F", "G", "H", "I"],
+        edges=(
+            ("A", "B", dt.LINK | dt.RUN),
+            ("A", "C", dt.BUILD),
+            ("A", "D", dt.BUILD | dt.RUN),
+            ("A", "H", dt.LINK),
+            ("A", "I", dt.RUN),
+            ("B", "D", dt.BUILD | dt.LINK),
+            ("C", "E", dt.BUILD | dt.LINK | dt.RUN),
+            ("D", "F", dt.LINK),
+            ("D", "G", dt.BUILD | dt.RUN),
+            ("H", "B", dt.LINK),
+        ),
+    )
+    starting_points = traverse.with_artificial_edges([my_graph[root] for root in roots])
+    visitor = traverse.MixedDepthVisitor(direct=dt.BUILD, transitive=dt.LINK)
+
+    if order == "pre":
+        edges = traverse.traverse_depth_first_edges_generator(
+            starting_points, visitor, post_order=False, root=include_root
+        )
+    elif order == "post":
+        edges = traverse.traverse_depth_first_edges_generator(
+            starting_points, visitor, post_order=True, root=include_root
+        )
+    elif order == "breadth":
+        edges = traverse.traverse_breadth_first_edges_generator(
+            starting_points, visitor, root=include_root
+        )
+
+    artificial_edges = [(None, root) for root in roots] if include_root else []
+    simple_edges = [
+        (None if edge.parent is None else edge.parent.name, edge.spec.name) for edge in edges
+    ]
+
+    # make sure that every edge is listed exactly once and that the right edges are listed
+    assert len(simple_edges) == len(set(simple_edges))
+    assert set(simple_edges) == {
+        # the roots
+        *artificial_edges,
+        ("A", "B"),
+        ("A", "C"),
+        ("A", "D"),
+        ("A", "H"),
+        ("B", "D"),
+        ("D", "F"),
+        ("H", "B"),
+    }
diff --git a/lib/spack/spack/traverse.py b/lib/spack/spack/traverse.py
index 158f4d892e..880d7f71fa 100644
--- a/lib/spack/spack/traverse.py
+++ b/lib/spack/spack/traverse.py
@@ -4,7 +4,7 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
 from collections import defaultdict
-from typing import NamedTuple, Union
+from typing import Any, Callable, List, NamedTuple, Set, Union
 
 import spack.deptypes as dt
 import spack.spec
@@ -115,6 +115,64 @@ class CoverEdgesVisitor:
         return self.visitor.neighbors(item)
 
 
+class MixedDepthVisitor:
+    """Visits all unique edges of the sub-DAG induced by direct dependencies of type ``direct``
+    and transitive dependencies of type ``transitive``. An example use for this is traversing build
+    type dependencies non-recursively, and link dependencies recursively."""
+
+    def __init__(
+        self,
+        *,
+        direct: dt.DepFlag,
+        transitive: dt.DepFlag,
+        key: Callable[["spack.spec.Spec"], Any] = id,
+    ) -> None:
+        self.direct_type = direct
+        self.transitive_type = transitive
+        self.key = key
+        self.seen: Set[Any] = set()
+        self.seen_roots: Set[Any] = set()
+
+    def accept(self, item: EdgeAndDepth) -> bool:
+        # Do not accept duplicate root nodes. This only happens if the user starts iterating from
+        # multiple roots and lists one of the roots multiple times.
+        if item.edge.parent is None:
+            node_id = self.key(item.edge.spec)
+            if node_id in self.seen_roots:
+                return False
+            self.seen_roots.add(node_id)
+        return True
+
+    def neighbors(self, item: EdgeAndDepth) -> List[EdgeAndDepth]:
+        # If we're here through an artificial source node, it's a root, and we return all
+        # direct_type  and transitive_type edges. If we're here through a transitive_type edge, we
+        # return all transitive_type edges. To avoid returning the same edge twice:
+        # 1. If we had already encountered the current node through a transitive_type edge, we
+        #    don't need to return transitive_type edges again.
+        # 2. If we encounter the current node through a direct_type edge, and we had already seen
+        #    it through a transitive_type edge, only return the non-transitive_type, direct_type
+        #    edges.
+        node_id = self.key(item.edge.spec)
+        seen = node_id in self.seen
+        is_root = item.edge.parent is None
+        follow_transitive = is_root or bool(item.edge.depflag & self.transitive_type)
+        follow = self.direct_type if is_root else dt.NONE
+
+        if follow_transitive and not seen:
+            follow |= self.transitive_type
+            self.seen.add(node_id)
+        elif follow == dt.NONE:
+            return []
+
+        edges = item.edge.spec.edges_to_dependencies(depflag=follow)
+
+        # filter direct_type edges already followed before becuase they were also transitive_type.
+        if seen:
+            edges = [edge for edge in edges if not edge.depflag & self.transitive_type]
+
+        return sort_edges(edges)
+
+
 def get_visitor_from_args(
     cover, direction, depflag: Union[dt.DepFlag, dt.DepTypes], key=id, visited=None, visitor=None
 ):
@@ -342,9 +400,7 @@ def traverse_topo_edges_generator(edges, visitor, key=id, root=True, all_edges=F
     # maps parent identifier to a list of edges, where None is a special identifier
     # for the artificial root/source.
     node_to_edges = defaultdict(list)
-    for edge in traverse_breadth_first_edges_generator(
-        edges, CoverEdgesVisitor(visitor, key=key), root=True, depth=False
-    ):
+    for edge in traverse_breadth_first_edges_generator(edges, visitor, root=True, depth=False):
         in_edge_count[key(edge.spec)] += 1
         parent_id = key(edge.parent) if edge.parent is not None else None
         node_to_edges[parent_id].append(edge)
@@ -422,9 +478,9 @@ def traverse_edges(
     elif order not in ("post", "pre", "breadth"):
         raise ValueError(f"Unknown order {order}")
 
-    # In topo traversal we need to construct a sub-DAG including all edges even if we are yielding
-    # a subset of them, hence "paths".
-    _cover = "paths" if order == "topo" else cover
+    # In topo traversal we need to construct a sub-DAG including all unique edges even if we are
+    # yielding a subset of them, hence "edges".
+    _cover = "edges" if order == "topo" else cover
     visitor = get_visitor_from_args(_cover, direction, deptype, key, visited)
     root_edges = with_artificial_edges(specs)