Merge pull request #106 from serpilliere/dominator

Dominator

3 files changed, 109 insertions, 40 deletions

diff --git a/miasm2/core/graph.py b/miasm2/core/graph.py
index b25907e5..c59634d1 100644
--- a/miasm2/core/graph.py
+++ b/miasm2/core/graph.py
@@ -119,20 +119,6 @@ class DiGraph(object):
                     out.append(path + [b])
         return out
 
-    def get_all_parents(self, node):
-        """Return every parents nodes for a given @node"""
-
-        todo = set([node])
-        done = set()
-        while todo:
-            node = todo.pop()
-            if node in done:
-                continue
-            done.add(node)
-            for parent in self.predecessors(node):
-                todo.add(parent)
-        return done
-
     def node2str(self, n):
         return str(n)
 
@@ -162,42 +148,87 @@ shape = "box"
         return out
 
 
+    def _reachable_nodes(self, head, next_cb):
+        """Generic algorithm to compute every nodes reachable from/to node
+        @head"""
 
-    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])
-
+        todo = set([head])
+        reachable = set()
+        while todo:
+            node = todo.pop()
+            if node in reachable:
+                continue
+            reachable.add(node)
+            yield node
+            for next_node in next_cb(node):
+                todo.add(next_node)
+
+    def reachable_sons(self, head):
+        """Compute every nodes reachable from node @head"""
+        return self._reachable_nodes(head, self.successors_iter)
+
+    def reachable_parents(self, leaf):
+        """Compute every parents of node @leaf"""
+        return self._reachable_nodes(leaf, self.predecessors_iter)
+
+    def _compute_generic_dominators(self, head, reachable_cb, prev_cb, next_cb):
+        """Generic algorithm to compute either the dominators or postdominators
+        of the graph.
+        @head: the head/leaf of the graph
+        @reachable_cb: sons/parents of the head/leaf
+        @prev_cb: return predecessors/succesors of a node
+        @next_cb: return succesors/predecessors of a node
+        """
+
+        nodes = set(reachable_cb(head))
+        dominators = {}
+        for node in nodes:
+            dominators[node] = set(nodes)
+
+        dominators[head] = set([head])
         modified = True
-        todo = nodes.difference(heads)
+        todo = set(nodes)
 
         while todo:
             node = todo.pop()
 
             # Heads state must not be changed
-            if node in heads:
+            if node == head:
                 continue
 
             # Compute intersection of all predecessors'dominators
             new_dom = None
-            for pred in self.predecessors(node):
+            for pred in prev_cb(node):
+                if not pred in nodes:
+                    continue
                 if new_dom is None:
                     new_dom = set(dominators[pred])
                 new_dom.intersection_update(dominators[pred])
-            if new_dom is None:
-                new_dom = set()
+
+            # We are not a head to we have at least one dominator
+            assert(new_dom is not None)
+
             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):
+            for succ in next_cb(node):
                 todo.add(succ)
-        return dict(dominators)
+        return dominators
+
+    def compute_dominators(self, head):
+        """Compute the dominators of the graph"""
+        return self._compute_generic_dominators(head,
+                                                self.reachable_sons,
+                                                self.predecessors_iter,
+                                                self.successors_iter)
+
+    def compute_postdominators(self, leaf):
+        """Compute the postdominators of the graph"""
+        return self._compute_generic_dominators(leaf,
+                                                self.reachable_parents,
+                                                self.successors_iter,
+                                                self.predecessors_iter)
diff --git a/miasm2/ir/analysis.py b/miasm2/ir/analysis.py
index 202f7631..c6fc4948 100644
--- a/miasm2/ir/analysis.py
+++ b/miasm2/ir/analysis.py
@@ -253,7 +253,7 @@ class ira:
                     break
             if has_all_son:
                 continue
-            parents = self.g.get_all_parents(node)
+            parents = self.g.reachable_parents(node)
             for parent in parents:
                 irb = self.blocs[parent]
                 for var_w in irb.w:
diff --git a/test/core/graph.py b/test/core/graph.py
index eb896cfe..40f34f5c 100644
--- a/test/core/graph.py
+++ b/test/core/graph.py
@@ -32,10 +32,48 @@ 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]))
+dominators = g1.compute_dominators(1)
+assert(dominators == {1: set([1]),
+                      2: set([1, 2]),
+                      3: set([1, 2, 3]),
+                      4: set([1, 2, 4]),
+                      5: set([1, 2, 5]),
+                      6: set([1, 2, 6])})
+
+# Regression test with multiple heads
+g2 = DiGraph()
+g2.add_edge(1, 2)
+g2.add_edge(2, 3)
+g2.add_edge(3, 4)
+g2.add_edge(5, 6)
+g2.add_edge(6, 3)
+
+dominators = g2.compute_dominators(5)
+assert(dominators == {3: set([3, 5, 6]),
+                      4: set([3, 4, 5, 6]),
+                      5: set([5]),
+                      6: set([5, 6])})
+
+
+postdominators = g1.compute_postdominators(6)
+assert(postdominators == {1: set([1, 2, 6]),
+                          2: set([2, 6]),
+                          3: set([2, 3, 5, 6]),
+                          4: set([2, 4, 5, 6]),
+                          5: set([2, 5, 6]),
+                          6: set([6])})
+
+postdominators = g1.compute_postdominators(5)
+assert(postdominators == {1: set([1, 2, 5]),
+                          2: set([2, 5]),
+                          3: set([3, 5]),
+                          4: set([4, 5]),
+                          5: set([5])})
+
+postdominators = g2.compute_postdominators(4)
+assert(postdominators == {1: set([1, 2, 3, 4]),
+                          2: set([2, 3, 4]),
+                          3: set([3, 4]),
+                          4: set([4]),
+                          5: set([3, 4, 5, 6]),
+                          6: set([3, 4, 6])})