IR: Make IRBlock immutable

7 files changed, 113 insertions, 118 deletions

diff --git a/miasm2/arch/aarch64/sem.py b/miasm2/arch/aarch64/sem.py
index 79c72d32..d5209e3e 100644
--- a/miasm2/arch/aarch64/sem.py
+++ b/miasm2/arch/aarch64/sem.py
@@ -782,14 +782,14 @@ class ir_aarch64l(IntermediateRepresentation):
         src = self.expr_fix_regs_for_mode(e.src)
         return m2_expr.ExprAff(dst, src)
 
-    def irbloc_fix_regs_for_mode(self, irbloc, mode=64):
-        for idx, assignblk in enumerate(irbloc.irs):
+    def irbloc_fix_regs_for_mode(self, irblock, mode=64):
+        irs = []
+        for assignblk in irblock.irs:
             new_assignblk = dict(assignblk)
             for dst, src in assignblk.iteritems():
                 del(new_assignblk[dst])
                 # Special case for 64 bits:
                 # If destination is a 32 bit reg, zero extend the 64 bit reg
-
                 if (isinstance(dst, m2_expr.ExprId) and
                     dst.size == 32 and
                     dst in replace_regs):
@@ -799,27 +799,24 @@ class ir_aarch64l(IntermediateRepresentation):
                 dst = self.expr_fix_regs_for_mode(dst)
                 src = self.expr_fix_regs_for_mode(src)
                 new_assignblk[dst] = src
-            irbloc.irs[idx] = AssignBlock(new_assignblk, assignblk.instr)
-        if irbloc.dst is not None:
-            irbloc.dst = self.expr_fix_regs_for_mode(irbloc.dst)
+            irs.append(AssignBlock(new_assignblk, assignblk.instr))
+        return IRBlock(irblock.label, irs)
 
     def mod_pc(self, instr, instr_ir, extra_ir):
         "Replace PC by the instruction's offset"
         cur_offset = m2_expr.ExprInt(instr.offset, 64)
+        pc_fixed = {self.pc: cur_offset}
         for i, expr in enumerate(instr_ir):
             dst, src = expr.dst, expr.src
             if dst != self.pc:
-                dst = dst.replace_expr({self.pc: cur_offset})
-            src = src.replace_expr({self.pc: cur_offset})
+                dst = dst.replace_expr(pc_fixed)
+            src = src.replace_expr(pc_fixed)
             instr_ir[i] = m2_expr.ExprAff(dst, src)
-        for irblock in extra_ir:
-            for irs in irblock.irs:
-                for i, expr in enumerate(irs):
-                    dst, src = expr.dst, expr.src
-                    if dst != self.pc:
-                        dst = dst.replace_expr({self.pc: cur_offset})
-                    src = src.replace_expr({self.pc: cur_offset})
-                    irs[i] = m2_expr.ExprAff(dst, src)
+
+        for idx, irblock in enumerate(extra_ir):
+            extra_ir[idx] = irblock.modify_exprs(lambda expr: expr.replace_expr(pc_fixed) \
+                                                 if expr != self.pc else expr,
+                                                 lambda expr: expr.replace_expr(pc_fixed))
 
 
     def del_dst_zr(self, instr, instr_ir, extra_ir):
@@ -827,11 +824,16 @@ class ir_aarch64l(IntermediateRepresentation):
         regs_to_fix = [WZR, XZR]
         instr_ir = [expr for expr in instr_ir if expr.dst not in regs_to_fix]
 
+        new_irblocks = []
         for irblock in extra_ir:
-            for i, irs in enumerate(irblock.irs):
-                irblock.irs[i] = [expr for expr in irs if expr.dst not in regs_to_fix]
-
-        return instr_ir, extra_ir
+            irs = []
+            for assignblk in irblock.irs:
+                new_dsts = {dst:src for dst, src in assignblk.iteritems()
+                                if dst not in regs_to_fix}
+                irs.append(AssignBlock(new_dsts, assignblk.instr))
+            new_irblocks.append(IRBlock(irblock.label, irs))
+
+        return instr_ir, new_irblocks
 
 
 class ir_aarch64b(ir_aarch64l):
diff --git a/miasm2/arch/arm/sem.py b/miasm2/arch/arm/sem.py
index 710cdc9f..29b25538 100644
--- a/miasm2/arch/arm/sem.py
+++ b/miasm2/arch/arm/sem.py
@@ -1,5 +1,5 @@
 from miasm2.expression.expression import *
-from miasm2.ir.ir import IntermediateRepresentation, IRBlock
+from miasm2.ir.ir import IntermediateRepresentation, IRBlock, AssignBlock
 from miasm2.arch.arm.arch import mn_arm, mn_armt
 from miasm2.arch.arm.regs import *
 
