IR: Make IRBlock immutable

3 files changed, 34 insertions, 31 deletions

diff --git a/miasm2/arch/mips32/ira.py b/miasm2/arch/mips32/ira.py
index a2eab4fb..e342a6fd 100644
--- a/miasm2/arch/mips32/ira.py
+++ b/miasm2/arch/mips32/ira.py
@@ -11,12 +11,13 @@ class ir_a_mips32l(ir_mips32l, ira):
         ir_mips32l.__init__(self, symbol_pool)
         self.ret_reg = self.arch.regs.V0
 
-    def pre_add_instr(self, block, instr, irb_cur, ir_blocks_all, gen_pc_updt):
+    def pre_add_instr(self, block, instr, assignments, ir_blocks_all, gen_pc_updt):
         # Avoid adding side effects, already done in post_add_bloc
-        return irb_cur
+        return False
 
     def post_add_bloc(self, block, ir_blocks):
         IntermediateRepresentation.post_add_bloc(self, block, ir_blocks)
+        new_irblocks = []
         for irb in ir_blocks:
             pc_val = None
             lr_val = None
@@ -25,13 +26,15 @@ class ir_a_mips32l(ir_mips32l, ira):
                 lr_val = assignblk.get(self.arch.regs.RA, lr_val)
 
             if pc_val is None or lr_val is None:
+                new_irblocks.append(irb)
                 continue
             if not expr_is_int_or_label(lr_val):
+                new_irblocks.append(irb)
                 continue
             if expr_is_label(lr_val):
                 lr_val = ExprInt(lr_val.name.offset, 32)
 
-            instr = block.irs[-2].instr
+            instr = block.lines[-2]
             if lr_val.arg != instr.offset + 8:
                 raise ValueError("Wrong arg")
 
@@ -42,9 +45,9 @@ class ir_a_mips32l(ir_mips32l, ira):
             irs.append(AssignBlock([ExprAff(self.IRDst,
                                             ExprId(lbl, size=self.pc.size))],
                                    instr))
-            nblock = IRBlock(new_lbl, irs)
-            self.blocks[new_lbl] = nblock
-            irb.dst = ExprId(new_lbl, size=self.pc.size)
+            new_irblocks.append(IRBlock(new_lbl, irs))
+            new_irblocks.append(irb.set_dst(ExprId(new_lbl, size=self.pc.size)))
+        return new_irblocks
 
     def get_out_regs(self, _):
         return set([self.ret_reg, self.sp])
diff --git a/miasm2/arch/mips32/jit.py b/miasm2/arch/mips32/jit.py
index 9b46589f..493da595 100644
--- a/miasm2/arch/mips32/jit.py
+++ b/miasm2/arch/mips32/jit.py
@@ -5,7 +5,7 @@ from miasm2.core import asmblock
 from miasm2.core.utils import pck32, upck32
 from miasm2.arch.mips32.sem import ir_mips32l, ir_mips32b
 from miasm2.jitter.codegen import CGen
-from miasm2.ir.ir import AssignBlock
+from miasm2.ir.ir import AssignBlock, IRBlock
 import miasm2.expression.expression as m2_expr
 
 log = logging.getLogger('jit_mips32')
@@ -40,24 +40,27 @@ class mipsCGen(CGen):
 
     def block2assignblks(self, block):
         irblocks_list = super(mipsCGen, self).block2assignblks(block)
-        for instr, irblocks in zip(block.lines, irblocks_list):
-            if not instr.breakflow():
-                continue
-            for irblock in irblocks:
-                for idx, assignblock in enumerate(irblock.irs):
+        for irblocks in irblocks_list:
+            for blk_idx, irblock in enumerate(irblocks):
+                has_breakflow = any(assignblock.instr.breakflow() for assignblock in irblock.irs)
+                if not has_breakflow:
+                    continue
+
+                irs = []
+                for assignblock in irblock.irs:
                     if self.ir_arch.pc not in assignblock:
+                        irs.append(AssignBlock(assignments, assignblock.instr))
                         continue
-                    new_assignblock = dict(assignblock)
+                    assignments = dict(assignblock)
                     # Add internal branch destination
-                    new_assignblock[self.delay_slot_dst] = assignblock[
+                    assignments[self.delay_slot_dst] = assignblock[
                         self.ir_arch.pc]
-                    new_assignblock[self.delay_slot_set] = m2_expr.ExprInt(1, 32)
+                    assignments[self.delay_slot_set] = m2_expr.ExprInt(1, 32)
                     # Replace IRDst with next instruction
-                    new_assignblock[self.ir_arch.IRDst] = m2_expr.ExprId(
-                        self.ir_arch.get_next_instr(instr))
-                    irblock.dst = m2_expr.ExprId(
-                        self.ir_arch.get_next_instr(instr))
-                    irblock.irs[idx] = AssignBlock(new_assignblock, assignblock.instr)
+                    assignments[self.ir_arch.IRDst] = m2_expr.ExprId(
+                        self.ir_arch.get_next_instr(assignblock.instr))
+                    irs.append(AssignBlock(assignments, assignblock.instr))
+                irblocks[blk_idx] = IRBlock(irblock.label, irs)
 
         return irblocks_list
 
diff --git a/miasm2/arch/mips32/sem.py b/miasm2/arch/mips32/sem.py
index bc050b38..645f9a4f 100644
--- a/miasm2/arch/mips32/sem.py
+++ b/miasm2/arch/mips32/sem.py
@@ -441,17 +441,14 @@ class ir_mips32l(IntermediateRepresentation):
         args = instr.args
         instr_ir, extra_ir = get_mnemo_expr(self, instr, *args)
 
-        for i, x in enumerate(instr_ir):
-            x = m2_expr.ExprAff(x.dst, x.src.replace_expr(
-                {self.pc: m2_expr.ExprInt(instr.offset + 4, 32)}))
-            instr_ir[i] = x
-        for irblock in extra_ir:
-            for irs in irblock.irs:
-                for i, x in enumerate(irs):
-                    x = m2_expr.ExprAff(x.dst, x.src.replace_expr(
-                        {self.pc: m2_expr.ExprInt(instr.offset + 4, 32)}))
-                    irs[i] = x
-        return instr_ir, extra_ir
+        pc_fixed = {self.pc: m2_expr.ExprInt(instr.offset + 4, 32)}
+
+        instr_ir = [m2_expr.ExprAff(expr.dst, expr.src.replace_expr(pc_fixed))
+                    for expr in instr_ir]
+
+        new_extra_ir = [irblock.modify_exprs(mod_src=lambda expr: expr.replace_expr(pc_fixed))
+                        for irblock in extra_ir]
+        return instr_ir, new_extra_ir
 
     def get_next_instr(self, instr):
         return self.symbol_pool.getby_offset_create(instr.offset  + 4)