AssignBlock

7 files changed, 86 insertions, 162 deletions

diff --git a/miasm2/arch/aarch64/ira.py b/miasm2/arch/aarch64/ira.py
index cf44f42c..3f610360 100644
--- a/miasm2/arch/aarch64/ira.py
+++ b/miasm2/arch/aarch64/ira.py
@@ -2,7 +2,7 @@
 #-*- coding:utf-8 -*-
 
 from miasm2.expression.expression import *
-from miasm2.ir.ir import ir, irbloc
+from miasm2.ir.ir import ir, irbloc, AssignBlock
 from miasm2.ir.analysis import ira
 from miasm2.arch.aarch64.sem import ir_aarch64l, ir_aarch64b
 from miasm2.arch.aarch64.regs import *
@@ -35,23 +35,14 @@ class ir_a_aarch64l(ir_a_aarch64l_base):
         b.rw[-1][1].add(self.arch.regs.of)
         b.rw[-1][1].add(self.arch.regs.cf)
 
-    def call_effects(self, ad):
-        irs = [[ExprAff(self.ret_reg, ExprOp('call_func_ret', ad, self.sp)),
-                ExprAff(self.sp, ExprOp('call_func_stack', ad, self.sp)),
-                ]]
-        return irs
-
     def post_add_bloc(self, bloc, ir_blocs):
         ir.post_add_bloc(self, bloc, ir_blocs)
         for irb in ir_blocs:
             pc_val = None
             lr_val = None
-            for exprs in irb.irs:
-                for e in exprs:
-                    if e.dst == PC:
-                        pc_val = e.src
-                    if e.dst == LR:
-                        lr_val = e.src
+            for assignblk in irb.irs:
+                pc_val = assignblk.get(PC, pc_val)
+                lr_val = assignblk.get(LR, lr_val)
             if pc_val is None or lr_val is None:
                 continue
             if not isinstance(lr_val, ExprInt):
@@ -60,10 +51,13 @@ class ir_a_aarch64l(ir_a_aarch64l_base):
             l = bloc.lines[-1]
             if lr_val.arg != l.offset + l.l:
                 continue
+
+            # CALL
             lbl = bloc.get_next()
             new_lbl = self.gen_label()
             irs = self.call_effects(pc_val)
-            irs.append([ExprAff(self.IRDst, ExprId(lbl, size=self.pc.size))])
+            irs.append(AssignBlock([ExprAff(self.IRDst,
+                                            ExprId(lbl, size=self.pc.size))]))
             nbloc = irbloc(new_lbl, irs)
             nbloc.lines = [l] * len(irs)
             self.blocs[new_lbl] = nbloc
diff --git a/miasm2/arch/aarch64/sem.py b/miasm2/arch/aarch64/sem.py
index 9c896095..b198bc43 100644
--- a/miasm2/arch/aarch64/sem.py
+++ b/miasm2/arch/aarch64/sem.py
@@ -1,5 +1,5 @@
 from miasm2.expression import expression as m2_expr
-from miasm2.ir.ir import ir, irbloc
+from miasm2.ir.ir import ir, irbloc, AssignBlock
 from miasm2.arch.aarch64.arch import mn_aarch64, conds_expr, replace_regs
 from miasm2.arch.aarch64.regs import *
 from miasm2.core.sembuilder import SemBuilder
@@ -778,19 +778,21 @@ class ir_aarch64l(ir):
         return m2_expr.ExprAff(dst, src)
 
     def irbloc_fix_regs_for_mode(self, irbloc, mode=64):
-        for irs in irbloc.irs:
-            for i, e in enumerate(irs):
-                """
-                special case for 64 bits:
-                if destination is a 32 bit reg, zero extend the 64 bit reg
-                """
-                if (isinstance(e.dst, m2_expr.ExprId) and
-                        e.dst.size == 32 and
-                        e.dst in replace_regs):
-                    src = self.expr_fix_regs_for_mode(e.src)
-                    dst = replace_regs[e.dst].arg
-                    e = m2_expr.ExprAff(dst, src.zeroExtend(64))
-                irs[i] = self.expr_fix_regs_for_mode(e)
+        for assignblk in irbloc.irs:
+            for dst, src in assignblk.items():
+                del(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):
+                    src = src.zeroExtend(64)
+                    dst = replace_regs[dst].arg
+
+                dst = self.expr_fix_regs_for_mode(dst)
+                src = self.expr_fix_regs_for_mode(src)
+                assignblk[dst] = src
         irbloc.dst = self.expr_fix_regs_for_mode(irbloc.dst)
 
     def mod_pc(self, instr, instr_ir, extra_ir):
