Merge pull request #537 from commial/fix/ida-examples

Fix/ida examples

1 files changed, 9 insertions, 6 deletions

diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py
index d2552c82..6ff4304a 100644
--- a/example/ida/graph_ir.py
+++ b/example/ida/graph_ir.py
@@ -11,6 +11,8 @@ from miasm2.expression.expression import *
 from miasm2.analysis.data_analysis import inter_bloc_flow, \
     intra_bloc_flow_symbexec
 from miasm2.analysis.data_flow import dead_simp
+from miasm2.ir.ir import AssignBlock
+
 from utils import guess_machine, expr2colorstr
 
 
@@ -133,12 +135,13 @@ for block in ab:
 
 print "IR ok... %x" % ad
 
-for irb in ir_arch.blocks.values():
-    for assignblk in irb.irs:
-        for dst, src in assignblk.items():
-            del(assignblk[dst])
-            dst, src = expr_simp(dst), expr_simp(src)
-            assignblk[dst] = src
+for irb in ir_arch.blocks.itervalues():
+    for i, assignblk in enumerate(irb.irs):
+        new_assignblk = {
+            expr_simp(dst): expr_simp(src)
+            for dst, src in assignblk.iteritems()
+        }
+        irb.irs[i] = AssignBlock(new_assignblk, instr=assignblk.instr)
 
 out = ir_arch.graph.dot()
 open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out)