1 files changed, 156 insertions, 0 deletions
diff --git a/miasm2/expression/expression_helper.py b/miasm2/expression/expression_helper.py
index 9c98b08f..2ce5c034 100644
--- a/miasm2/expression/expression_helper.py
+++ b/miasm2/expression/expression_helper.py
@@ -17,6 +17,10 @@
 #
 
 # Expressions manipulation functions
+import re
+import itertools
+import collections
+
 import miasm2.expression.expression as m2_expr
 
 
@@ -195,3 +199,155 @@ def get_missing_interval(all_intervals, i_min=0, i_max=32):
             missing_i.append((last_pos, start))
         last_pos = stop
     return missing_i
+
+
+class Variables_Identifier(object):
+    """Identify variables in an expression.
+    Returns:
+    - variables with their corresponding values
+    - original expression with variables translated
+    """
+
+    var_identifier = re.compile("v\d+")
+
+    def __init__(self, expr):
+        """Set the expression @expr to handle and launch variable identification
+        process"""
+
+        # Init
+        self.var_indice = itertools.count()
+        self.var_asked = set()
+        self._vars = {} # VarID -> Expr
+
+        # Launch recurrence
+        self.find_variables_rec(expr)
+
+        # Compute inter-variable dependencies
+        has_change = True
+        while has_change:
+            has_change = False
+            for var_id, var_value in self._vars.items():
+                cur = var_value
+
+                # Do not replace with itself
+                to_replace = {v_val:v_id for v_id, v_val in self._vars.items()
+                              if v_id != var_id}
+                var_value = var_value.replace_expr(to_replace)
+
+                if cur != var_value:
+                    # Force @self._vars update
+                    has_change = True
+                    self._vars[var_id] = var_value
+                    break
+
+        # Replace in the original equation
+        self._equation = expr.replace_expr({v_val: v_id for v_id, v_val
+                                            in self._vars.items()})
+
+        # Compute variables dependencies
+        self._vars_ordered = collections.OrderedDict()
+        todo = set(self._vars.keys())
+        needs = {}
+
+        ## Build initial needs
+        for var_id, var_expr in self._vars.items():
+            needs[var_id] = [var_name
+                             for var_name in var_expr.get_r(mem_read=True)
+                             if self.is_var_identifier(var_name)]
+
+        ## Build order list
+        while todo:
+            done = set()
+            for var_id in todo:
+                all_met = True
+                for need in needs[var_id]:
+                    if need not in self._vars_ordered:
+                        # A dependency is not met
+                        all_met = False
+                        break
+
+                if not all_met:
+                    continue
+
+                # All dependencies are already met, add current
+                self._vars_ordered[var_id] = self._vars[var_id]
+                done.add(var_id)
+
+            # Update the todo list
+            for element_done in done:
+                todo.remove(element_done)
+
+    @classmethod
+    def is_var_identifier(cls, expr):
+        "Return True iff expr seems to be a variable identifier"
+        if not isinstance(expr, m2_expr.ExprId):
+            return False
+
+        match = cls.var_identifier.match(expr.name)
+        return match is not None and match.group(0) == expr.name
+
+    def find_variables_rec(self, expr):
+        """Recursive method called by find_variable to expand @expr.
+        Set @var_names and @var_values.
+        This implementation is faster than an expression visitor because
+        we do not rebuild each expression.
+        """
+
+        if (expr in self.var_asked):
+            # Expr has already been asked
+
+            if (expr not in self._vars.values()):
+                # Create var
+                identifier = m2_expr.ExprId("v%s" % self.var_indice.next(),
+                                    size = expr.size)
+                self._vars[identifier] = expr
+
+            # Recursion stop case
+            return
+        else:
+            # First time for @expr
+            self.var_asked.add(expr)
+
+        if isinstance(expr, m2_expr.ExprOp):
+            for a in expr.args:
+                self.find_variables_rec(a)
+
+        elif isinstance(expr, m2_expr.ExprInt):
+            pass
+
+        elif isinstance(expr, m2_expr.ExprId):
+            pass
+
+        elif isinstance(expr, m2_expr.ExprMem):
+            self.find_variables_rec(expr.arg)
+
+        elif isinstance(expr, m2_expr.ExprCompose):
+            for a in expr.args:
+                self.find_variables_rec(list(a)[0])
+
+        elif isinstance(expr, m2_expr.ExprSlice):
+            self.find_variables_rec(expr.arg)
+
+        elif isinstance(expr, m2_expr.ExprCond):
+            self.find_variables_rec(expr.cond)
+            self.find_variables_rec(expr.src1)
+            self.find_variables_rec(expr.src2)
+
+        else:
+            raise NotImplementedError("Type not handled: %s" % expr)
+
+    @property
+    def vars(self):
+        return self._vars_ordered
+
+    @property
+    def equation(self):
+        return self._equation
+
+    def __str__(self):
+        "Display variables and final equation"
+        out = ""
+        for var_id, var_expr in self.vars.items():
+            out += "%s = %s\n" % (var_id, var_expr)
+        out += "Final: %s" % self.equation
+        return out