diff --git a/miasm2/arch/arm/ira.py b/miasm2/arch/arm/ira.py
index b918a2e6..b279fab5 100644
--- a/miasm2/arch/arm/ira.py
+++ b/miasm2/arch/arm/ira.py
@@ -2,7 +2,7 @@
 #-*- coding:utf-8 -*-
 
 from miasm2.expression.expression import *
-from miasm2.ir.ir import ir, irbloc
+from miasm2.ir.ir import ir, irbloc, AssignBlock
 from miasm2.ir.analysis import ira
 from miasm2.arch.arm.sem import ir_arml, ir_armtl, ir_armb, ir_armtb
 from miasm2.arch.arm.regs import *
@@ -33,26 +33,14 @@ class ir_a_arml(ir_a_arml_base):
         b.rw[-1][1].add(self.arch.regs.of)
         b.rw[-1][1].add(self.arch.regs.cf)
 
-    def call_effects(self, ad):
-        irs = [[ExprAff(self.ret_reg, ExprOp('call_func_ret', ad, self.sp)),
-                ExprAff(self.sp, ExprOp('call_func_stack', ad, self.sp)),
-                ]]
-        return irs
-
     def post_add_bloc(self, bloc, ir_blocs):
         ir.post_add_bloc(self, bloc, ir_blocs)
-        # flow_graph = DiGraph()
         for irb in ir_blocs:
-            # print 'X'*40
-            # print irb
             pc_val = None
             lr_val = None
-            for exprs in irb.irs:
-                for e in exprs:
-                    if e.dst == PC:
-                        pc_val = e.src
-                    if e.dst == LR:
-                        lr_val = e.src
+            for assignblk in irb.irs:
+                pc_val = assignblk.get(PC, pc_val)
+                lr_val = assignblk.get(LR, lr_val)
             if pc_val is None or lr_val is None:
                 continue
             if not isinstance(lr_val, ExprInt):
@@ -61,51 +49,18 @@ class ir_a_arml(ir_a_arml_base):
             l = bloc.lines[-1]
             if lr_val.arg != l.offset + l.l:
                 continue
-            # print 'IS CALL!'
+
+            # CALL
             lbl = bloc.get_next()
             new_lbl = self.gen_label()
             irs = self.call_effects(pc_val)
-            irs.append([ExprAff(self.IRDst, ExprId(lbl, size=self.pc.size))])
+            irs.append(AssignBlock([ExprAff(self.IRDst,
+                                            ExprId(lbl, size=self.pc.size))]))
             nbloc = irbloc(new_lbl, irs)
-            nbloc.lines = [l]*len(irs)
+            nbloc.lines = [l] * len(irs)
             self.blocs[new_lbl] = nbloc
             irb.dst = ExprId(new_lbl, size=self.pc.size)
 
-        """
-        if not bloc.lines:
-            return
-        l = bloc.lines[-1]
-        sub_call_dst = None
-        if not l.is_subcall():
-            return
-        sub_call_dst = l.args[0]
-        if self.ExprIsLabel(sub_call_dst):
-            sub_call_dst = sub_call_dst.name
-        for b in ir_blocs:
-            l = b.lines[-1]
-            sub_call_dst_b = None
-            sub_call_dst_b = l.args[0]
-            #if self.ExprIsLabel(sub_call_dst_b):
-            #    sub_call_dst_b = sub_call_dst.name
-            #if str(b.dst) == str(sub_call_dst_b):
-            #    pass
-            if not l.is_subcall():
-                continue
-            if b.dst != sub_call_dst_b:
-                continue
-            sub_call_dst_b = l.args[0]
-            if self.ExprIsLabel(sub_call_dst_b):
-                sub_call_dst_b = sub_call_dst.name
-            lbl = bloc.get_next()
-            new_lbl = self.gen_label()
-            irs = self.call_effects(l.args[0])
-            nbloc = irbloc(new_lbl, ExprId(lbl, size=self.pc.size), irs)
-            nbloc.lines = [l]
-            self.blocs[new_lbl] = nbloc
-            b.dst = ExprId(new_lbl, size=self.pc.size)
-        return
-        """
-
     def get_out_regs(self, b):
         return set([self.ret_reg, self.sp])
 