@@ -1055,7 +1055,7 @@ def add_condition_expr(ir, instr, cond, instr_ir):
             break
     if not has_irdst:
         instr_ir.append(ExprAff(ir.IRDst, lbl_next))
-    e_do = IRBlock(lbl_do.name, [instr_ir])
+    e_do = IRBlock(lbl_do.name, [AssignBlock(instr_ir, instr)])
     e = [ExprAff(ir.IRDst, dst_cond)]
     return e, [e_do]
 
@@ -1246,20 +1246,15 @@ class ir_arml(IntermediateRepresentation):
                                   args[-1].args[0],
                                   args[-1].args[-1][:8].zeroExtend(32))
         instr_ir, extra_ir = get_mnemo_expr(self, instr, *args)
-        # if self.name.startswith('B'):
-        #    return instr_ir, extra_ir
-        for i, x in enumerate(instr_ir):
-            x = ExprAff(x.dst, x.src.replace_expr(
-                {self.pc: ExprInt(instr.offset + 8, 32)}))
-            instr_ir[i] = x
-        for irblock in extra_ir:
-            for irs in irblock.irs:
-                for i, x in enumerate(irs):
-                    x = ExprAff(x.dst, x.src.replace_expr(
-                        {self.pc: ExprInt(instr.offset + 8, 32)}))
-                    irs[i] = x
-        # return out_ir, extra_ir
-        return instr_ir, extra_ir
+
+        pc_fixed = {self.pc: ExprInt(instr.offset + 8, 32)}
+        for i, expr in enumerate(instr_ir):
+            instr_ir[i] = ExprAff(expr.dst, expr.src.replace_expr(pc_fixed))
+
+        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
 
 
 class ir_armb(ir_arml):
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)
diff --git a/miasm2/arch/x86/jit.py b/miasm2/arch/x86/jit.py
index 6d9be8ac..9acab5ed 100644
--- a/miasm2/arch/x86/jit.py
+++ b/miasm2/arch/x86/jit.py
@@ -45,7 +45,7 @@ class jitter_x86_16(jitter):
         self.ir_arch.irbloc_fix_regs_for_mode = self.ir_archbloc_fix_regs_for_mode
 
     def ir_archbloc_fix_regs_for_mode(self, irblock, attrib=64):
-        self.orig_irbloc_fix_regs_for_mode(irblock, 64)
+        return self.orig_irbloc_fix_regs_for_mode(irblock, 64)
 
     def push_uint16_t(self, value):
         self.cpu.SP -= self.ir_arch.sp.size / 8
@@ -78,7 +78,7 @@ class jitter_x86_32(jitter):
         self.ir_arch.irbloc_fix_regs_for_mode = self.ir_archbloc_fix_regs_for_mode
 
     def ir_archbloc_fix_regs_for_mode(self, irblock, attrib=64):
-        self.orig_irbloc_fix_regs_for_mode(irblock, 64)
+        return self.orig_irbloc_fix_regs_for_mode(irblock, 64)
 
     def push_uint32_t(self, value):
         self.cpu.ESP -= self.ir_arch.sp.size / 8
@@ -183,7 +183,7 @@ class jitter_x86_64(jitter):
         self.ir_arch.irbloc_fix_regs_for_mode = self.ir_archbloc_fix_regs_for_mode
 
     def ir_archbloc_fix_regs_for_mode(self, irblock, attrib=64):
-        self.orig_irbloc_fix_regs_for_mode(irblock, 64)
+        return self.orig_irbloc_fix_regs_for_mode(irblock, 64)
 
     def push_uint64_t(self, value):
         self.cpu.RSP -= self.ir_arch.sp.size / 8
diff --git a/miasm2/arch/x86/sem.py b/miasm2/arch/x86/sem.py
index e32b8001..e1847fe7 100644
--- a/miasm2/arch/x86/sem.py
+++ b/miasm2/arch/x86/sem.py
@@ -263,7 +263,7 @@ def gen_fcmov(ir, instr, cond, arg1, arg2, mov_if):
     e_do, extra_irs = [m2_expr.ExprAff(arg1, arg2)], []
     e_do.append(m2_expr.ExprAff(ir.IRDst, lbl_skip))
     e.append(m2_expr.ExprAff(ir.IRDst, m2_expr.ExprCond(cond, dstA, dstB)))
-    return e, [IRBlock(lbl_do.name, [e_do])]
+    return e, [IRBlock(lbl_do.name, [AssignBlock(e_do, instr)])]
 
 
 def gen_cmov(ir, instr, cond, dst, src, mov_if):
