IR: move graph relative code from IRA to IR

2 files changed, 99 insertions, 98 deletions

diff --git a/miasm2/ir/analysis.py b/miasm2/ir/analysis.py
index 31f6294c..dcdac29a 100644
--- a/miasm2/ir/analysis.py
+++ b/miasm2/ir/analysis.py
@@ -4,7 +4,6 @@
 import logging
 
 from miasm2.ir.symbexec import symbexec
-from miasm2.core.graph import DiGraph
 from miasm2.expression.expression \
     import ExprAff, ExprCond, ExprId, ExprInt, ExprMem
 
@@ -20,103 +19,6 @@ class ira:
         """Returns ids of all registers used in the IR"""
         return self.arch.regs.all_regs_ids + [self.IRDst]
 
-    def sort_dst(self, todo, done):
-        out = set()
-        while todo:
-            dst = todo.pop()
-            if self.ExprIsLabel(dst):
-                done.add(dst)
-            elif isinstance(dst, ExprMem) or isinstance(dst, ExprInt):
-                done.add(dst)
-            elif isinstance(dst, ExprCond):
-                todo.add(dst.src1)
-                todo.add(dst.src2)
-            elif isinstance(dst, ExprId):
-                out.add(dst)
-            else:
-                done.add(dst)
-        return out
-
-    def dst_trackback(self, b):
-        dst = b.dst
-        todo = set([dst])
-        done = set()
-
-        for irs in reversed(b.irs):
-            if len(todo) == 0:
-                break
-            out = self.sort_dst(todo, done)
-            found = set()
-            follow = set()
-            for i in irs:
-                if not out:
-                    break
-                for o in out:
-                    if i.dst == o:
-                        follow.add(i.src)
-                        found.add(o)
-                for o in found:
-                    out.remove(o)
-
-            for o in out:
-                if o not in found:
-                    follow.add(o)
-            todo = follow
-
-        return done
-
-    def gen_graph(self, link_all = True):
-        """
-        Gen irbloc digraph
-        @link_all: also gen edges to non present irblocs
-        """
-        self.g = DiGraph()
-        for lbl, b in self.blocs.items():
-            # print 'add', lbl
-            self.g.add_node(lbl)
-            # dst = self.get_bloc_dst(b)
-            dst = self.dst_trackback(b)
-            # print "\tdst", dst
-            for d in dst:
-                if isinstance(d, ExprInt):
-                    d = ExprId(
-                        self.symbol_pool.getby_offset_create(int(d.arg)))
-                if self.ExprIsLabel(d):
-                    if d.name in self.blocs or link_all is True:
-                        self.g.add_edge(lbl, d.name)
-
-    def graph(self):
-        """Output the graphviz script"""
-        out = """
-    digraph asm_graph {
-    size="80,50";
-    node [
-    fontsize = "16",
-    shape = "box"
-    ];
-        """
-        all_lbls = {}
-        for lbl in self.g.nodes():
-            if lbl not in self.blocs:
-                continue
-            irb = self.blocs[lbl]
-            ir_txt = [str(lbl)]
-            for irs in irb.irs:
-                for l in irs:
-                    ir_txt.append(str(l))
-                ir_txt.append("")
-            ir_txt.append("")
-            all_lbls[hash(lbl)] = "\l\\\n".join(ir_txt)
-        for l, v in all_lbls.items():
-            # print l, v
-            out += '%s [label="%s"];\n' % (l, v)
-
-        for a, b in self.g.edges():
-            # print 'edge', a, b, hash(a), hash(b)
-            out += '%s -> %s;\n' % (hash(a), hash(b))
-        out += '}'
-        return out
-
     def remove_dead_instr(self, irb, useful):
         """Remove dead affectations using previous reaches analysis
         @irb: irbloc instance
diff --git a/miasm2/ir/ir.py b/miasm2/ir/ir.py
index e051dc8c..31994fae 100644
--- a/miasm2/ir/ir.py
+++ b/miasm2/ir/ir.py
@@ -25,6 +25,7 @@ from miasm2.expression.expression_helper import get_missing_interval
 from miasm2.core import asmbloc
 from miasm2.expression.simplifications import expr_simp
 from miasm2.core.asmbloc import asm_symbol_pool
+from miasm2.core.graph import DiGraph
 
 
 class irbloc(object):
@@ -335,5 +336,103 @@ class ir(object):
             b.get_rw(regs_ids)
 
     def ExprIsLabel(self, l):
+        #TODO : use expression helper
         return isinstance(l, m2_expr.ExprId) and isinstance(l.name,
                                                             asmbloc.asm_label)
+
+    def sort_dst(self, todo, done):
+        out = set()
+        while todo:
+            dst = todo.pop()
+            if self.ExprIsLabel(dst):
+                done.add(dst)
+            elif isinstance(dst, m2_expr.ExprMem) or isinstance(dst, m2_expr.ExprInt):
+                done.add(dst)
+            elif isinstance(dst, m2_expr.ExprCond):
+                todo.add(dst.src1)
+                todo.add(dst.src2)
+            elif isinstance(dst, m2_expr.ExprId):
+                out.add(dst)
+            else:
+                done.add(dst)
+        return out
+
+    def dst_trackback(self, b):
+        dst = b.dst
+        todo = set([dst])
+        done = set()
+
+        for irs in reversed(b.irs):
+            if len(todo) == 0:
+                break
+            out = self.sort_dst(todo, done)
+            found = set()
+            follow = set()
+            for i in irs:
+                if not out:
+                    break
+                for o in out:
+                    if i.dst == o:
+                        follow.add(i.src)
+                        found.add(o)
+                for o in found:
+                    out.remove(o)
+
+            for o in out:
+                if o not in found:
+                    follow.add(o)
+            todo = follow
+
+        return done
+
+    def gen_graph(self, link_all = True):
+        """
+        Gen irbloc digraph
+        @link_all: also gen edges to non present irblocs
+        """
+        self.g = DiGraph()
+        for lbl, b in self.blocs.items():
+            # print 'add', lbl
+            self.g.add_node(lbl)
+            # dst = self.get_bloc_dst(b)
+            dst = self.dst_trackback(b)
+            # print "\tdst", dst
+            for d in dst:
+                if isinstance(d, m2_expr.ExprInt):
+                    d = m2_expr.ExprId(
+                        self.symbol_pool.getby_offset_create(int(d.arg)))
+                if self.ExprIsLabel(d):
+                    if d.name in self.blocs or link_all is True:
+                        self.g.add_edge(lbl, d.name)
+
+    def graph(self):
+        """Output the graphviz script"""
+        out = """
+    digraph asm_graph {
+    size="80,50";
+    node [
+    fontsize = "16",
+    shape = "box"
+    ];
+        """
+        all_lbls = {}
+        for lbl in self.g.nodes():
+            if lbl not in self.blocs:
+                continue
+            irb = self.blocs[lbl]
+            ir_txt = [str(lbl)]
+            for irs in irb.irs:
+                for l in irs:
+                    ir_txt.append(str(l))
+                ir_txt.append("")
+            ir_txt.append("")
+            all_lbls[hash(lbl)] = "\l\\\n".join(ir_txt)
+        for l, v in all_lbls.items():
+            # print l, v
+            out += '%s [label="%s"];\n' % (l, v)
+
+        for a, b in self.g.edges():
+            # print 'edge', a, b, hash(a), hash(b)
+            out += '%s -> %s;\n' % (hash(a), hash(b))
+        out += '}'
+        return out