third_party/python/taskcluster_taskgraph/taskgraph/graph.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4
   5
   6 import collections
   7 from dataclasses import dataclass
   8 from typing import FrozenSet
   9
  10
  11 @dataclass(frozen=True)
  12 class Graph:
  13     """Generic representation of a directed acyclic graph with labeled edges
  14     connecting the nodes. Graph operations are implemented in a functional
  15     manner, so the data structure is immutable.
  16
  17     It permits at most one edge of a given name between any set of nodes.  The
  18     graph is not checked for cycles, and methods may hang or otherwise fail if
  19     given a cyclic graph.
  20
  21     The `nodes` and `edges` attributes may be accessed in a read-only fashion.
  22     The `nodes` attribute is a set of node names, while `edges` is a set of
  23     `(left, right, name)` tuples representing an edge named `name` going from
  24     node `left` to node `right`..
  25     """
  26
  27     nodes: FrozenSet
  28     edges: FrozenSet
  29
  30     def transitive_closure(self, nodes, reverse=False):
  31         """Return the transitive closure of <nodes>: the graph containing all
  32         specified nodes as well as any nodes reachable from them, and any
  33         intervening edges.
  34
  35         If `reverse` is true, the "reachability" will be reversed and this
  36         will return the set of nodes that can reach the specified nodes.
  37
  38         Example:
  39
  40         .. code-block::
  41
  42             a ------> b ------> c
  43                       |
  44                       `-------> d
  45
  46         transitive_closure([b]).nodes == set([a, b])
  47         transitive_closure([c]).nodes == set([c, b, a])
  48         transitive_closure([c], reverse=True).nodes == set([c])
  49         transitive_closure([b], reverse=True).nodes == set([b, c, d])
  50         """
  51         assert isinstance(nodes, set)
  52         if not (nodes <= self.nodes):
  53             raise Exception(
  54                 f"Unknown nodes in transitive closure: {nodes - self.nodes}"
  55             )
  56
  57         # generate a new graph by expanding along edges until reaching a fixed
  58         # point
  59         new_nodes, new_edges = nodes, set()
  60         nodes, edges = set(), set()
  61         while (new_nodes, new_edges) != (nodes, edges):
  62             nodes, edges = new_nodes, new_edges
  63             add_edges = {
  64                 (left, right, name)
  65                 for (left, right, name) in self.edges
  66                 if (right if reverse else left) in nodes
  67             }
  68             add_nodes = {(left if reverse else right) for (left, right, _) in add_edges}
  69             new_nodes = nodes | add_nodes
  70             new_edges = edges | add_edges
  71         return Graph(new_nodes, new_edges)
  72
  73     def _visit(self, reverse):
  74         queue = collections.deque(sorted(self.nodes))
  75         links_by_node = self.reverse_links_dict() if reverse else self.links_dict()
  76         seen = set()
  77         while queue:
  78             node = queue.popleft()
  79             if node in seen:
  80                 continue
  81             links = links_by_node[node]
  82             if all((n in seen) for n in links):
  83                 seen.add(node)
  84                 yield node
  85             else:
  86                 queue.extend(n for n in links if n not in seen)
  87                 queue.append(node)
  88
  89     def visit_postorder(self):
  90         """
  91         Generate a sequence of nodes in postorder, such that every node is
  92         visited *after* any nodes it links to.
  93
  94         Behavior is undefined (read: it will hang) if the graph contains a
  95         cycle.
  96         """
  97         return self._visit(False)
  98
  99     def visit_preorder(self):
 100         """
 101         Like visit_postorder, but in reverse: evrey node is visited *before*
 102         any nodes it links to.
 103         """
 104         return self._visit(True)
 105
 106     def links_dict(self):
 107         """
 108         Return a dictionary mapping each node to a set of the nodes it links to
 109         (omitting edge names)
 110         """
 111         links = collections.defaultdict(set)
 112         for left, right, _ in self.edges:
 113             links[left].add(right)
 114         return links
 115
 116     def named_links_dict(self):
 117         """
 118         Return a two-level dictionary mapping each node to a dictionary mapping
 119         edge names to labels.
 120         """
 121         links = collections.defaultdict(dict)
 122         for left, right, name in self.edges:
 123             links[left][name] = right
 124         return links
 125
 126     def reverse_links_dict(self):
 127         """
 128         Return a dictionary mapping each node to a set of the nodes linking to
 129         it (omitting edge names)
 130         """
 131         links = collections.defaultdict(set)
 132         for left, right, _ in self.edges:
 133             links[right].add(left)
 134         return links