Merge pull request #38 from serpilliere/graph_add_dominators

Graph: add dominators computation

2 files changed, 75 insertions, 0 deletions

diff --git a/miasm2/core/graph.py b/miasm2/core/graph.py
index d29517cc..2b873ad9 100644
--- a/miasm2/core/graph.py
+++ b/miasm2/core/graph.py
@@ -1,4 +1,7 @@
+from collections import defaultdict
+
 class DiGraph:
+    """Implementation of directed graph"""
 
     def __init__(self):
         self._nodes = set()
@@ -84,6 +87,14 @@ class DiGraph:
     def leaves(self):
         return [x for x in self.leaves_iter()]
 
+    def heads_iter(self):
+        for node in self._nodes:
+            if len(self._nodes_from[node]) == 0:
+                yield node
+
+    def heads(self):
+        return [node for node in self.heads_iter()]
+
     def roots_iter(self):
         for n in self._nodes:
             if len(self._nodes_from[n]) == 0:
@@ -149,3 +160,44 @@ shape = "box"
                                                 self.edge2str(a, b))
         out += "}"
         return out
+
+
+
+    def compute_dominators(self):
+        """Compute dominators of the graph"""
+        nodes = self.nodes()
+        heads = self.heads()
+
+        dominators = defaultdict(set)
+        for node in self.nodes():
+            dominators[node].update(nodes)
+        for node in heads:
+            dominators[node] = set([node])
+
+        modified = True
+        todo = nodes.difference(heads)
+
+        while todo:
+            node = todo.pop()
+
+            # Heads state must not be changed
+            if node in heads:
+                continue
+
+            # Compute intersection of all predecessors'dominators
+            new_dom = None
+            for pred in self.predecessors(node):
+                if new_dom is None:
+                    new_dom = set(dominators[pred])
+                new_dom.intersection_update(dominators[pred])
+            if new_dom is None:
+                new_dom = set()
+            new_dom.update(set([node]))
+
+            # If intersection has changed, add sons to the todo list
+            if new_dom == dominators[node]:
+                continue
+            dominators[node] = new_dom
+            for succ in self.successors(node):
+                todo.add(succ)
+        return dict(dominators)
diff --git a/test/core/graph.py b/test/core/graph.py
index a419a686..eb896cfe 100644
--- a/test/core/graph.py
+++ b/test/core/graph.py
@@ -16,3 +16,26 @@ print [x for x in g.predecessors('a')]
 print [x for x in g.predecessors('b')]
 print [x for x in g.predecessors('c')]
 print [x for x in g.successors('c')]
+
+
+"""
+Test from: https://en.wikipedia.org/wiki/Dominator_(graph_theory)
+"""
+
+g1 = DiGraph()
+g1.add_edge(1, 2)
+g1.add_edge(2, 3)
+g1.add_edge(2, 4)
+g1.add_edge(3, 5)
+g1.add_edge(4, 5)
+g1.add_edge(5, 2)
+g1.add_edge(2, 6)
+
+
+dominators = g1.compute_dominators()
+assert(dominators[1] == set([1]))
+assert(dominators[2] == set([1, 2]))
+assert(dominators[3] == set([1, 2, 3]))
+assert(dominators[4] == set([1, 2, 4]))
+assert(dominators[5] == set([1, 2, 5]))
+assert(dominators[6] == set([1, 2, 6]))