Merge pull request #814 from serpilliere/fix_graph_reduce

Fix graph reduce

1 files changed, 13 insertions, 7 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