Merge pull request #173 from serpilliere/walk_dominator

Walk dominator

2 files changed, 107 insertions, 0 deletions

diff --git a/miasm2/core/graph.py b/miasm2/core/graph.py
index efb938e6..47be1864 100644
--- a/miasm2/core/graph.py
+++ b/miasm2/core/graph.py
@@ -229,3 +229,74 @@ shape = "box"
                                                 self.reachable_parents,
                                                 self.successors_iter,
                                                 self.predecessors_iter)
+
+    @staticmethod
+    def _walk_generic_dominator(node, gen_dominators, succ_cb):
+        """Generic algorithm to return an iterator of the ordered list of
+        @node's dominators/post_dominator.
+
+        The function doesn't return the self reference in dominators.
+        @node: The start node
+        @gen_dominators: The dictionnary containing at least node's
+        dominators/post_dominators
+        @succ_cb: return predecessors/succesors of a node
+
+        """
+        # Init
+        done = set()
+        if node not in gen_dominators:
+            # We are in a branch which doesn't reach head
+            return
+        node_gen_dominators = set(gen_dominators[node])
+        todo = set([node])
+
+        # Avoid working on itself
+        node_gen_dominators.remove(node)
+
+        # For each level
+        while node_gen_dominators:
+            new_node = None
+
+            # Worklist pattern
+            while todo:
+                node = todo.pop()
+                if node in done:
+                    continue
+                if node in node_gen_dominators:
+                    new_node = node
+                    break
+
+                # Avoid loops
+                done.add(node)
+
+                # Look for the next level
+                for pred in succ_cb(node):
+                    todo.add(pred)
+
+            # Return the node; it's the next starting point
+            assert(new_node is not None)
+            yield new_node
+            node_gen_dominators.remove(new_node)
+            todo = set([new_node])
+
+    def walk_dominators(self, node, dominators):
+        """Return an iterator of the ordered list of @node's dominators
+        The function doesn't return the self reference in dominators.
+        @node: The start node
+        @dominators: The dictionnary containing at least node's dominators
+        """
+        return self._walk_generic_dominator(node,
+                                            dominators,
+                                            self.predecessors_iter)
+
+    def walk_postdominators(self, node, postdominators):
+        """Return an iterator of the ordered list of @node's postdominators
+        The function doesn't return the self reference in postdominators.
+        @node: The start node
+        @postdominators: The dictionnary containing at least node's
+        postdominators
+
+        """
+        return self._walk_generic_dominator(node,
+                                            postdominators,
+                                            self.successors_iter)
diff --git a/test/core/graph.py b/test/core/graph.py
index 40f34f5c..eb320542 100644
--- a/test/core/graph.py
+++ b/test/core/graph.py
@@ -40,6 +40,13 @@ assert(dominators == {1: set([1]),
                       5: set([1, 2, 5]),
                       6: set([1, 2, 6])})
 
+assert(list(g1.walk_dominators(1, dominators)) == [])
+assert(list(g1.walk_dominators(2, dominators)) == [1])
+assert(list(g1.walk_dominators(3, dominators)) == [2, 1])
+assert(list(g1.walk_dominators(4, dominators)) == [2, 1])
+assert(list(g1.walk_dominators(5, dominators)) == [2, 1])
+assert(list(g1.walk_dominators(6, dominators)) == [2, 1])
+
 # Regression test with multiple heads
 g2 = DiGraph()
 g2.add_edge(1, 2)
@@ -55,6 +62,13 @@ assert(dominators == {3: set([3, 5, 6]),
                       6: set([5, 6])})
 
 
+assert(list(g2.walk_dominators(1, dominators)) == [])
+assert(list(g2.walk_dominators(2, dominators)) == [])
+assert(list(g2.walk_dominators(3, dominators)) == [6, 5])
+assert(list(g2.walk_dominators(4, dominators)) == [3, 6, 5])
+assert(list(g2.walk_dominators(5, dominators)) == [])
+assert(list(g2.walk_dominators(6, dominators)) == [5])
+
 postdominators = g1.compute_postdominators(6)
 assert(postdominators == {1: set([1, 2, 6]),
                           2: set([2, 6]),
@@ -63,6 +77,14 @@ assert(postdominators == {1: set([1, 2, 6]),
                           5: set([2, 5, 6]),
                           6: set([6])})
 
+assert(list(g1.walk_postdominators(1, postdominators)) == [2, 6])
+assert(list(g1.walk_postdominators(2, postdominators)) == [6])
+assert(list(g1.walk_postdominators(3, postdominators)) == [5, 2, 6])
+assert(list(g1.walk_postdominators(4, postdominators)) == [5, 2, 6])
+assert(list(g1.walk_postdominators(5, postdominators)) == [2, 6])
+assert(list(g1.walk_postdominators(6, postdominators)) == [])
+
+
 postdominators = g1.compute_postdominators(5)
 assert(postdominators == {1: set([1, 2, 5]),
                           2: set([2, 5]),
@@ -70,6 +92,13 @@ assert(postdominators == {1: set([1, 2, 5]),
                           4: set([4, 5]),
                           5: set([5])})
 
+assert(list(g1.walk_postdominators(1, postdominators)) == [2, 5])
+assert(list(g1.walk_postdominators(2, postdominators)) == [5])
+assert(list(g1.walk_postdominators(3, postdominators)) == [5])
+assert(list(g1.walk_postdominators(4, postdominators)) == [5])
+assert(list(g1.walk_postdominators(5, postdominators)) == [])
+assert(list(g1.walk_postdominators(6, postdominators)) == [])
+
 postdominators = g2.compute_postdominators(4)
 assert(postdominators == {1: set([1, 2, 3, 4]),
                           2: set([2, 3, 4]),
@@ -77,3 +106,10 @@ assert(postdominators == {1: set([1, 2, 3, 4]),
                           4: set([4]),
                           5: set([3, 4, 5, 6]),
                           6: set([3, 4, 6])})
+
+assert(list(g2.walk_postdominators(1, postdominators)) == [2, 3, 4])
+assert(list(g2.walk_postdominators(2, postdominators)) == [3, 4])
+assert(list(g2.walk_postdominators(3, postdominators)) == [4])
+assert(list(g2.walk_postdominators(4, postdominators)) == [])
+assert(list(g2.walk_postdominators(5, postdominators)) == [6, 3, 4])
+assert(list(g2.walk_postdominators(6, postdominators)) == [3, 4])