From 60ee1817b0b91fcfc3f610dc5ae81fbf19ea7bb8 Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Wed, 29 Jul 2026 15:08:38 -0700
Subject: [PATCH] perf: avoid redundant work in Graph.transitive_closure

The adjacency map is built over every edge before the traversal starts,
which an empty node set never needs. Return early instead: on a
20k-node, 2M-edge graph an empty closure drops from ~1.3s to
microseconds.

Resolve `reverse` to a pair of tuple indices once rather than branching
on it per edge in both loops, worth 1.1-1.4x on sparse graphs. Accept a
frozenset in the input assert, since Graph.nodes is itself a frozenset
and g.transitive_closure(g.nodes) previously tripped it.

The four worked examples in the docstring had forward and reverse
swapped. They have been wrong since 2eec8615 imported the module in
2019, were never covered by a test, and contradict the prose above them,
which was correct.

Add coverage for the diamond, frozenset and unknown-node paths.
---
 src/taskgraph/graph.py | 29 ++++++++++++++++++-----------
 test/test_graph.py     | 27 +++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/src/taskgraph/graph.py b/src/taskgraph/graph.py
index 576325514..a59875295 100644
--- a/src/taskgraph/graph.py
+++ b/src/taskgraph/graph.py
@@ -50,32 +50,39 @@ def transitive_closure(self, nodes, reverse=False):
                       |
                       `-------> d
 
-        transitive_closure([b]).nodes == set([a, b])
-        transitive_closure([c]).nodes == set([c, b, a])
-        transitive_closure([c], reverse=True).nodes == set([c])
-        transitive_closure([b], reverse=True).nodes == set([b, c, d])
+        transitive_closure([b]).nodes == set([b, c, d])
+        transitive_closure([c]).nodes == set([c])
+        transitive_closure([c], reverse=True).nodes == set([c, b, a])
+        transitive_closure([b], reverse=True).nodes == set([b, a])
         """
-        assert isinstance(nodes, set)
+        assert isinstance(nodes, (set, frozenset))
         if not (nodes <= self.nodes):
             raise Exception(
                 f"Unknown nodes in transitive closure: {nodes - self.nodes}"
             )
 
-        # Build an adjacency map keyed by the node to expand from. This reduces
-        # traversal below from O(V·E) -> O(V+E).
-        adjacency = collections.defaultdict(list)
+        # An empty closure reaches nothing, and the adjacency map below costs
+        # O(E) to build.
+        if not nodes:
+            return Graph(set(), set())
+
+        # Index the edges by the endpoint the traversal expands from, keeping the
+        # traversal O(V+E). `src` is the endpoint an edge is looked up by, `dst`
+        # the one it expands to.
+        src, dst = (1, 0) if reverse else (0, 1)
+        adjacency = {}
         for edge in self.edges:
-            left, right, _ = edge
-            adjacency[right if reverse else left].append(edge)
+            adjacency.setdefault(edge[src], []).append(edge)
 
         new_nodes = set(nodes)
         new_edges = set()
         queue = collections.deque(nodes)
         while queue:
             node = queue.popleft()
+            # Nodes with no edges to expand from have no entry in the map.
             for edge in adjacency.get(node, ()):
                 new_edges.add(edge)
-                neighbor = edge[0] if reverse else edge[1]
+                neighbor = edge[dst]
                 if neighbor not in new_nodes:
                     new_nodes.add(neighbor)
                     queue.append(neighbor)
diff --git a/test/test_graph.py b/test/test_graph.py
index 81d339baa..9820b0a5d 100644
--- a/test/test_graph.py
+++ b/test/test_graph.py
@@ -194,6 +194,33 @@ def test_transitive_closure_loopy_reverse(self):
         "reverse transitive closure of a loop is the whole loop"
         self.assertEqual(self.loopy.transitive_closure({"A"}, reverse=True), self.loopy)
 
+    def test_transitive_closure_diamond(self):
+        "transitive closure of a diamond reaches shared descendants by both paths"
+        self.assertEqual(
+            self.diamonds.transitive_closure({"A"}),
+            Graph(
+                {"A", "D", "F", "G", "I", "J"},
+                {
+                    ("A", "D", "L"),
+                    ("A", "F", "L"),
+                    ("D", "F", "L"),
+                    ("D", "G", "L"),
+                    ("F", "I", "L"),
+                    ("G", "I", "L"),
+                    ("G", "J", "L"),
+                },
+            ),
+        )
+
+    def test_transitive_closure_frozenset(self):
+        "transitive closure accepts a frozenset, as returned by Graph.nodes"
+        self.assertEqual(self.tree.transitive_closure(self.tree.nodes), self.tree)
+
+    def test_transitive_closure_unknown_nodes(self):
+        "transitive closure raises when given nodes not in the graph"
+        with pytest.raises(Exception, match="Unknown nodes"):
+            self.tree.transitive_closure({"z"})
+
     def test_visit_postorder_empty(self):
         "postorder visit of an empty graph is empty"
         self.assertEqual(list(Graph(set(), set()).visit_postorder()), [])