diff --git a/miasm2/arch/mips32/ira.py b/miasm2/arch/mips32/ira.py
index f88172fb..de508e41 100644
--- a/miasm2/arch/mips32/ira.py
+++ b/miasm2/arch/mips32/ira.py
@@ -2,7 +2,7 @@
 #-*- coding:utf-8 -*-
 
 from miasm2.expression.expression import *
-from miasm2.ir.ir import ir, irbloc
+from miasm2.ir.ir import ir, irbloc, AssignBlock
 from miasm2.ir.analysis import ira
 from miasm2.arch.mips32.sem import ir_mips32l, ir_mips32b
 from miasm2.arch.mips32.regs import *
@@ -18,26 +18,15 @@ class ir_a_mips32l(ir_mips32l, ira):
     def set_dead_regs(self, b):
         pass
 
-    def call_effects(self, ad):
-        irs = [[ExprAff(self.ret_reg, ExprOp('call_func_ret', ad, self.sp)),
-                ExprAff(self.sp, ExprOp('call_func_stack', ad, self.sp)),
-                ]]
-        return irs
-
     def post_add_bloc(self, bloc, ir_blocs):
         ir.post_add_bloc(self, bloc, ir_blocs)
         for irb in ir_blocs:
-            # print 'X'*40
-            # print irb
             pc_val = None
             lr_val = None
-            for exprs in irb.irs:
-                for e in exprs:
-                    if e.dst == PC:
-                        pc_val = e.src
-                    if e.dst == RA:
-                        lr_val = e.src
-            #print "XXX", pc_val, lr_val
+            for assignblk in irb.irs:
+                pc_val = assignblk.get(PC, pc_val)
+                lr_val = assignblk.get(RA, lr_val)
+
             if pc_val is None or lr_val is None:
                 continue
             if not expr_is_int_or_label(lr_val):
@@ -46,18 +35,17 @@ class ir_a_mips32l(ir_mips32l, ira):
                 lr_val = ExprInt32(lr_val.name.offset)
 
             l = bloc.lines[-2]
-            #print 'TEST', l, hex(lr_val.arg), hex(l.offset + 8)
-            #print lr_val.arg, hex(l.offset + l.l)
             if lr_val.arg != l.offset + 8:
                 raise ValueError("Wrong arg")
 
-            # print 'IS CALL!'
+            # CALL
             lbl = bloc.get_next()
             new_lbl = self.gen_label()
             irs = self.call_effects(pc_val)
-            irs.append([ExprAff(self.IRDst, ExprId(lbl, size=self.pc.size))])
+            irs.append(AssignBlock([ExprAff(self.IRDst,
+                                            ExprId(lbl, size=self.pc.size))]))
             nbloc = irbloc(new_lbl, irs)
-            nbloc.lines = [l]
+            nbloc.lines = [l] * len(irs)
             self.blocs[new_lbl] = nbloc
             irb.dst = ExprId(new_lbl, size=self.pc.size)
 
diff --git a/miasm2/arch/msp430/ira.py b/miasm2/arch/msp430/ira.py
index 26a53a1e..f9da81a7 100644
--- a/miasm2/arch/msp430/ira.py
+++ b/miasm2/arch/msp430/ira.py
@@ -2,7 +2,7 @@
 #-*- coding:utf-8 -*-
 
 from miasm2.expression.expression import *
-from miasm2.ir.ir import ir, irbloc
+from miasm2.ir.ir import ir, irbloc, AssignBlock
 from miasm2.ir.analysis import ira
 from miasm2.arch.msp430.sem import ir_msp430
 from miasm2.arch.msp430.regs import *
@@ -35,39 +35,27 @@ class ir_a_msp430(ir_a_msp430_base):
         b.rw[-1][1].add(self.arch.regs.cpuoff)
         b.rw[-1][1].add(self.arch.regs.gie)
 