@@ -283,7 +283,7 @@ def gen_cmov(ir, instr, cond, dst, src, mov_if):
     e_do, extra_irs = mov(ir, instr, dst, src)
     e_do.append(m2_expr.ExprAff(ir.IRDst, lbl_skip))
     e.append(m2_expr.ExprAff(ir.IRDst, m2_expr.ExprCond(cond, dstA, dstB)))
-    return e, [IRBlock(lbl_do.name, [e_do])]
+    return e, [IRBlock(lbl_do.name, [AssignBlock(e_do, instr)])]
 
 
 def mov(_, instr, dst, src):
@@ -504,7 +504,7 @@ def _rotate_tpl(ir, instr, dst, src, op, left=False, include_cf=False):
     e_do.append(m2_expr.ExprAff(ir.IRDst, lbl_skip))
     e.append(m2_expr.ExprAff(
         ir.IRDst, m2_expr.ExprCond(shifter, lbl_do, lbl_skip)))
-    return (e, [IRBlock(lbl_do.name, [e_do])])
+    return (e, [IRBlock(lbl_do.name, [AssignBlock(e_do, instr)])])
 
 
 def l_rol(ir, instr, dst, src):
@@ -601,7 +601,7 @@ def _shift_tpl(op, ir, instr, a, b, c=None, op_inv=None, left=False,
     e_do.append(m2_expr.ExprAff(ir.IRDst, lbl_skip))
     e.append(m2_expr.ExprAff(ir.IRDst, m2_expr.ExprCond(shifter, lbl_do,
                                                         lbl_skip)))
-    return e, [IRBlock(lbl_do.name, [e_do])]
+    return e, [IRBlock(lbl_do.name, [AssignBlock(e_do, instr)])]
 
 
 def sar(ir, instr, dst, src):
@@ -949,7 +949,7 @@ def cmps(ir, instr, size):
     e0.append(m2_expr.ExprAff(b.arg,
                               b.arg + m2_expr.ExprInt(size / 8, b.arg.size)))
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e0 = IRBlock(lbl_df_0.name, [e0])
+    e0 = IRBlock(lbl_df_0.name, [AssignBlock(e0, instr)])
 
     e1 = []
     e1.append(m2_expr.ExprAff(a.arg,
@@ -957,7 +957,7 @@ def cmps(ir, instr, size):
     e1.append(m2_expr.ExprAff(b.arg,
                               b.arg - m2_expr.ExprInt(size / 8, b.arg.size)))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e1 = IRBlock(lbl_df_1.name, [e1])
+    e1 = IRBlock(lbl_df_1.name, [AssignBlock(e1, instr)])
 
     e.append(m2_expr.ExprAff(ir.IRDst,
                              m2_expr.ExprCond(df, lbl_df_1, lbl_df_0)))
@@ -978,13 +978,13 @@ def scas(ir, instr, size):
     e0.append(m2_expr.ExprAff(a.arg,
                               a.arg + m2_expr.ExprInt(size / 8, a.arg.size)))
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e0 = IRBlock(lbl_df_0.name, [e0])
+    e0 = IRBlock(lbl_df_0.name, [AssignBlock(e0, instr)])
 
     e1 = []
     e1.append(m2_expr.ExprAff(a.arg,
                               a.arg - m2_expr.ExprInt(size / 8, a.arg.size)))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e1 = IRBlock(lbl_df_1.name, [e1])
+    e1 = IRBlock(lbl_df_1.name, [AssignBlock(e1, instr)])
 
     e.append(m2_expr.ExprAff(ir.IRDst,
                              m2_expr.ExprCond(df, lbl_df_1, lbl_df_0)))
@@ -1455,13 +1455,13 @@ def div(ir, instr, src1):
     do_div = []
     do_div += e
     do_div.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    blk_div = IRBlock(lbl_div.name, [do_div])
+    blk_div = IRBlock(lbl_div.name, [AssignBlock(do_div, instr)])
 
     do_except = []
     do_except.append(m2_expr.ExprAff(exception_flags, m2_expr.ExprInt(
         EXCEPT_DIV_BY_ZERO, exception_flags.size)))
     do_except.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    blk_except = IRBlock(lbl_except.name, [do_except])
+    blk_except = IRBlock(lbl_except.name, [AssignBlock(do_except, instr)])
 
     e = []
     e.append(m2_expr.ExprAff(ir.IRDst,
@@ -1501,13 +1501,13 @@ def idiv(ir, instr, src1):
     do_div = []
     do_div += e
     do_div.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    blk_div = IRBlock(lbl_div.name, [do_div])
+    blk_div = IRBlock(lbl_div.name, [AssignBlock(do_div, instr)])
 
     do_except = []
     do_except.append(m2_expr.ExprAff(exception_flags, m2_expr.ExprInt(
         EXCEPT_DIV_BY_ZERO, exception_flags.size)))
     do_except.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    blk_except = IRBlock(lbl_except.name, [do_except])
+    blk_except = IRBlock(lbl_except.name, [AssignBlock(do_except, instr)])
 
     e = []
     e.append(m2_expr.ExprAff(ir.IRDst,
@@ -1667,12 +1667,12 @@ def stos(ir, instr, size):
     e0 = []
     e0.append(m2_expr.ExprAff(addr_o, addr_p))
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e0 = IRBlock(lbl_df_0.name, [e0])
+    e0 = IRBlock(lbl_df_0.name, [AssignBlock(e0, instr)])
 
     e1 = []
     e1.append(m2_expr.ExprAff(addr_o, addr_m))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e1 = IRBlock(lbl_df_1.name, [e1])
+    e1 = IRBlock(lbl_df_1.name, [AssignBlock(e1, instr)])
 
     e = []
     e.append(m2_expr.ExprAff(ir.ExprMem(addr, size), b))
@@ -1702,12 +1702,12 @@ def lods(ir, instr, size):
     e0 = []
     e0.append(m2_expr.ExprAff(addr_o, addr_p))
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e0 = IRBlock(lbl_df_0.name, [e0])
+    e0 = IRBlock(lbl_df_0.name, [AssignBlock(e0, instr)])
 
     e1 = []
     e1.append(m2_expr.ExprAff(addr_o, addr_m))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e1 = IRBlock(lbl_df_1.name, [e1])
+    e1 = IRBlock(lbl_df_1.name, [AssignBlock(e1, instr)])
 
     e = []
     if instr.mode == 64 and b.size == 32:
@@ -1744,13 +1744,13 @@ def movs(ir, instr, size):
     e0.append(m2_expr.ExprAff(a, a + m2_expr.ExprInt(size / 8, a.size)))
     e0.append(m2_expr.ExprAff(b, b + m2_expr.ExprInt(size / 8, b.size)))
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e0 = IRBlock(lbl_df_0.name, [e0])
+    e0 = IRBlock(lbl_df_0.name, [AssignBlock(e0, instr)])
 
     e1 = []
     e1.append(m2_expr.ExprAff(a, a - m2_expr.ExprInt(size / 8, a.size)))
     e1.append(m2_expr.ExprAff(b, b - m2_expr.ExprInt(size / 8, b.size)))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    e1 = IRBlock(lbl_df_1.name, [e1])
+    e1 = IRBlock(lbl_df_1.name, [AssignBlock(e1, instr)])
 
     e.append(m2_expr.ExprAff(ir.IRDst,
                              m2_expr.ExprCond(df, lbl_df_1, lbl_df_0)))
@@ -2783,8 +2783,8 @@ def bsr_bsf(ir, instr, dst, src, op_name):
     e_src_not_null.append(m2_expr.ExprAff(dst, m2_expr.ExprOp(op_name, src)))
     e_src_not_null.append(aff_dst)
 
-    return e, [IRBlock(lbl_src_null.name, [e_src_null]),
-               IRBlock(lbl_src_not_null.name, [e_src_not_null])]
+    return e, [IRBlock(lbl_src_null.name, [AssignBlock(e_src_null, instr)]),
+               IRBlock(lbl_src_not_null.name, [AssignBlock(e_src_not_null, instr)])]
 
 
 def bsf(ir, instr, dst, src):
@@ -3682,7 +3682,8 @@ def ps_rl_ll(ir, instr, dst, src, op, size):
     e_do = []
     e.append(m2_expr.ExprAff(dst[0:dst.size], m2_expr.ExprCompose(*slices)))
     e_do.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
-    return e, [IRBlock(lbl_do.name, [e_do]), IRBlock(lbl_zero.name, [e_zero])]
+    return e, [IRBlock(lbl_do.name, [AssignBlock(e_do, instr)]),
+               IRBlock(lbl_zero.name, [AssignBlock(e_zero, instr)])]
 
 
 def psrlw(ir, instr, dst, src):
@@ -4598,11 +4599,10 @@ class ir_x86_16(IntermediateRepresentation):
         lbl_skip = m2_expr.ExprId(self.get_next_label(instr), self.IRDst.size)
         lbl_next = m2_expr.ExprId(self.get_next_label(instr), self.IRDst.size)
 
-        for irblock in extra_ir:
-            for ir in irblock.irs:
-                for i, e in enumerate(ir):
-                    src = e.src.replace_expr({lbl_next: lbl_end})
-                    ir[i] = m2_expr.ExprAff(e.dst, src)
+        fix_next_lbl = {lbl_next: lbl_end}
+        new_extra_ir = [irblock.modify_exprs(mod_src=lambda expr: expr.replace_expr(fix_next_lbl))
+                        for irblock in extra_ir]
+
         cond_bloc = []
         cond_bloc.append(m2_expr.ExprAff(c_reg,
                                          c_reg - m2_expr.ExprInt(1,
@@ -4610,14 +4610,14 @@ class ir_x86_16(IntermediateRepresentation):
         cond_bloc.append(m2_expr.ExprAff(self.IRDst, m2_expr.ExprCond(c_cond,
                                                                       lbl_skip,
                                                                       lbl_do)))
-        cond_bloc = IRBlock(lbl_end.name, [cond_bloc])
+        cond_bloc = IRBlock(lbl_end.name, [AssignBlock(cond_bloc, instr)])
         e_do = instr_ir
 
-        c = IRBlock(lbl_do.name, [e_do])
+        c = IRBlock(lbl_do.name, [AssignBlock(e_do, instr)])
         c.except_automod = False
         e_n = [m2_expr.ExprAff(self.IRDst, m2_expr.ExprCond(c_reg, lbl_do,
                                                             lbl_skip))]
-        return e_n, [cond_bloc, c] + extra_ir
+        return e_n, [cond_bloc, c] + new_extra_ir
 
     def expr_fix_regs_for_mode(self, e, mode=64):
         return e.replace_expr(replace_regs[mode])
@@ -4627,8 +4627,9 @@ class ir_x86_16(IntermediateRepresentation):
         src = self.expr_fix_regs_for_mode(e.src, mode)
         return m2_expr.ExprAff(dst, src)
 
-    def irbloc_fix_regs_for_mode(self, irbloc, mode=64):
-        for idx, assignblk in enumerate(irbloc.irs):
+    def irbloc_fix_regs_for_mode(self, irblock, mode=64):
+        irs = []
+        for assignblk in irblock.irs:
             new_assignblk = dict(assignblk)
             for dst, src in assignblk.iteritems():
                 del new_assignblk[dst]
@@ -4643,9 +4644,8 @@ class ir_x86_16(IntermediateRepresentation):
                 dst = self.expr_fix_regs_for_mode(dst, mode)
                 src = self.expr_fix_regs_for_mode(src, mode)
                 new_assignblk[dst] = src
-            irbloc.irs[idx] = AssignBlock(new_assignblk, assignblk.instr)
-        if irbloc.dst is not None:
-            irbloc.dst = self.expr_fix_regs_for_mode(irbloc.dst, mode)
+            irs.append(AssignBlock(new_assignblk, assignblk.instr))
+        return IRBlock(irblock.label, irs)
 
 
 class ir_x86_32(ir_x86_16):
@@ -4677,21 +4677,16 @@ class ir_x86_64(ir_x86_16):
 
     def mod_pc(self, instr, instr_ir, extra_ir):
         # fix RIP for 64 bit
+        pc_fixed = {self.pc: m2_expr.ExprInt(instr.offset + instr.l, 64)}
+
         for i, expr in enumerate(instr_ir):
             dst, src = expr.dst, expr.src
             if dst != self.pc:
-                dst = dst.replace_expr(
-                    {self.pc: m2_expr.ExprInt(instr.offset + instr.l, 64)})
-            src = src.replace_expr(
-                {self.pc: m2_expr.ExprInt(instr.offset + instr.l, 64)})
+                dst = dst.replace_expr(pc_fixed)
+            src = src.replace_expr(pc_fixed)
             instr_ir[i] = m2_expr.ExprAff(dst, src)
-        for irblock in extra_ir:
-            for irs in irblock.irs:
-                for i, expr in enumerate(irs):
-                    dst, src = expr.dst, expr.src
-                    if dst != self.pc:
-                        new_pc = m2_expr.ExprInt(instr.offset + instr.l, 64)
-                        dst = dst.replace_expr({self.pc: new_pc})
-                    src = src.replace_expr(
-                        {self.pc: m2_expr.ExprInt(instr.offset + instr.l, 64)})
-                    irs[i] = m2_expr.ExprAff(dst, src)
+
+        for idx, irblock in enumerate(extra_ir):
+            extra_ir[idx] = irblock.modify_exprs(lambda expr: expr.replace_expr(pc_fixed) \
+                                                 if expr != self.pc else expr,
+                                                 lambda expr: expr.replace_expr(pc_fixed))