Ir: make AssignBlock immutable

1 files changed, 134 insertions, 121 deletions

diff --git a/miasm2/jitter/codegen.py b/miasm2/jitter/codegen.py
index 09a6fecf..8a03c667 100644
--- a/miasm2/jitter/codegen.py
+++ b/miasm2/jitter/codegen.py
@@ -1,5 +1,9 @@
+"""
+Module to generate C code for a given native @block
+"""
+
 import miasm2.expression.expression as m2_expr
-from miasm2.ir.ir import IRBlock
+from miasm2.ir.ir import IRBlock, AssignBlock
 from miasm2.ir.translators import Translator
 from miasm2.core.asmblock import expr_is_label, AsmBlockBad, AsmLabel
 
@@ -131,17 +135,22 @@ class CGen(object):
     def assignblk_to_irbloc(self, instr, assignblk):
         """
         Ensure IRDst is always set in the head @assignblk of the @instr
-        @assignblk: Assignblk instance
         @instr: an instruction instance
+        @assignblk: Assignblk instance
         """
+        new_assignblk = dict(assignblk)
         if self.ir_arch.IRDst not in assignblk:
-            assignblk[self.ir_arch.IRDst] = m2_expr.ExprInt(
-                instr.offset + instr.l,
-                self.ir_arch.IRDst.size)
-
-        return IRBlock(self.ir_arch.get_instr_label(instr), [assignblk])
+            offset = instr.offset + instr.l
+            dst = m2_expr.ExprInt(offset, self.ir_arch.IRDst.size)
+            new_assignblk[self.ir_arch.IRDst] = dst
+        irs = [AssignBlock(new_assignblk)]
+        return IRBlock(self.ir_arch.get_instr_label(instr), irs)
 
     def block2assignblks(self, block):
+        """
+        Return the list of irblocks for a native @block
+        @block: AsmBlock
+        """
         irblocks_list = []
         for instr in block.lines:
             assignblk_head, assignblks_extra = self.ir_arch.instr2ir(instr)
@@ -155,16 +164,13 @@ class CGen(object):
             irblocks_list.append(irblocks)
         return irblocks_list
 
-    def gen_mem_prefetch(self, assignblk, mems_to_prefetch):
-        out = []
-        for expr, prefetcher in sorted(mems_to_prefetch.iteritems()):
-            str_src = self.id_to_c(expr)
-            str_dst = self.id_to_c(prefetcher)
-            out.append('%s = %s;' % (str_dst, str_src))
-        assignblk.C_prefetch = out
-        return out
-
     def add_local_var(self, dst_var, dst_index, expr):
+        """
+        Add local varaible used to store temporay result
+        @dst_var: dictionnary of Expr -> local_var_expr
+        @dst_index : dictionnary of size -> local var count
+        @expr: Expression source
+        """
         size = expr.size
         if size < 8:
             size = 8
@@ -176,17 +182,51 @@ class CGen(object):
         dst_var[expr] = dst
         return dst
 
-    def gen_assignments(self, assignblk, prefetchers):
-        out_var = []
-        out_main = []
-        out_mem = []
-        out_updt = []
+    def get_mem_prefetch(self, assignblk):
+        """
+        Generate temporary variables used to fetch memory used in the @assignblk
+        Return a dictionnary: ExprMem -> temporary variable
+        @assignblk: AssignBlock instance
+        """
+        mem_index = {8: 0, 16: 0, 32: 0, 64: 0, 128:0}
+        mem_var = {}
+
+        # Prefetch memory read
+        for expr in assignblk.get_r(mem_read=True):
+            if not isinstance(expr, m2_expr.ExprMem):
+                continue
+            var_num = mem_index[expr.size]
+            mem_index[expr.size] += 1
+            var = m2_expr.ExprId(
+                "prefetch_%.2d_%.2d" % (expr.size, var_num), expr.size)
+            mem_var[expr] = var
+
+        # Generate memory prefetch
+        return mem_var
+
+    def gen_c_assignments(self, assignblk):
+        """
+        Return C informations used to generate the C code of the @assignblk
+        @assignblk: an AssignBlock instance
+        """
+        c_var = []
+        c_main = []
+        c_mem = []
+        c_updt = []
+        c_prefetch = []
 
-        dst_index = {8: 0, 16: 0, 32: 0, 64: 0}
+        dst_index = {8: 0, 16: 0, 32: 0, 64: 0, 128:0}
         dst_var = {}
 
