IDA: fix depgraph for immutable AssignBlock

1 files changed, 12 insertions, 8 deletions

diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
index faec2857..bfd7cabc 100644
--- a/example/ida/depgraph.py
+++ b/example/ida/depgraph.py
@@ -7,6 +7,7 @@ from miasm2.expression import expression as m2_expr
 
 from miasm2.expression.simplifications import expr_simp
 from miasm2.analysis.depgraph import DependencyGraph
+from miasm2.ir.ir import AssignBlock
 
 from utils import guess_machine
 
@@ -22,9 +23,11 @@ class depGraphSettingsForm(Form):
 
         self.address = ScreenEA()
         cur_block = list(ira.getby_offset(self.address))[0]
-        for line_nb, l in enumerate(cur_block.lines):
-            if l.offset == self.address:
+        line_nb = None
+        for line_nb, assignblk in enumerate(cur_block.irs):
+            if assignblk.instr.offset == self.address:
                 break
+        assert line_nb is not None
         cur_label = str(cur_block.label)
         labels = sorted(map(str, ira.blocks.keys()))
         regs = sorted(ir_arch.arch.regs.all_regs_ids_byname.keys())
@@ -102,7 +105,7 @@ Method to use:
     def elements(self):
         value = self.cbReg.value
         if value in self.stk_args:
-            line = self.ira.blocks[self.label].lines[self.line_nb]
+            line = self.ira.blocks[self.label].irs[self.line_nb].instr
             arg_num = self.stk_args[value]
             stk_high = m2_expr.ExprInt(GetSpd(line.offset), ir_arch.sp.size)
             stk_off = m2_expr.ExprInt(self.ira.sp.size/8 * arg_num, ir_arch.sp.size)
@@ -167,17 +170,18 @@ for irb in ir_arch.blocks.values():
     fix_stack = irb.label.offset is not None and settings.unalias_stack
     for i, assignblk in enumerate(irb.irs):
         if fix_stack:
-            stk_high = m2_expr.ExprInt(GetSpd(irb.lines[i].offset), ir_arch.sp.size)
+            stk_high = m2_expr.ExprInt(GetSpd(irb.irs[i].instr.offset), ir_arch.sp.size)
             fix_dct = {ir_arch.sp: mn.regs.regs_init[ir_arch.sp] + stk_high}
 
-        for dst, src in assignblk.items():
-            del(assignblk[dst])
+        new_assignblk = {}
+        for dst, src in assignblk.iteritems():
             if fix_stack:
                 src = src.replace_expr(fix_dct)
                 if dst != ir_arch.sp:
                     dst = dst.replace_expr(fix_dct)
             dst, src = expr_simp(dst), expr_simp(src)
-            assignblk[dst] = src
+            new_assignblk[dst] = src
+        irb.irs[i] = AssignBlock(new_assignblk, instr=assignblk.instr)
 
 # Get dependency graphs
 dg = settings.depgraph
@@ -215,7 +219,7 @@ def treat_element():
 
     for node in graph.relevant_nodes:
         try:
-            offset = ir_arch.blocks[node.label].lines[node.line_nb].offset
+            offset = ir_arch.blocks[node.label].irs[node.line_nb].instr.offset
         except IndexError:
             print "Unable to highlight %s" % node
             continue