traverse: pass key correctly (#44460)

Fixes a bug where custom keys to identify nodes were not passed
correctly.

2 files changed, 25 insertions, 2 deletions

diff --git a/lib/spack/spack/test/traverse.py b/lib/spack/spack/test/traverse.py
index 0876c65cc2..79e05c0b2c 100644
--- a/lib/spack/spack/test/traverse.py
+++ b/lib/spack/spack/test/traverse.py
@@ -272,6 +272,29 @@ def test_breadth_first_versus_depth_first_tree(abstract_specs_chain):
     ]
 
 
+@pytest.mark.parametrize("cover", ["nodes", "edges"])
+@pytest.mark.parametrize("depth_first", [True, False])
+def test_tree_traversal_with_key(cover, depth_first, abstract_specs_chain):
+    """Compare two multisource traversals of the same DAG. In one case the DAG consists of unique
+    Spec instances, in the second case there are identical copies of nodes and edges. Traversal
+    should be equivalent when nodes are identified by dag_hash."""
+    a = abstract_specs_chain["chain-a"]
+    c = abstract_specs_chain["chain-c"]
+    kwargs = {"cover": cover, "depth_first": depth_first}
+    dag_hash = lambda s: s.dag_hash()
+
+    # Traverse DAG spanned by a unique set of Spec instances
+    first = traverse.traverse_tree([a, c], key=id, **kwargs)
+
+    # Traverse equivalent DAG with copies of Spec instances included, keyed by dag hash.
+    second = traverse.traverse_tree([a, c.copy()], key=dag_hash, **kwargs)
+
+    # Check that the same nodes are discovered at the same depth
+    node_at_depth_first = [(depth, dag_hash(edge.spec)) for (depth, edge) in first]
+    node_at_depth_second = [(depth, dag_hash(edge.spec)) for (depth, edge) in second]
+    assert node_at_depth_first == node_at_depth_second
+
+
 def test_breadth_first_versus_depth_first_printing(abstract_specs_chain):
     """Test breadth-first versus depth-first tree printing."""
     s = abstract_specs_chain["chain-a"]
diff --git a/lib/spack/spack/traverse.py b/lib/spack/spack/traverse.py
index 3ac2bfe24c..091ebfe193 100644
--- a/lib/spack/spack/traverse.py
+++ b/lib/spack/spack/traverse.py
@@ -563,10 +563,10 @@ def traverse_tree(
     # identical to DFS, which is much more efficient then.
     if not depth_first and cover == "edges":
         edges, parents = breadth_first_to_tree_edges(specs, deptype, key)
-        return traverse_breadth_first_tree_edges(None, edges, parents)
+        return traverse_breadth_first_tree_edges(None, edges, parents, key)
     elif not depth_first and cover == "nodes":
         edges = breadth_first_to_tree_nodes(specs, deptype, key)
-        return traverse_breadth_first_tree_nodes(None, edges)
+        return traverse_breadth_first_tree_nodes(None, edges, key)
 
     return traverse_edges(specs, order="pre", cover=cover, deptype=deptype, key=key, depth=True)