+        prefetchers = self.get_mem_prefetch(assignblk)
+
+        for expr, prefetcher in sorted(prefetchers.iteritems()):
+            str_src = self.id_to_c(expr)
+            str_dst = self.id_to_c(prefetcher)
+            c_prefetch.append('%s = %s;' % (str_dst, str_src))
+
         for var in prefetchers.itervalues():
-            out_var.append("uint%d_t %s;" % (var.size, var))
+            c_var.append("uint%d_t %s;" % (var.size, var))
 
         for dst, src in sorted(assignblk.iteritems()):
             src = src.replace_expr(prefetchers)
@@ -196,10 +236,10 @@ class CGen(object):
                 new_dst = self.add_local_var(dst_var, dst_index, dst)
                 if dst in self.ir_arch.arch.regs.regs_flt_expr:
                     # Dont mask float affectation
-                    out_main.append(
+                    c_main.append(
                         '%s = (%s);' % (self.id_to_c(new_dst), self.id_to_c(src)))
                 else:
-                    out_main.append(
+                    c_main.append(
                         '%s = (%s)&0x%X;' % (self.id_to_c(new_dst),
                                              self.id_to_c(src),
                                              SIZE_TO_MASK[src.size]))
@@ -207,48 +247,17 @@ class CGen(object):
                 ptr = dst.arg.replace_expr(prefetchers)
                 new_dst = m2_expr.ExprMem(ptr, dst.size)
                 str_dst = self.id_to_c(new_dst).replace('MEM_LOOKUP', 'MEM_WRITE')
-                out_mem.append('%s, %s);' % (str_dst[:-1], self.id_to_c(src)))
+                c_mem.append('%s, %s);' % (str_dst[:-1], self.id_to_c(src)))
             else:
                 raise ValueError("Unknown dst")
 
         for dst, new_dst in dst_var.iteritems():
             if dst is self.ir_arch.IRDst:
                 continue
-            out_updt.append('%s = %s;' % (self.id_to_c(dst), self.id_to_c(new_dst)))
-            out_var.append("uint%d_t %s;" % (new_dst.size, new_dst))
-
-        assignblk.C_var = out_var
-        assignblk.C_main = out_main
-        assignblk.C_mem = out_mem
-        assignblk.C_updt = out_updt
+            c_updt.append('%s = %s;' % (self.id_to_c(dst), self.id_to_c(new_dst)))
+            c_var.append("uint%d_t %s;" % (new_dst.size, new_dst))
 
-    def gen_c_assignblk(self, assignblk):
-        mem_read, mem_write = False, False
-
-        mem_index = {8: 0, 16: 0, 32: 0, 64: 0}
-        mem_var = {}
-        prefetch_index = {8: 0, 16: 0, 32: 0, 64: 0}
-
-        # Prefetch memory read
-        for expr in assignblk.get_r(mem_read=True):
-            if not isinstance(expr, m2_expr.ExprMem):
-                continue
-            mem_read = True
-            var_num = mem_index[expr.size]
-            mem_index[expr.size] += 1
-            var = m2_expr.ExprId(
-                "prefetch_%.2d_%.2d" % (expr.size, var_num), expr.size)
-            mem_var[expr] = var
-
-        # Check if assignblk can write mem
-        mem_write = any(isinstance(expr, m2_expr.ExprMem)
-                        for expr in assignblk.get_w())
-
-        assignblk.mem_write = mem_write
-        assignblk.mem_read = mem_read
-
-        # Generate memory prefetch
-        return mem_var
+        return c_prefetch, c_var, c_main, c_mem, c_updt
 
     def gen_check_memory_exception(self, address):
         dst = self.dst_to_c(address)
@@ -313,12 +322,12 @@ class CGen(object):
 
         return out
 
-    def gen_pre_code(self, attrib):
+    def gen_pre_code(self, instr_attrib):
         out = []
 
-        if attrib.log_mn:
-            out.append('printf("%.8X %s\\n");' % (attrib.instr.offset,
-                                                  attrib.instr))
+        if instr_attrib.log_mn:
+            out.append('printf("%.8X %s\\n");' % (instr_attrib.instr.offset,
+                                                  instr_attrib.instr))
         return out
 
     def gen_post_code(self, attrib):
@@ -357,7 +366,7 @@ class CGen(object):
         """
         Generate code for possible @dst2index.
 
-        @attrib: an Attributs instance
+        @attrib: an Attributes instance
         @instr_offsets: list of instructions offsets
         @dst2index: link from destination to index
         """
@@ -377,54 +386,57 @@ class CGen(object):
                     continue
 
             out.append('\tcase %d:' % index)
+
             out += self.gen_goto_code(attrib, instr_offsets, dst)
             out.append('\t\tbreak;')
         out.append('};')
         return out
 
-    def gen_c_code(self, assignblk, c_dst):
+    def gen_c_code(self, attrib, c_dst, c_assignmnts):
         """
-        Generate the C code for @assignblk.
-        @assignblk: an Assignblk instance
+        Generate the C code for assignblk.
+        @attrib: Attributes instance
         @c_dst: irdst C code
         """
+
+        c_prefetch, c_var, c_main, c_mem, c_updt = c_assignmnts
         out = []
         out.append("{")
         out.append("// var")
-        out += assignblk.C_var
+        out += c_var
         out.append("// Prefetch")
-        out += assignblk.C_prefetch
+        out += c_prefetch
         out.append("// Dst")
         out += c_dst
         out.append("// Main")
-        out += assignblk.C_main
+        out += c_main
 
         out.append("// Check op/mem exceptions")
 
         # Check memory access if assignblk has memory read
-        if assignblk.C_prefetch:
-            out += self.gen_check_memory_exception(assignblk.instr_addr)
+        if c_prefetch:
+            out += self.gen_check_memory_exception(attrib.instr.offset)
 
         # Check if operator raised exception flags
-        if assignblk.op_set_exception:
-            out += self.gen_check_cpu_exception(assignblk.instr_addr)
+        if attrib.op_set_exception:
+            out += self.gen_check_cpu_exception(attrib.instr.offset)
 
         out.append("// Mem updt")
-        out += assignblk.C_mem
+        out += c_mem
 
         out.append("// Check exception Mem write")
         # Check memory write exceptions
-        if assignblk.mem_write:
-            out += self.gen_check_memory_exception(assignblk.instr_addr)
+        if attrib.mem_write:
+            out += self.gen_check_memory_exception(attrib.instr.offset)
 
         out.append("// Updt")
-        out += assignblk.C_updt
+        out += c_updt
 
         out.append("// Checks exception")
 
         # Check post assignblk exception flags
-        if assignblk.set_exception:
-            out += self.gen_check_cpu_exception(assignblk.instr_addr)
+        if attrib.set_exception:
+            out += self.gen_check_cpu_exception(attrib.instr.offset)
 
         out.append("}")
 
@@ -436,29 +448,27 @@ class CGen(object):
         return any(operator.startswith(except_op)
                    for except_op in self.IMPLICIT_EXCEPTION_OP)
 
-    def get_caracteristics(self, irblock):
+    def get_caracteristics(self, assignblk, attrib):
         """
-        Get the carateristics of each assignblk in the @irblock
-        @irblock: an irbloc instance
+        Set the carateristics in @attrib according to the @assignblk
+        @assignblk: an AssignBlock instance
+        @attrib: an Attributes instance
         """
 
-        for assignblk in irblock.irs:
-            assignblk.mem_read, assignblk.mem_write = False, False
-            assignblk.op_set_exception = False
-            # Check explicit exception raising
-            assignblk.set_exception = self.ir_arch.arch.regs.exception_flags in assignblk
-
-            element_read = assignblk.get_r(mem_read=True)
-            # Check implicit exception raising
-            assignblk.op_set_exception = any(self.is_exception_operator(operator)
-                                             for elem in assignblk.values()
-                                             for operator in m2_expr.get_expr_ops(elem))
-            # Check mem read
-            assignblk.mem_read = any(isinstance(expr, m2_expr.ExprMem)
-                                     for expr in element_read)
-            # Check mem write
-            assignblk.mem_write = any(isinstance(dst, m2_expr.ExprMem)
-                                      for dst in assignblk)
+        # Check explicit exception raising
+        attrib.set_exception = self.ir_arch.arch.regs.exception_flags in assignblk
+
+        element_read = assignblk.get_r(mem_read=True)
+        # Check implicit exception raising
+        attrib.op_set_exception = any(self.is_exception_operator(operator)
+                                      for elem in assignblk.values()
+                                      for operator in m2_expr.get_expr_ops(elem))
+        # Check mem read
+        attrib.mem_read = any(isinstance(expr, m2_expr.ExprMem)
+                              for expr in element_read)
+        # Check mem write
+        attrib.mem_write = any(isinstance(dst, m2_expr.ExprMem)
+                               for dst in assignblk)
 
     def get_attributes(self, instr, irblocks, log_mn=False, log_regs=False):
         """
@@ -469,17 +479,24 @@ class CGen(object):
         @log_regs: generate code to log registers states
         """
 
-        attrib = Attributes(log_mn, log_regs)
+        instr_attrib = Attributes()
+        instr_attrib.instr = instr
+        irblocks_attributes = []
 
         for irblock in irblocks:
+            attributes = []
+            irblocks_attributes.append(attributes)
             for assignblk in irblock.irs:
-                self.get_caracteristics(irblock)
-                attrib.mem_read |= assignblk.mem_read
-                attrib.mem_write |= assignblk.mem_write
-                attrib.set_exception |= assignblk.set_exception
-                attrib.op_set_exception |= assignblk.op_set_exception
-        attrib.instr = instr
-        return attrib
+                attrib = Attributes(log_mn, log_regs)
+                attributes.append(attrib)
+                self.get_caracteristics(assignblk, attrib)
+                attrib.instr = instr
+                instr_attrib.mem_read |= attrib.mem_read
+                instr_attrib.mem_write |= attrib.mem_write
+                instr_attrib.op_set_exception |= attrib.op_set_exception
+                instr_attrib.set_exception |= attrib.set_exception
+
+        return instr_attrib, irblocks_attributes
 
     def gen_bad_block(self):
         """
@@ -503,12 +520,11 @@ class CGen(object):
         lbl_start = self.ir_arch.symbol_pool.getby_offset_create(instr_offsets[0])
         return (self.CODE_INIT % self.label_to_jitlabel(lbl_start)).split("\n"), instr_offsets
 
-    def gen_irblock(self, attrib, instr_offsets, instr, irblock):
+    def gen_irblock(self, instr_attrib, attributes, instr_offsets, irblock):
         """
         Generate the C code for an @irblock
-        @instr: the current instruction to translate
         @irblock: an irbloc instance
-        @attrib: an Attributs instance
+        @attributes: an Attributes instance list
         """
 
         out = []
@@ -518,17 +534,14 @@ class CGen(object):
                 c_dst, dst2index = self.gen_assignblk_dst(irblock.dst)
             else:
                 c_dst = []
-            assignblk.instr_addr = instr.offset
-            prefetchers = self.gen_c_assignblk(assignblk)
-            self.gen_mem_prefetch(assignblk, prefetchers)
-            self.gen_assignments(assignblk, prefetchers)
 
-            out += self.gen_c_code(assignblk, c_dst)
+            c_assignmnts = self.gen_c_assignments(assignblk)
+            out += self.gen_c_code(attributes[index], c_dst, c_assignmnts)
 
         if dst2index:
             out.append("// Set irdst")
             # Gen goto on irdst set
-            out += self.gen_dst_goto(attrib, instr_offsets, dst2index)
+            out += self.gen_dst_goto(instr_attrib, instr_offsets, dst2index)
 
         return out
 
@@ -556,7 +569,7 @@ class CGen(object):
         out, instr_offsets = self.gen_init(block)
 
         for instr, irblocks in zip(block.lines, irblocks_list):
-            attrib = self.get_attributes(instr, irblocks, log_mn, log_regs)
+            instr_attrib, irblocks_attributes = self.get_attributes(instr, irblocks, log_mn, log_regs)
 
             for index, irblock in enumerate(irblocks):
                 self.ir_arch.irbloc_fix_regs_for_mode(
@@ -569,8 +582,8 @@ class CGen(object):
                     out.append("%-40s // %.16X %s" %
                                (self.label_to_jitlabel(irblock.label) + ":", instr.offset, instr))
                 if index == 0:
-                    out += self.gen_pre_code(attrib)
-                out += self.gen_irblock(attrib, instr_offsets, instr, irblock)
+                    out += self.gen_pre_code(instr_attrib)
+                out += self.gen_irblock(instr_attrib, irblocks_attributes[index], instr_offsets, irblock)
 
         out += self.gen_finalize(block)
         return ['\t' + line for line in out]