Merge pull request #1011 from icecr4ck/add_simp_compose

Add compose_and_mask simplification pass

3 files changed, 37 insertions, 2 deletions

diff --git a/miasm/expression/simplifications.py b/miasm/expression/simplifications.py
index 585a9c6b..8f63ab91 100644
--- a/miasm/expression/simplifications.py
+++ b/miasm/expression/simplifications.py
@@ -58,7 +58,7 @@ class ExpressionSimplifier(object):
             simplifications_common.simp_test_signext_inf,
             simplifications_common.simp_test_zeroext_inf,
             simplifications_common.simp_cond_inf_eq_unsigned_zero,
-
+            simplifications_common.simp_compose_and_mask,
         ],
 
         m2_expr.ExprSlice: [
diff --git a/miasm/expression/simplifications_common.py b/miasm/expression/simplifications_common.py
index cda9c5e2..69d56997 100644
--- a/miasm/expression/simplifications_common.py
+++ b/miasm/expression/simplifications_common.py
@@ -607,7 +607,6 @@ def simp_compose(e_s, expr):
         return ExprCond(cond, arg1, arg2)
     return ExprCompose(*args)
 
-
 def simp_cond(_, expr):
     """
     Common simplifications on ExprCond.
@@ -1554,3 +1553,34 @@ def simp_add_multiple(_, expr):
     if len(out) == 1:
         return out[0]
     return ExprOp('+', *out)
+
+def simp_compose_and_mask(_, expr):
+    """
+    {X 0 8, Y 8 32} & 0xFF => zeroExt(X)
+    {X 0 8, Y 8 16, Z 16 32} & 0xFFFF => {X 0 8, Y 8 16, 0x0 16 32}
+    {X 0 8, 0x123456 8 32} & 0xFFFFFF => {X 0 8, 0x1234 8 24, 0x0 24 32}
+    """
+    if not expr.is_op('&'):
+        return expr
+    # handle the case where arg2 = arg1.mask
+    if len(expr.args) != 2:
+        return expr
+    arg1, arg2 = expr.args
+    if not arg1.is_compose():
+        return expr
+    if not arg2.is_int():
+        return expr
+    int2 = int(arg2)
+    if (int2 + 1) & int2 != 0:
+        return expr
+    mask_size = int2.bit_length() + 7 // 8
+    out = []
+    for offset, arg in arg1.iter_args():
+        if offset == mask_size:
+            return ExprCompose(*out).zeroExtend(expr.size)
+        elif mask_size > offset and mask_size < offset+arg.size and arg.is_int():
+            out.append(ExprSlice(arg, 0, mask_size-offset))
+            return ExprCompose(*out).zeroExtend(expr.size)
+        else:
+            out.append(arg)
+    return expr
diff --git a/test/expression/simplifications.py b/test/expression/simplifications.py
index 1a22c43d..e0b666da 100644
--- a/test/expression/simplifications.py
+++ b/test/expression/simplifications.py
@@ -192,6 +192,11 @@ to_test = [(ExprInt(1, 32) - ExprInt(1, 32), ExprInt(0, 32)),
             ExprOp('&', a, ExprInt(0x0FFFFFFF, 32))),
            (ExprOp('<<', ExprOp('>>', a, ExprInt(0x4, 32)), ExprInt(0x4, 32)),
             ExprOp('&', a, ExprInt(0xFFFFFFF0, 32))),
+
+           (ExprCompose(ExprId("a", 8), ExprId("b", 24)) & ExprInt(0xFF, 32), ExprCompose(ExprId("a", 8), ExprInt(0x0, 24))),
+           (ExprCompose(ExprId("a", 8), ExprInt(0x12, 8), ExprId("b", 16)) & ExprInt(0xFFFF, 32), ExprCompose(ExprId("a", 8), ExprInt(0x12, 24))),
+           (ExprCompose(ExprId("a", 8), ExprInt(0x1234, 16), ExprId("b", 8)) & ExprInt(0xFFFF, 32), ExprCompose(ExprId("a", 8), ExprInt(0x34, 24))),
+
            (a[:32], a),
            (a[:8][:8], a[:8]),
            (a[:16][:8], a[:8]),