VariablesIdentifier: Handle corner cases (using var ident on already computed eq)

1 files changed, 34 insertions, 3 deletions

diff --git a/test/expression/expression_helper.py b/test/expression/expression_helper.py
index e1674fdc..514a9a51 100644
--- a/test/expression/expression_helper.py
+++ b/test/expression/expression_helper.py
@@ -30,6 +30,7 @@ class TestExpressionExpressionHelper(unittest.TestCase):
 
         # Test the result
         new_expr = vi.equation
+
         ## Force replace in the variable dependency order
         for var_id, var_value in reversed(vi.vars.items()):
             new_expr = new_expr.replace_expr({var_id: var_value})
@@ -38,16 +39,46 @@ class TestExpressionExpressionHelper(unittest.TestCase):
         # Test prefix
         vi = Variables_Identifier(exprf, var_prefix="prefix_v")
 
-        # Use __str__
+        ## Use __str__
         print vi
 
-        # Test the result
+        ## Test the result
         new_expr = vi.equation
-        ## Force replace in the variable dependency order
+        ### Force replace in the variable dependency order
         for var_id, var_value in reversed(vi.vars.items()):
             new_expr = new_expr.replace_expr({var_id: var_value})
         self.assertEqual(exprf, new_expr)
 
+        # Test an identify on an expression already containing identifier
+        vi = Variables_Identifier(exprf)
+        vi2 = Variables_Identifier(vi.equation)
+
+        ## Test the result
+        new_expr = vi2.equation
+        ### Force replace in the variable dependency order
+        for var_id, var_value in reversed(vi2.vars.items()):
+            new_expr = new_expr.replace_expr({var_id: var_value})
+        self.assertEqual(vi.equation, new_expr)
+
+        ## Corner case: each sub var depends on itself
+        mem1 = m2_expr.ExprMem(ebx, size=32)
+        mem2 = m2_expr.ExprMem(mem1, size=32)
+        cst2 = m2_expr.ExprInt32(-1)
+        expr_mini = ((eax ^ mem2 ^ cst2) & (mem2 ^ (eax + mem2)))[31:32]
+
+        ## Build
+        vi = Variables_Identifier(expr_mini)
+        vi2 = Variables_Identifier(vi.equation)
+
+        ## Test the result
+        new_expr = vi2.equation
+        ### Force replace in the variable dependency order
+        for var_id, var_value in reversed(vi2.vars.items()):
+            new_expr = new_expr.replace_expr({var_id: var_value})
+        self.assertEqual(vi.equation, new_expr)
+
+
+
 if __name__ == '__main__':
     testcase = TestExpressionExpressionHelper
     testsuite = unittest.TestLoader().loadTestsFromTestCase(testcase)