2 files changed, 17 insertions, 12 deletions
diff --git a/example/ida/ctype_propagation.py b/example/ida/ctype_propagation.py
index e8b52e3e..3c8a745a 100644
--- a/example/ida/ctype_propagation.py
+++ b/example/ida/ctype_propagation.py
@@ -230,11 +230,14 @@ def get_ira_call_fixer(ira):
             stk_after = idc.GetSpd(instr.offset + instr.l)
             stk_diff = stk_after - stk_before
             print hex(stk_diff)
-            return [AssignBlock([ExprAff(self.ret_reg, ExprOp('call_func_ret', ad)),
-                                 ExprAff(self.sp, self.sp + ExprInt(stk_diff, self.sp.size))
-                                 ],
-                                instr
-                                )]
+            call_assignblk = AssignBlock(
+                [
+                    ExprAff(self.ret_reg, ExprOp('call_func_ret', ad)),
+                    ExprAff(self.sp, self.sp + ExprInt(stk_diff, self.sp.size))
+                ],
+                instr
+            )
+            return [call_assignblk], []
 
     return iraCallStackFixer
 
diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py
index afd00d5c..3f9a43f2 100644
--- a/example/ida/graph_ir.py
+++ b/example/ida/graph_ir.py
@@ -8,7 +8,8 @@ import idautils
 from miasm2.core.bin_stream_ida import bin_stream_ida
 from miasm2.core.asmblock import is_int
 from miasm2.expression.simplifications import expr_simp
-from miasm2.analysis.data_flow import dead_simp
+from miasm2.analysis.data_flow import dead_simp, remove_empty_assignblks, \
+    merge_blocks
 from miasm2.ir.ir import AssignBlock, IRBlock
 
 from utils import guess_machine, expr2colorstr
@@ -101,7 +102,7 @@ def build_graph(verbose=False, simplify=False):
     start_addr = idc.ScreenEA()
 
     machine = guess_machine(addr=start_addr)
-    mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
+    dis_engine, ira = machine.dis_engine, machine.ira
 
     if verbose:
         print "Arch", dis_engine
@@ -131,6 +132,8 @@ def build_graph(verbose=False, simplify=False):
 
     asmcfg = mdis.dis_multiblock(start_addr)
 
+    entry_points = set([start_addr])
+
     if verbose:
         print "generating graph"
         open('asm_flow.dot', 'w').write(asmcfg.dot())
@@ -165,14 +168,13 @@ def build_graph(verbose=False, simplify=False):
         while modified:
             modified = False
             modified |= dead_simp(ir_arch, ircfg)
-            modified |= ircfg.remove_empty_assignblks()
-            modified |= ircfg.remove_jmp_blocks()
-            modified |= ircfg.merge_blocks()
+            modified |= remove_empty_assignblks(ircfg)
+            modified |= merge_blocks(ircfg, entry_points)
         title += " (simplified)"
 
-    g = GraphMiasmIR(ircfg, title, None)
+    graph = GraphMiasmIR(ircfg, title, None)
 
-    g.Show()
+    graph.Show()
 
 if __name__ == "__main__":
     build_graph(verbose=True, simplify=False)