Merge pull request #814 from serpilliere/fix_graph_reduce

Fix graph reduce

2 files changed, 128 insertions, 8 deletions

diff --git a/miasm2/analysis/data_flow.py b/miasm2/analysis/data_flow.py
index bba598e0..0a224319 100644
--- a/miasm2/analysis/data_flow.py
+++ b/miasm2/analysis/data_flow.py
@@ -354,7 +354,9 @@ def _relink_block_node(ircfg, loc_key, son_loc_key, replace_dct):
     Link loc_key's parents to parents directly to son_loc_key
     """
     for parent in set(ircfg.predecessors(loc_key)):
-        parent_block = ircfg.blocks[parent]
+        parent_block = ircfg.blocks.get(parent, None)
+        if parent_block is None:
+            continue
 
         new_block = parent_block.modify_exprs(
             lambda expr:expr.replace_expr(replace_dct),
@@ -472,18 +474,22 @@ def merge_blocks(ircfg, loc_key_entries):
         # Test jmp only block
         son = _test_jmp_only(ircfg, loc_key)
         if son is not None and loc_key not in loc_key_entries:
-            modified |= _remove_to_son(ircfg, loc_key, son)
-            todo.add(loc_key)
-            continue
+            ret = _remove_to_son(ircfg, loc_key, son)
+            modified |= ret
+            if ret:
+                todo.add(loc_key)
+                continue
 
         # Test head jmp only block
         if (son is not None and
             son not in loc_key_entries and
             son in ircfg.blocks):
             # jmp only test done previously
-            modified |= _remove_to_parent(ircfg, loc_key, son)
-            todo.add(loc_key)
-            continue
+            ret = _remove_to_parent(ircfg, loc_key, son)
+            modified |= ret
+            if ret:
+                todo.add(loc_key)
+                continue
 
 
     return modified
diff --git a/test/ir/reduce_graph.py b/test/ir/reduce_graph.py
index d8b78c90..68a3aea8 100644
--- a/test/ir/reduce_graph.py
+++ b/test/ir/reduce_graph.py
@@ -507,6 +507,118 @@ for irb in [G6_RES_IRB0, G6_RES_IRB2, G6_RES_IRB3  ]:
 
 
 
+
+########## G7 ##########
+# Input
+
+G7 = IRA.new_ircfg()
+
+G7_IRB0 = gen_irblock(
+    LBL0,
+    [
+        [
+            ExprAff(A, C),
+            ExprAff(IRDst, ExprLoc(LBL1, 32)),
+        ]
+    ]
+)
+
+G7_IRB1 = gen_irblock(
+    LBL1,
+    [
+        [
+            ExprAff(IRDst, ExprLoc(LBL1, 32)),
+        ]
+    ]
+)
+
+
+for irb in [G7_IRB0, G7_IRB1]:
+    G7.add_irblock(irb)
+
+
+# Result
+G7_RES = IRA.new_ircfg()
+
+
+
+G7_RES_IRB0 = gen_irblock(
+    LBL0,
+    [
+        [
+            ExprAff(A, C),
+            ExprAff(IRDst, ExprLoc(LBL1, 32)),
+        ]
+    ]
+)
+
+G7_RES_IRB1 = gen_irblock(
+    LBL1,
+    [
+        [
+            ExprAff(IRDst, ExprLoc(LBL1, 32)),
+        ]
+    ]
+)
+
+
+for irb in [G7_RES_IRB0, G7_RES_IRB1]:
+    G7_RES.add_irblock(irb)
+
+
+
+########## G8 ##########
+# Input
+
+G8 = IRA.new_ircfg()
+
+G8_IRB0 = gen_irblock(
+    LBL0,
+    [
+        [
+            ExprAff(IRDst, ExprLoc(LBL1, 32)),
+        ]
+    ]
+)
+
+G8_IRB1 = gen_irblock(
+    LBL1,
+    [
+        [
+            ExprAff(A, C),
+            ExprAff(IRDst, ExprLoc(LBL1, 32)),
+        ]
+    ]
+)
+
+
+for irb in [G8_IRB0, G8_IRB1]:
+    G8.add_irblock(irb)
+
+
+# Result
+G8_RES = IRA.new_ircfg()
+
+
+
+G8_RES_IRB0 = gen_irblock(
+    LBL0,
+    [
+        [
+            ExprAff(A, C),
+            ExprAff(IRDst, ExprLoc(LBL0, 32)),
+        ]
+    ]
+)
+
+
+for irb in [G8_RES_IRB0]:
+    G8_RES.add_irblock(irb)
+
+
+
+
+
 ################# Tests
 
 
@@ -518,14 +630,16 @@ for i, (g_test, g_ref) in enumerate(
             (G4, G4_RES),
             (G5, G5_RES),
             (G6, G6_RES),
+            (G7, G7_RES),
+            (G8, G8_RES),
         ], 1):
 
     heads = g_test.heads()
     print '*'*10, 'Test', i, "*"*10
     open('test_in_%d.dot' % i, 'w').write(g_test.dot())
+    open('test_ref_%d.dot' % i, 'w').write(g_ref.dot())
     merge_blocks(g_test, heads)
     open('test_out_%d.dot' % i, 'w').write(g_test.dot())
-    open('test_ref_%d.dot' % i, 'w').write(g_ref.dot())
 
     cmp_ir_graph(g_test, g_ref)
     print '\t', 'OK'