Add check to not slice an ExprId

2 files changed, 10 insertions, 9 deletions

diff --git a/miasm/expression/simplifications_common.py b/miasm/expression/simplifications_common.py
index 444308f0..f00733a0 100644
--- a/miasm/expression/simplifications_common.py
+++ b/miasm/expression/simplifications_common.py
@@ -5,7 +5,7 @@
 from future.utils import viewitems
 
 from miasm.expression.modint import mod_size2int, mod_size2uint
-from miasm.expression.expression import ExprInt, ExprSlice, ExprMem, \
+from miasm.expression.expression import ExprId, ExprInt, ExprSlice, ExprMem, \
     ExprCond, ExprOp, ExprCompose, TOK_INF_SIGNED, TOK_INF_UNSIGNED, \
     TOK_INF_EQUAL_SIGNED, TOK_INF_EQUAL_UNSIGNED, TOK_EQUAL
 from miasm.expression.expression_helper import parity, op_propag_cst, \
@@ -1556,9 +1556,9 @@ def simp_add_multiple(_, expr):
 
 def simp_compose_and_mask(_, expr):
     """
-    {X 0 8, Y 8 16} & 0xFF => X
-    {X 0 32} & 0xFFFF => X[0:16]
-    {X 0 8, Y 8 24, Z 24 32} & 0xFFFFFF => X|Y
+    {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
@@ -1570,8 +1570,10 @@ def simp_compose_and_mask(_, expr):
         return expr
     if not arg2.is_int():
         return expr
-    int2 = arg2.arg
+    int2 = int(arg2)
     if (int2 + 1) & int2 != 0:
         return expr
-    mask_size = int2.arg.bit_length() + 7 // 8
+    mask_size = int2.bit_length() + 7 // 8
+    if not mask_size in [arg[0] for arg in ExprCompose(ExprId("a", 8), ExprInt(0x1234, 16), ExprId("b", 8)).iter_args()]:
+        return expr
     return ExprSlice(arg1, 0, mask_size).zeroExtend(expr.size)
diff --git a/test/expression/simplifications.py b/test/expression/simplifications.py
index b1e590ec..ab235543 100644
--- a/test/expression/simplifications.py
+++ b/test/expression/simplifications.py
@@ -193,10 +193,9 @@ to_test = [(ExprInt(1, 32) - ExprInt(1, 32), ExprInt(0, 32)),
            (ExprOp('<<', ExprOp('>>', a, ExprInt(0x4, 32)), ExprInt(0x4, 32)),
             ExprOp('&', a, ExprInt(0xFFFFFFF0, 32))),
 
-           (ExprCompose(ExprInt(0x1234, 16), ExprId("a", 16)) & ExprInt(0xFF, 32), ExprInt(0x34, 32)),
+           (ExprCompose(ExprId("a", 8), ExprId("b", 24)) & ExprInt(0xFF, 32), ExprCompose(ExprId("a", 8), ExprInt(0x0, 24))),
            (ExprCompose(ExprInt(0x12, 8), ExprInt(0x34, 8)) & ExprInt(0xFFFF, 16), ExprInt(0x3412, 16)),
-           (ExprCompose(ExprInt(0x12, 8), ExprInt(0x3456, 16), ExprInt(0x78, 8)) & ExprInt(0xFFFFFF, 32), ExprInt(0x345612, 32)),
-           (ExprCompose(ExprInt(0x1234, 16), ExprId("a", 8), ExprInt(0x67, 8)) & ExprInt(0xFFFFFFFF, 32), ExprCompose(ExprInt(0x1234, 16), ExprId("a", 8), ExprInt(0x67, 8))),
+           (ExprCompose(ExprInt(0x12, 8), ExprInt(0x345678, 24)) & ExprInt(0xFFFFFF, 32), ExprInt(0x567812, 32)),
 
            (a[:32], a),
            (a[:8][:8], a[:8]),