Expression: add ExprReduce

3 files changed, 271 insertions, 0 deletions

diff --git a/example/expression/expr_reduce.py b/example/expression/expr_reduce.py
new file mode 100644
index 00000000..b5fc96c8
--- /dev/null
+++ b/example/expression/expr_reduce.py
@@ -0,0 +1,93 @@
+from miasm2.expression.expression import ExprId, ExprInt, ExprMem
+from miasm2.expression.expression_reduce import ExprReducer
+
+
+class StructLookup(ExprReducer):
+    """
+    ExprReduce example.
+    This example retrieve the nature of a given expression
+    Input:
+    ECX is a pointer on a structure STRUCT_A
+
+    Reduction rules:
+    ECX              -> FIELD_A_PTR
+    ECX + CST        -> FIELD_A_PTR
+    ECX + CST*CST... -> FIELD_A_PTR
+    @ECX             -> FIELD_A
+    @(ECX + CST)     -> FIELD_A
+    """
+    CST = "CST"
+    FIELD_A_PTR = "FIELD_A_PTR"
+    FIELD_A = "FIELD_A"
+
+    def reduce_int(self, node, _):
+        """
+        Reduction: int -> CST
+        """
+        if node.expr.is_int():
+            return self.CST
+        return None
+
+    def reduce_ptr_struct(self, node, _):
+        """
+        Reduction: ECX -> FIELD_A_PTR
+        """
+        if node.expr.is_id("ECX"):
+            return self.FIELD_A_PTR
+        return None
+
+    def reduce_ptr_plus_int(self, node, _):
+        """
+        Reduction: ECX + CST -> FIELD_A_PTR
+        """
+        if not node.expr.is_op('+'):
+            return None
+        if [arg.info for arg in node.args] == [self.FIELD_A_PTR, self.CST]:
+            return self.FIELD_A_PTR
+        return None
+
+    def reduce_cst_op(self, node, _):
+        """
+        Reduction: CST + CST -> CST
+        """
+        if not node.expr.is_op():
+            return None
+        if set(arg.info for arg in node.args) == set([self.CST]):
+            return self.CST
+        return None
+
+    def reduce_at_struct_ptr(self, node, _):
+        """
+        Reduction: @FIELD_A_PTR -> FIELD_A
+        """
+        if not node.expr.is_mem():
+            return None
+        return self.FIELD_A
+
+    reduction_rules = [reduce_int,
+                       reduce_ptr_struct,
+                       reduce_ptr_plus_int,
+                       reduce_cst_op,
+                       reduce_at_struct_ptr
+                      ]
+
+
+def test():
+    struct_lookup = StructLookup()
+
+    ptr = ExprId('ECX')
+    int4 = ExprInt(4, 32)
+    tests = [
+        (ptr, StructLookup.FIELD_A_PTR),
+        (ptr + int4, StructLookup.FIELD_A_PTR),
+        (ptr + int4 * int4, StructLookup.FIELD_A_PTR),
+        (ExprMem(ptr), StructLookup.FIELD_A),
+        (ExprMem(ptr + int4 * int4), StructLookup.FIELD_A),
+    ]
+
+    for expr_in, result in tests:
+        assert struct_lookup.reduce(expr_in).info == result
+
+
+if __name__ == "__main__":
+    test()
diff --git a/miasm2/expression/expression_reduce.py b/miasm2/expression/expression_reduce.py
new file mode 100644
index 00000000..f4a68be4
--- /dev/null
+++ b/miasm2/expression/expression_reduce.py
@@ -0,0 +1,177 @@
+"""
+Expression reducer:
+Apply reduction rules to an Expression ast
+"""
+
+import logging
+from miasm2.expression.expression import ExprInt, ExprId, ExprOp, ExprSlice,\
+    ExprCompose, ExprMem, ExprCond
+
+log_reduce = logging.getLogger("expr_reduce")
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(logging.Formatter("%(levelname)-5s: %(message)s"))
+log_reduce.addHandler(console_handler)
+log_reduce.setLevel(logging.WARNING)
+
+
+class ExprNode(object):
+    """Clone of Expression object with additionnal information"""
+
+    def __init__(self, expr):
+        self.expr = expr
+        # Generic field to store custom node information
+        self.info = None
+
+        self.arg, self.args = None, None
+        self.cond, self.src1, self.src2 = None, None, None
+
+    def __repr__(self):
+        expr = self.expr
+        if self.info is not None:
+            out = repr(self.info)
+        elif isinstance(expr, (ExprInt, ExprId)):
+            out = str(None)
+        elif isinstance(expr, ExprMem):
+            out = "@%d[%r]" % (self.expr.size, self.arg)
+        elif isinstance(expr, ExprSlice):
+            out = "%r[%d:%d]" % (self.arg, expr.start, expr.stop)
+        elif isinstance(expr, ExprOp):
+            if len(self.args) == 1:
+                out = "%s(%s)" % (expr.op, self.args[0])
+            else:
+                out = expr.op.join([repr(arg) for arg in self.args])
+        elif isinstance(expr, ExprCompose):
+            out = "{%s}" % ', '.join([repr(arg) for arg in self.args])
+        elif isinstance(expr, ExprCond):
+            out = "%r?%r:%r" % (self.cond, self.src1, self.src2)
+        else:
+            raise TypeError("Unknown node Type %r", type(expr))
+        return out
+
+
+class ExprReducer(object):
+    """Apply reduction rules to an expr
+
+    reduction_rules: list of ordered reduction rules
+
+    List of function representing reduction rules
+    Function API:
+    reduction_xxx(self, node, lvl=0)
+    with:
+    * node: the ExprNode to qualify
+    * lvl: [optional] the recursion level
+    Returns:
+    * None if the reduction rule is not applied
+    * the resulting information to store in the ExprNode.info
+
+    allow_none_result: allow missing reduction rules
+    """
+
+    reduction_rules = []
+    allow_none_result = False
+
+    def expr2node(self, expr):
+        """Build ExprNode mirror of @expr
+
+        @expr: Expression to analyze
+        """
+
+        if isinstance(expr, (ExprId, ExprInt)):
+            node = ExprNode(expr)
+        elif isinstance(expr, (ExprMem, ExprSlice)):
+            son = self.expr2node(expr.arg)
+            node = ExprNode(expr)
+            node.arg = son
+        elif isinstance(expr, ExprOp):
+            sons = [self.expr2node(arg) for arg in expr.args]
+            node = ExprNode(expr)
+            node.args = sons
+        elif isinstance(expr, ExprCompose):
+            sons = [self.expr2node(arg) for arg in expr.args]
+            node = ExprNode(expr)
+            node.args = sons
+        elif isinstance(expr, ExprCond):
+            node = ExprNode(expr)
+            node.cond = self.expr2node(expr.cond)
+            node.src1 = self.expr2node(expr.src1)
+            node.src2 = self.expr2node(expr.src2)
+        else:
+            raise TypeError("Unknown Expr Type %r", type(expr))
+        return node
+
+    def reduce(self, expr):
+        """Returns an ExprNode tree mirroring @expr tree. The ExprNode is
+        computed by applying reduction rules to the expression @expr
+
+        @expr: an Expression
+        """
+
+        node = self.expr2node(expr)
+        return self.categorize(node, 0)
+
+    def categorize(self, node, lvl=0):
+        """Recursively apply rules to @node
+
+        @node: ExprNode to analyze
+        @lvl: actual recusion level
+        """
+
+        expr = node.expr
+        log_reduce.debug("\t" * lvl + "Reduce...: %s", node.expr)
+        if isinstance(expr, (ExprId, ExprInt)):
+            pass
+        elif isinstance(expr, ExprMem):
+            arg = self.categorize(node.arg, lvl + 1)
+            node = ExprNode(ExprMem(arg.expr, expr.size))
+            node.arg = arg
+        elif isinstance(expr, ExprSlice):
+            arg = self.categorize(node.arg, lvl + 1)
+            node = ExprNode(ExprSlice(arg.expr, expr.start, expr.stop))
+            node.arg = arg
+        elif isinstance(expr, ExprOp):
+            new_args = []
+            for arg in node.args:
+                new_a = self.categorize(arg, lvl + 1)
+                assert new_a.expr.size == arg.expr.size
+                new_args.append(new_a)
+            node = ExprNode(ExprOp(expr.op, *[x.expr for x in new_args]))
+            node.args = new_args
+            expr = node.expr
+        elif isinstance(expr, ExprCompose):
+            new_args = []
+            new_expr_args = []
+            for arg in node.args:
+                arg = self.categorize(arg, lvl + 1)
+                new_args.append(arg)
+                new_expr_args.append(arg.expr)
+            new_expr = ExprCompose(*new_expr_args)
+            node = ExprNode(new_expr)
+            node.args = new_args
+        elif isinstance(expr, ExprCond):
+            cond = self.categorize(node.cond, lvl + 1)
+            src1 = self.categorize(node.src1, lvl + 1)
+            src2 = self.categorize(node.src2, lvl + 1)
+            node = ExprNode(ExprCond(cond.expr, src1.expr, src2.expr))
+            node.cond, node.src1, node.src2 = cond, src1, src2
+        else:
+            raise TypeError("Unknown Expr Type %r", type(expr))
+
+        node.info = self.apply_rules(node, lvl)
+        log_reduce.debug("\t" * lvl + "Reduce result: %s %r",
+                         node.expr, node.info)
+        return node
+
+    def apply_rules(self, node, lvl=0):
+        """Find and apply reduction rules to @node
+
+        @node: ExprNode to analyse
+        @lvl: actuel recusion level
+        """
+
+        for rule in self.reduction_rules:
+            ret = rule(self, node, lvl)
+            if ret is not None:
+                log_reduce.debug("\t" * lvl + "Rule found: %r", rule)
+                return ret
+        if not self.allow_none_result:
+            raise RuntimeError('Missing reduction rule for %r' % node.expr)
diff --git a/test/test_all.py b/test/test_all.py
index 2f8c6421..e405bcec 100755
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -524,6 +524,7 @@ for script in [["basic_op.py"],
                ["simplification_add.py"],
                ["expr_random.py"],
                ["expr_translate.py"],
+               ["expr_reduce.py"],
                ]:
     testset += ExampleExpression(script)