graph.py

   1 import sys
   2 import copy
   3 from collections import defaultdict
   4
   5
   6 class Graph:
   7     """Nodes in a graph are identified by IDs and contain property dictionaries.
   8     A graph maps node ID to node properties:
   9
  10     node_props = G[node_id]
  11
  12     A specific node property thus can be accessed as:
  13
  14     prop_val = G[node_id][prop_name]
  15
  16     The "default" property, a "value" of node, is by convention in a property named
  17     "val".
  18
  19     Edges in a graph are identified by a pair of node IDs. They also have associated
  20     property dictionary:
  21
  22     edge_props = G[(node_id1, node_id2)]
  23
  24     """
  25
  26     directed = True
  27
  28     def __init__(self):
  29         self._nodes = {}
  30         self._edges = {}
  31         self._succ = defaultdict(list)
  32         self._pred = defaultdict(list)
  33         self.first_node = None
  34         # Graph-level properties
  35         self.props = {}
  36
  37     def add_node(self, node, **attrs):
  38         """Add node to a graph. node is an ID of a node (usually lightweight
  39         scalar value, but can be any immutable value). Arbitrary attributes
  40         can be associated with a node, e.g. "val" attribute for node's "value".
  41         """
  42         if node in self._nodes:
  43             self._nodes[node].update(attrs)
  44         else:
  45             self._nodes[node] = attrs
  46         # Many algos need proper-entry graph, and if graph is not
  47         # (e.g., has a loop backedge to entry), it's not possible
  48         # to detect such entry without explicit pointing or
  49         # heuristics. We use the latter here - first node added is
  50         # an entry.
  51         if self.first_node is None:
  52             self.first_node = node
  53
  54     def remove_node(self, node):
  55         for s in self._succ[node][:]:
  56             self.remove_edge(node, s)
  57         for p in self._pred[node][:]:
  58             self.remove_edge(p, node)
  59         del self._nodes[node]
  60         del self._succ[node]
  61         del self._pred[node]
  62
  63     def node(self, n):
  64         return self._nodes[n]
  65
  66     # Allow to index graph to access node or edge data
  67     # graph[node], graph[from_n, to_n]
  68     def __getitem__(self, idx):
  69         if isinstance(idx, tuple):
  70             return self.edge(*idx)
  71         else:
  72             return self.node(idx)
  73
  74     def set_node_attr(self, node, attr, val):
  75         self._nodes[node][attr] = val
  76
  77     def get_node_attr(self, node, attr, default=Ellipsis):
  78         if default is Ellipsis:
  79             return self._nodes[node][attr]
  80         else:
  81             return self._nodes[node].get(attr, default)
  82
  83
  84     def add_edge(self, from_node, to_node, **attrs):
  85         """Add edge between 2 nodes. If any of the nodes does not exist,
  86         it will be created."""
  87         if (from_node, to_node) in self._edges:
  88             self._edges[(from_node, to_node)].update(attrs)
  89         else:
  90             self.add_node(from_node)
  91             self.add_node(to_node)
  92             self._edges[(from_node, to_node)] = attrs
  93             self._succ[from_node].append(to_node)
  94             self._pred[to_node].append(from_node)
  95
  96     set_edge = add_edge
  97
  98     def remove_edge(self, from_node, to_node):
  99         del self._edges[(from_node, to_node)]
 100         self._succ[from_node].remove(to_node)
 101         self._pred[to_node].remove(from_node)
 102
 103     def edge(self, from_node, to_node):
 104         return self._edges[(from_node, to_node)]
 105
 106     def has_edge(self, from_node, to_node):
 107         return (from_node, to_node) in self._edges
 108
 109     def is_back_edge(self, from_node, to_node):
 110         raise NotImplementedError
 111         # This is not correct a test. And edge like that can be a cross edge
 112         # too. To properly distinguish back edge from cross, combination of
 113         # pre-order number and post-order number is required.
 114         #return self[from_node]["postno"] < self[to_node]["postno"]
 115
 116     def is_empty(self):
 117         return not self._nodes
 118
 119     def succ(self, n):
 120         "Return successors of a node."
 121         return self._succ[n][:]
 122
 123     def sorted_succ(self, n):
 124         """Return successors ordered the way that successor with labeled
 125         edge comes first. Assumes 2 succesors."""
 126         succ = self.succ(n)
 127         if len(succ) < 2:
 128             return succ
 129         assert len(succ) == 2
 130         if self.edge(n, succ[0]).get("cond") is None:
 131             succ = [succ[1], succ[0]]
 132         return succ
 133
 134     def pred(self, n):
 135         "Return predecessors of a node."
 136         return self._pred[n][:]
 137
 138     def degree_out(self, n):
 139         return len(self._succ[n])
 140
 141     def degree_in(self, n):
 142         return len(self._pred[n])
 143
 144     def __contains__(self, val):
 145         if isinstance(val, tuple):
 146             return val in self._edges
 147         else:
 148             return val in self._nodes
 149
 150     def nodes(self):
 151         "Iterate over nodes."
 152         return self._nodes.keys()
 153
 154     def nodes_props(self):
 155         "Iterate over pairs of (node, node_properties)."
 156         return self._nodes.items()
 157
 158     def iter_sorted_nodes(self):
 159         return sorted(self._nodes.items(), key=lambda x: x[0])
 160
 161     def edges(self):
 162         "Iterate over edges."
 163         return self._edges.keys()
 164
 165     def edges_props(self):
 166         "Iterate over pairs of (edge, edge_properties)."
 167         return self._edges.items()
 168
 169     def entries(self):
 170         # Will also return single disconnected nodes (which are entries
 171         # by all means).
 172         entries = [n for n in self._nodes if not self._pred[n]]
 173         return entries
 174
 175     def entry(self):
 176         e = self.entries()
 177         assert len(e) == 1, "Expected single entry, instead multiple: %r" % e
 178         return e[0]
 179
 180     def exits(self):
 181         # TODO: Will also return single disconnected nodes
 182         return [n for n in self._nodes if not self._succ[n]]
 183
 184     def exit(self):
 185         e = self.exits()
 186         assert len(e) == 1, "Expected single exit, instead multiple: %r" % e
 187         return e[0]
 188
 189     def move_pred(self, from_node, to_node):
 190         """Move all predecessor edges from one node to another. Useful in
 191         graph transformations involving node folding."""
 192         for p in self.pred(from_node):
 193             self.add_edge(p, to_node, **self._edges[(p, from_node)])
 194             self.remove_edge(p, from_node)
 195
 196     def move_succ(self, from_node, to_node):
 197         """Move all succesor edges from one node to another. Useful in
 198         graph transformations involving node folding."""
 199         for p in self.succ(from_node):
 200             self.add_edge(to_node, p, **self._edges[(from_node, p)])
 201             self.remove_edge(from_node, p)
 202
 203     def __repr__(self):
 204         return "<Graph nodes=%r edges=%r pred=%r succ=%r>" % (self._nodes, self._edges, self._pred, self._succ)
 205
 206     def print_nodes(self, stream=sys.stdout):
 207         "Print nodes of a graph with all attributes. Useful for debugging."
 208         for i, info in self.iter_sorted_nodes():
 209             print("%s\t%s" % (i, info))
 210
 211     def reset_numbering(self, prop="postno"):
 212         for n, info in self._nodes.items():
 213             info[prop] = None
 214
 215     def _number_postorder_bidi(self, node, num, prop_name, next_func):
 216         self.set_node_attr(node, prop_name, True)
 217         next_nodes = next_func(node)
 218         # TODO: If not using reverse, numbering changes to less natural, but
 219         # algos which depend on postorder should not?
 220         next_nodes.sort(reverse=True)
 221         for n in next_nodes:
 222             if not self.get_node_attr(n, prop_name, None):
 223                 num = self._number_postorder_bidi(n, num, prop_name, next_func)
 224         self.set_node_attr(node, prop_name, num)
 225         if prop_name == "postno":
 226             self._postorder[num - 1] = node
 227         #print("Setting %s to %s" % (node, num))
 228         num += 1
 229         return num
 230
 231     def number_postorder(self, node=None, num=1):
 232         "Number nodes in depth-first search post-order order."
 233         self.reset_numbering()
 234         self._postorder = [None] * len(self._nodes)
 235         if node is None:
 236             node = self.first_node
 237         return self._number_postorder_bidi(node, 1, "postno", self.succ)
 238
 239     def number_postorder_from_exit(self, node):
 240         "Number nodes in depth-first search post-order, from exit."
 241         self.reset_numbering("postno_exit")
 242         return self._number_postorder_bidi(node, 1, "postno_exit", self.pred)
 243
 244     def number_postorder_forest(self):
 245         "Number nodes in depth-first search post-order order."
 246         self.reset_numbering()
 247         entries = self.entries()
 248         num = 1
 249         for e in entries:
 250             num = self.number_postorder(e, num)
 251
 252     def iter_postorder(self):
 253         return iter(self._postorder)
 254         #return sorted(self._nodes.items(), key=lambda x: x[1]["postno"])
 255
 256     def iter_rev_postorder(self):
 257         return reversed(self._postorder)
 258         #return sorted(self._nodes.items(), key=lambda x: -x[1]["postno"])
 259
 260     def copy(self):
 261         # TODO: not optimal
 262         return copy.deepcopy(self)
 263
 264     def __eq__(self, other):
 265         if self._nodes != other._nodes:
 266             return False
 267         if self._edges != other._edges:
 268             return False
 269         return True
 270
 271
 272 def find_all_nodes_on_path(cfg, from_n, to_n):
 273     """Find set of nodes which lie on all paths from from_n to to_n.
 274     from_n and to_b can be a same node to process a loop. Result includes
 275     from_n and excludes to_n, if they are different nodes (but node
 276     looping to itself gives result of itself).
 277     """
 278     on_path = set()
 279
 280     stack = [(from_n, [from_n])]
 281     while stack:
 282         cur_n, path = stack.pop()
 283         succ = cfg.succ(cur_n)
 284         postno = cfg.get_node_attr(cur_n, "postno")
 285         assert postno is not None
 286         print("Considering %s, postno: %s, succ: %s, cur path: %s" % (cur_n, postno, succ, path))
 287
 288         if not succ:
 289             print("Found sink path, pruning")
 290
 291         for s in succ:
 292             if s == to_n:
 293                 on_path.update(path)
 294                 print("Found full path, joining all path nodes")
 295             else:
 296                 no = cfg.get_node_attr(s, "postno", None)
 297                 if no is None:
 298                     print("Unmarked node detected (edge %s->%s)" % (cur_n, s))
 299                     assert False
 300                 if no < postno:
 301                     stack.append((s, path + [s]))
 302                 else:
 303                     # prune path
 304                     print("Found far backedge (%s, %s), pruning this path" % (cur_n, s))
 305     return on_path