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