-    def call_effects(self, ad):
-        irs = [[ExprAff(self.ret_reg, ExprOp('call_func_ret', ad, self.sp)),
-                ExprAff(self.sp, ExprOp('call_func_stack', ad, self.sp)),
-                ]]
-        return irs
-
     def post_add_bloc(self, bloc, ir_blocs):
         ir.post_add_bloc(self, bloc, ir_blocs)
-        # flow_graph = DiGraph()
-
         l = bloc.lines[-1]
         if not l.is_subcall():
             return
 
         for irb in ir_blocs:
-            # print 'X'*40
-            # print irb
             pc_val = None
-            for exprs in irb.irs:
-                for e in exprs:
-                    if e.dst == PC:
-                        pc_val = e.src
+            for assignblk in irb.irs:
+                pc_val = assignblk.get(PC, pc_val)
             if pc_val is None:
                 continue
 
             l = bloc.lines[-1]
-            # print str(l), 'IS CALL!'
             lbl = bloc.get_next()
             new_lbl = self.gen_label()
             irs = self.call_effects(pc_val)
-            irs.append([ExprAff(self.IRDst, ExprId(lbl, size=self.pc.size))])
+            irs.append(AssignBlock([ExprAff(self.IRDst,
+                                            ExprId(lbl, size=self.pc.size))]))
             nbloc = irbloc(new_lbl, irs)
-            nbloc.lines = [l]
+            nbloc.lines = [l] * len(irs)
             self.blocs[new_lbl] = nbloc
             irb.dst = ExprId(new_lbl, size=self.pc.size)
 
diff --git a/miasm2/arch/x86/ira.py b/miasm2/arch/x86/ira.py
index b7a1f19f..d496f380 100644
--- a/miasm2/arch/x86/ira.py
+++ b/miasm2/arch/x86/ira.py
@@ -4,7 +4,7 @@
 from miasm2.expression.expression import ExprAff, ExprOp, ExprId
 from miasm2.core.graph import DiGraph
 from miasm2.core.asmbloc import expr_is_label
-from miasm2.ir.ir import ir, irbloc
+from miasm2.ir.ir import ir, irbloc, AssignBlock
 from miasm2.ir.analysis import ira
 from miasm2.arch.x86.sem import ir_x86_16, ir_x86_32, ir_x86_64
 
@@ -32,12 +32,6 @@ class ir_a_x86_16(ir_x86_16, ira):
         for b in leaves:
             self.set_dead_regs(b)
 
-    def call_effects(self, ad):
-        irs = [[ExprAff(self.ret_reg, ExprOp('call_func_ret', ad, self.sp)),
-                ExprAff(self.sp, ExprOp('call_func_stack', ad, self.sp)),
-                ]]
-        return irs
-
     def post_add_bloc(self, bloc, ir_blocs):
         ir.post_add_bloc(self, bloc, ir_blocs)
         if not bloc.lines:
@@ -49,8 +43,8 @@ class ir_a_x86_16(ir_x86_16, ira):
         sub_call_dst = l.args[0]
         if expr_is_label(sub_call_dst):
             sub_call_dst = sub_call_dst.name
-        for b in ir_blocs:
-            l = b.lines[-1]
+        for irb in ir_blocs:
+            l = irb.lines[-1]
             sub_call_dst = None
             if not l.is_subcall():
                 continue
@@ -60,12 +54,13 @@ class ir_a_x86_16(ir_x86_16, ira):
             lbl = bloc.get_next()
             new_lbl = self.gen_label()
             irs = self.call_effects(l.args[0])
-            irs.append([ExprAff(self.IRDst, ExprId(lbl, size=self.pc.size))])
+            irs.append(AssignBlock([ExprAff(self.IRDst,
+                                            ExprId(lbl, size=self.pc.size))]))
 
             nbloc = irbloc(new_lbl, irs)
-            nbloc.lines = [l]
+            nbloc.lines = [l] * len(irs)
             self.blocs[new_lbl] = nbloc
-            b.dst = ExprId(new_lbl, size=self.pc.size)
+            irb.dst = ExprId(new_lbl, size=self.pc.size)
         return
 
 
@@ -98,15 +93,16 @@ class ir_a_x86_64(ir_x86_64, ir_a_x86_16):
         self.ret_reg = self.arch.regs.RAX
 
     def call_effects(self, ad):
