Merge pull request #1159 from serpilliere/graph_weakly_connected

Graph: add weakly connected components

2 files changed, 65 insertions, 0 deletions

diff --git a/miasm/core/graph.py b/miasm/core/graph.py
index 01f580a3..8bb4371d 100644
--- a/miasm/core/graph.py
+++ b/miasm/core/graph.py
@@ -732,6 +732,49 @@ class DiGraph(object):
                         yield scc
 
 
+    def compute_weakly_connected_components(self):
+        """
+        Return the weakly connected components
+        """
+        remaining = set(self.nodes())
+        components = []
+        while remaining:
+            node = remaining.pop()
+            todo = set()
+            todo.add(node)
+            component = set()
+            done = set()
+            while todo:
+                node = todo.pop()
+                if node in done:
+                    continue
+                done.add(node)
+                remaining.discard(node)
+                component.add(node)
+                todo.update(self.predecessors(node))
+                todo.update(self.successors(node))
+            components.append(component)
+        return components
+
+
+
+    def replace_node(self, node, new_node):
+        """
+        Replace @node by @new_node
+        """
+
+        predecessors = self.predecessors(node)
+        successors = self.successors(node)
+        self.del_node(node)
+        for predecessor in predecessors:
+            if predecessor == node:
+                predecessor = new_node
+            self.add_uniq_edge(predecessor, new_node)
+        for successor in successors:
+            if successor == node:
+                successor = new_node
+            self.add_uniq_edge(new_node, successor)
+
 class DiGraphSimplifier(object):
 
     """Wrapper on graph simplification passes.
diff --git a/test/core/graph.py b/test/core/graph.py
index 3db5e523..ff27b780 100644
--- a/test/core/graph.py
+++ b/test/core/graph.py
@@ -286,3 +286,25 @@ assert sols[0] == {j1: 1,
                    j2: 2,
                    j3: 3}
 
+
+
+# Test replace_node
+graph = DiGraph()
+graph.add_edge(1, 2)
+graph.add_edge(2, 2)
+graph.add_edge(2, 3)
+
+graph.replace_node(2, 4)
+assert graph.nodes() == set([1, 3, 4])
+assert sorted(graph.edges()) == [(1, 4), (4, 3), (4, 4)]
+
+
+
+# Test compute_weakly_connected_components
+graph = DiGraph()
+graph.add_edge(1, 2)
+graph.add_edge(2, 2)
+graph.add_edge(3, 4)
+
+components = graph.compute_weakly_connected_components()
+assert sorted(components) == [set([1, 2]), set([3, 4])]