Graph: add dominators computation

1 files changed, 52 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)