-        irs = [[ExprAff(self.ret_reg, ExprOp('call_func_ret', ad, self.sp,
-                                             self.arch.regs.RCX,
-                                             self.arch.regs.RDX,
-                                             self.arch.regs.R8,
-                                             self.arch.regs.R9,
-                                             )),
-                ExprAff(self.sp, ExprOp('call_func_stack', ad, self.sp)),
-                ]]
-        return irs
+        return [AssignBlock([ExprAff(self.ret_reg, ExprOp('call_func_ret', ad,
+                                                          self.sp,
+                                                          self.arch.regs.RCX,
+                                                          self.arch.regs.RDX,
+                                                          self.arch.regs.R8,
+                                                          self.arch.regs.R9,
+                                                          )),
+                             ExprAff(self.sp, ExprOp('call_func_stack',
+                                                     ad, self.sp)),
+                ])]
 
     def sizeof_char(self):
         return 8
diff --git a/miasm2/arch/x86/sem.py b/miasm2/arch/x86/sem.py
index 18410bf9..c14d6f97 100644
--- a/miasm2/arch/x86/sem.py
+++ b/miasm2/arch/x86/sem.py
@@ -3620,11 +3620,11 @@ def ps_rl_ll(ir, instr, a, b, op, size):
               m2_expr.ExprAff(ir.IRDst, lbl_next)]
 
     e_do = []
+    slices = []
     for i in xrange(0, a.size, size):
-        e.append(m2_expr.ExprAff(a[i:i + size], m2_expr.ExprOp(op,
-                                                               a[i:i + size],
-                                                               count[:size])))
-
+        slices.append((m2_expr.ExprOp(op,a[i:i + size], count[:size]),
+                       i, i + size))
+    e.append(m2_expr.ExprAff(a[0:a.size], m2_expr.ExprCompose(slices)))
     e_do.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
     return e, [irbloc(lbl_do.name, [e_do]), irbloc(lbl_zero.name, [e_zero])]
 
@@ -3745,12 +3745,13 @@ def pcmpeqd(ir, instr, a, b):
 
 def punpck(ir, instr, a, b, size, off):
     e = []
+    slices = []
     for i in xrange(a.size / (2 * size)):
         src1 = a[size * i + off: size * i + off + size]
         src2 = b[size * i + off: size * i + off + size]
-        e.append(m2_expr.ExprAff(a[size * 2 * i: size * 2 * i + size], src1))
-        e.append(
-            m2_expr.ExprAff(a[size * (2 * i + 1): size * (2 * i + 1) + size], src2))
+        slices.append((src1, size * 2 * i, size * 2 * i + size))
+        slices.append((src2, size * (2 * i + 1), size * (2 * i + 1) + size))
+    e.append(m2_expr.ExprAff(a, m2_expr.ExprCompose(slices)))
     return e, []
 
 
@@ -4553,20 +4554,20 @@ class ir_x86_16(ir):
         return m2_expr.ExprAff(dst, src)
 
     def irbloc_fix_regs_for_mode(self, irbloc, mode=64):
-        for irs in irbloc.irs:
-            for i, e in enumerate(irs):
-                """
-                special case for 64 bits:
-                if destination is a 32 bit reg, zero extend the 64 bit reg
-                """
+        for assignblk in irbloc.irs:
+            for dst, src in assignblk.items():
+                del(assignblk[dst])
+                # Special case for 64 bits:
+                # If destination is a 32 bit reg, zero extend the 64 bit reg
                 if mode == 64:
-                    if (isinstance(e.dst, m2_expr.ExprId) and
-                            e.dst.size == 32 and
-                            e.dst in replace_regs[64]):
-                        src = self.expr_fix_regs_for_mode(e.src, mode)
-                        dst = replace_regs[64][e.dst].arg
-                        e = m2_expr.ExprAff(dst, src.zeroExtend(64))
-                irs[i] = self.expr_fix_regs_for_mode(e, mode)
+                    if (isinstance(dst, m2_expr.ExprId) and
+                            dst.size == 32 and
+                            dst in replace_regs[64]):
+                        src = src.zeroExtend(64)
+                        dst = replace_regs[64][dst].arg
+                dst = self.expr_fix_regs_for_mode(dst, mode)
+                src = self.expr_fix_regs_for_mode(src, mode)
+                assignblk[dst] = src
         irbloc.dst = self.expr_fix_regs_for_mode(irbloc.dst, mode)