ExprCompose: updt internal struct

1 files changed, 69 insertions, 97 deletions

diff --git a/miasm2/expression/expression_helper.py b/miasm2/expression/expression_helper.py
index e9176658..8babba70 100644
--- a/miasm2/expression/expression_helper.py
+++ b/miasm2/expression/expression_helper.py
@@ -34,103 +34,76 @@ def parity(a):
     return cpt
 
 
-def merge_sliceto_slice(args):
-    sources = {}
-    non_slice = {}
-    sources_int = {}
-    for a in args:
-        if isinstance(a[0], m2_expr.ExprInt):
-            # sources_int[a.start] = a
-            # copy ExprInt because we will inplace modify arg just below
-            # /!\ TODO XXX never ever modify inplace args...
-            sources_int[a[1]] = (m2_expr.ExprInt(int(a[0]),
-                                                 a[2] - a[1]),
-                                 a[1],
-                                 a[2])
-        elif isinstance(a[0], m2_expr.ExprSlice):
-            if not a[0].arg in sources:
-                sources[a[0].arg] = []
-            sources[a[0].arg].append(a)
+def merge_sliceto_slice(expr):
+    """
+    Apply basic factorisation on ExprCompose sub compoenents
+    @expr: ExprCompose
+    """
+
+    slices_raw = []
+    other_raw = []
+    integers_raw = []
+    for index, arg in expr.iter_args():
+        if isinstance(arg, m2_expr.ExprInt):
+            integers_raw.append((index, arg))
+        elif isinstance(arg, m2_expr.ExprSlice):
+            slices_raw.append((index, arg))
         else:
-            non_slice[a[1]] = a
-    # find max stop to determine size
-    max_size = None
-    for a in args:
-        if max_size is None or max_size < a[2]:
-            max_size = a[2]
-
-    # first simplify all num slices
-    final_sources = []
-    sorted_s = []
-    for x in sources_int.values():
-        x = list(x)
-        # mask int
-        v = x[0].arg & ((1 << (x[2] - x[1])) - 1)
-        x[0] = m2_expr.ExprInt_from(x[0], v)
-        x = tuple(x)
-        sorted_s.append((x[1], x))
-    sorted_s.sort()
-    while sorted_s:
-        start, v = sorted_s.pop()
-        out = [m2_expr.ExprInt(v[0].arg), v[1], v[2]]
-        size = v[2] - v[1]
-        while sorted_s:
-            if sorted_s[-1][1][2] != start:
+            other_raw.append((index, arg))
+
+    # Find max stop to determine size
+    max_size = sum([arg.size for arg in expr.args])
+
+    integers_merged = []
+    # Merge consecutive integers
+    while integers_raw:
+        index, arg = integers_raw.pop()
+        new_size = arg.size
+        value = int(arg)
+        while integers_raw:
+            prev_index, prev_value = integers_raw[-1]
+            # Check if intergers are consecutive
+            if prev_index + prev_value.size != index:
                 break
-            s_start, s_stop = sorted_s[-1][1][1], sorted_s[-1][1][2]
-            size += s_stop - s_start
-            a = m2_expr.mod_size2uint[size](
-                (int(out[0]) << (out[1] - s_start)) +
-                 int(sorted_s[-1][1][0]))
-            out[0] = m2_expr.ExprInt(a)
-            sorted_s.pop()
-            out[1] = s_start
-        out[0] = m2_expr.ExprInt(int(out[0]), size)
-        final_sources.append((start, tuple(out)))
-
-    final_sources_int = final_sources
-    # check if same sources have corresponding start/stop
-    # is slice AND is sliceto
-    simp_sources = []
-    for args in sources.values():
-        final_sources = []
-        sorted_s = []
-        for x in args:
-            sorted_s.append((x[1], x))
-        sorted_s.sort()
-        while sorted_s:
-            start, v = sorted_s.pop()
-            ee = v[0].arg[v[0].start:v[0].stop]
-            out = ee, v[1], v[2]
-            while sorted_s:
-                if sorted_s[-1][1][2] != start:
-                    break
-                if sorted_s[-1][1][0].stop != out[0].start:
-                    break
-
-                start = sorted_s[-1][1][1]
-                # out[0].start = sorted_s[-1][1][0].start
-                o_e, _, o_stop = out
-                o1, o2 = sorted_s[-1][1][0].start, o_e.stop
-                o_e = o_e.arg[o1:o2]
-                out = o_e, start, o_stop
-                # update _size
-                # out[0]._size = out[0].stop-out[0].start
-                sorted_s.pop()
-            out = out[0], start, out[2]
-
-            final_sources.append((start, out))
+            # Merge integers
+            index = prev_index
+            new_size += prev_value.size
+            value = value << prev_value.size
+            value |= int(prev_value)
+            integers_raw.pop()
+        integers_merged.append((index, m2_expr.ExprInt(value, new_size)))
+
+
+    slices_merged = []
+    # Merge consecutive slices
+    while slices_raw:
+        index, arg = slices_raw.pop()
+        value, slice_start, slice_stop = arg.arg, arg.start, arg.stop
+        while slices_raw:
+            prev_index, prev_value = slices_raw[-1]
+            # Check if slices are consecutive
+            if prev_index + prev_value.size != index:
+                break
+            # Check if slices can ben merged
+            if prev_value.arg != value:
+                break
+            if prev_value.stop != slice_start:
+                break
+            # Merge slices
+            index = prev_index
+            slice_start = prev_value.start
+            slices_raw.pop()
+        slices_merged.append((index, value[slice_start:slice_stop]))
 
-        simp_sources += final_sources
 
-    simp_sources += final_sources_int
+    new_args = slices_merged + integers_merged + other_raw
+    new_args.sort()
+    for i, (index, arg) in enumerate(new_args[:-1]):
+        assert index + arg.size == new_args[i+1][0]
+    ret = [arg[1] for arg in new_args]
 
-    for i, v in non_slice.items():
-        simp_sources.append((i, v))
+    return ret
 
-    simp_sources.sort()
-    simp_sources = [x[1] for x in simp_sources]
-    return simp_sources
 
 
 op_propag_cst = ['+', '*', '^', '&', '|', '>>',
@@ -327,8 +300,8 @@ class Variables_Identifier(object):
             self.find_variables_rec(expr.arg)
 
         elif isinstance(expr, m2_expr.ExprCompose):
-            for a in expr.args:
-                self.find_variables_rec(list(a)[0])
+            for arg in expr.args:
+                self.find_variables_rec(arg)
 
         elif isinstance(expr, m2_expr.ExprSlice):
             self.find_variables_rec(expr.arg)
@@ -646,15 +619,14 @@ def possible_values(expr):
     elif isinstance(expr, m2_expr.ExprCompose):
         # Generate each possibility for sub-argument, associated with the start
         # and stop bit
-        consvals_args = [map(lambda x: (x, arg[1], arg[2]),
-                             possible_values(arg[0]))
+        consvals_args = [map(lambda x: x, possible_values(arg))
                          for arg in expr.args]
         for consvals_possibility in itertools.product(*consvals_args):
             # Merge constraint of each sub-element
-            args_constraint = itertools.chain(*[consval[0].constraints
+            args_constraint = itertools.chain(*[consval.constraints
                                                 for consval in consvals_possibility])
             # Gen the corresponding constraints / ExprCompose
-            args = [consval[0].value for consval in consvals_possibility]
+            args = [consval.value for consval in consvals_possibility]
             consvals.add(
                 ConstrainedValue(frozenset(args_constraint),
                                  m2_expr.ExprCompose(*args)))