Graph: Add get_all_parents/del_node

1 files changed, 22 insertions, 0 deletions

diff --git a/miasm2/core/graph.py b/miasm2/core/graph.py
index 47047269..ad159771 100644
--- a/miasm2/core/graph.py
+++ b/miasm2/core/graph.py
@@ -27,6 +27,14 @@ class DiGraph:
         self._nodes_to[n] = []
         self._nodes_from[n] = []
 
+    def del_node(self, n):
+        if n in self._nodes:
+            self._nodes.remove(n)
+        for pred in self.predecessors(n):
+            self.del_edge(pred, n)
+        for succ in self.successors(n):
+            self.del_edge(n, succ)
+
     def add_edge(self, a, b):
         if not a in self._nodes:
             self.add_node(a)
@@ -97,6 +105,20 @@ class DiGraph:
                     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)