Repository: angr/claripy

Issue:
Calling solver max twice changes the result.
### Description

Calling `s.max(x)` twice in a row yields different results. I imagine this might be due to caching, as I can see the first invocation of `s.max(x)` calls out to `z3`'s backend, where as the second one does not. This invocation also adds extra constraints to the solver.

Note: calling `min` twice also has odd effects, though since the initial return value of `s.min(x)` is 0, the return value still returns correctly the second time.

### Steps to reproduce the bug

```python
>>> import claripy
>>> s = claripy.Solver()
>>> x = claripy.BVS("x", 32)
>>> s.max(x)
4294967295
>>> s.max(x)
0
```

### Environment

New venv with fresh clone of master: https://github.com/angr/claripy/commit/33483ea55d6074270f65e2e2b106016be814430f

### Additional context

Notice that `s.constraints` is changed by the max invocation:

```python
>>> import claripy
>>> s = claripy.Solver()
>>> x = claripy.BVS("x", 32)
>>> s.constraints
[]
>>> s.max(x)
4294967295
>>> s.constraints
[<Bool x_0_32 <= 0xffffffff>]
>>> s.max(x)
0
```

Discussion:

Patch:
diff --git a/claripy/backends/backend_z3.py b/claripy/backends/backend_z3.py
index 6a133e05..0ace4228 100644
--- a/claripy/backends/backend_z3.py
+++ b/claripy/backends/backend_z3.py
@@ -65,7 +65,7 @@ def _add_memory_pressure(p):
     This is not a problem for CPython since its GC is based on reference counting.
     """
 
-    global _is_pypy
+    global _is_pypy # pylint:disable=global-variable-not-assigned
     if _is_pypy:
         __pypy__.add_memory_pressure(p)
 
@@ -136,7 +136,6 @@ class BackendZ3(Backend):
 
     def __init__(self, reuse_z3_solver=None, ast_cache_size=10000):
         Backend.__init__(self, solver_required=True)
-        self._enable_simplification_cache = False
 
         # Per-thread Z3 solver
         # This setting is treated as a global setting and is not supposed to be changed during runtime, unless you know
@@ -257,30 +256,12 @@ class BackendZ3(Backend):
             self._tls.sym_cache = weakref.WeakValueDictionary()
             return self._tls.sym_cache
 
-    @property
-    def _simplification_cache_key(self):
-        try:
-            return self._tls.simplification_cache_key
-        except AttributeError:
-            self._tls.simplification_cache_key = weakref.WeakValueDictionary()
-            return self._tls.simplification_cache_key
-
-    @property
-    def _simplification_cache_val(self):
-        try:
-            return self._tls.simplification_cache_val
-        except AttributeError:
-            self._tls.simplification_cache_val = weakref.WeakValueDictionary()
-            return self._tls.simplification_cache_val
-
     def downsize(self):
         Backend.downsize(self)
 
         self._ast_cache.clear()
         self._var_cache.clear()
         self._sym_cache.clear()
-        self._simplification_cache_key.clear()
-        self._simplification_cache_val.clear()
 
     def _name(self, o): #pylint:disable=unused-argument
         l.warning("BackendZ3.name() called. This is weird.")
@@ -378,7 +359,7 @@ class BackendZ3(Backend):
     #
 
     @condom
-    def _convert(self, obj):  # pylint:disable=arguments-differ
+    def _convert(self, obj):  # pylint:disable=arguments-renamed
         if isinstance(obj, FSort):
             return z3.FPSortRef(z3.Z3_mk_fpa_sort(self._context.ref(), obj.exp, obj.mantissa), self._context)
         elif isinstance(obj, RM):
@@ -406,7 +387,7 @@ class BackendZ3(Backend):
             l.debug("BackendZ3 encountered unexpected type %s", type(obj))
             raise BackendError("unexpected type %s encountered in BackendZ3" % type(obj))
 
-    def call(self, *args, **kwargs):  # pylint;disable=arguments-differ
+    def call(self, *args, **kwargs):  # pylint;disable=arguments-renamed
         return Backend.call(self, *args, **kwargs)
 
     @condom
@@ -975,22 +956,6 @@ class BackendZ3(Backend):
         if expr._simplified:
             return expr
 
-        if self._enable_simplification_cache:
-            try:
-                k = self._simplification_cache_key[expr._cache_key]
-                #print "HIT WEAK KEY CACHE"
-                return k
-            except KeyError:
-                pass
-            try:
-                k = self._simplification_cache_val[expr._cache_key]
-                #print "HIT WEAK VALUE CACHE"
-                return k
-            except KeyError:
-                pass
-
-            #print "MISS CACHE"
-
         #l.debug("SIMPLIFYING EXPRESSION")
 
         expr_raw = self.convert(expr)
@@ -1016,9 +981,6 @@ class BackendZ3(Backend):
         o = self._abstract(s)
         o._simplified = Base.FULL_SIMPLIFY
 
-        if self._enable_simplification_cache:
-            self._simplification_cache_val[expr._cache_key] = o
-            self._simplification_cache_key[expr._cache_key] = o
         return o
 
     def _is_false(self, e, extra_constraints=(), solver=None, model_callback=None):
diff --git a/claripy/frontend_mixins/model_cache_mixin.py b/claripy/frontend_mixins/model_cache_mixin.py
index 57965461..b773b2ee 100644
--- a/claripy/frontend_mixins/model_cache_mixin.py
+++ b/claripy/frontend_mixins/model_cache_mixin.py
@@ -1,3 +1,4 @@
+from typing import Tuple
 import weakref
 import itertools
 
@@ -10,6 +11,7 @@ class ModelCache:
     def __init__(self, model):
         self.model = model
         self.replacements = weakref.WeakKeyDictionary()
+        self.constraint_only_replacements = weakref.WeakKeyDictionary()
 
     def __hash__(self):
         if not hasattr(self, '_hash'):
@@ -25,6 +27,7 @@ class ModelCache:
     def __setstate__(self, s):
         self.model = s[0]
         self.replacements = weakref.WeakKeyDictionary()
+        self.constraint_only_replacements = weakref.WeakKeyDictionary()
 
     #
     # Splitting support
@@ -50,12 +53,29 @@ class ModelCache:
             a
         )
 
-    def eval_ast(self, ast):
-        """Eval the ast, replacing symbols by their last value in the model.
+    def _leaf_op_existonly(self, a):
+        return (
+            all_operations.BVV(self.model[a.args[0]], a.length) if a.op == 'BVS' else
+            all_operations.BoolV(self.model[a.args[0]]) if a.op == 'BoolS' else
+            all_operations.FPV(self.model[a.args[0]], a.args[1]) if a.op == 'FPS' else
+            all_operations.StringV(self.model[a.args[0]]) if a.op == 'StringS' else
+            a
+        )
+
+    def eval_ast(self, ast, allow_unconstrained: bool=True):
         """
-        # If there was no last value, it was not constrained, so we can use
-        # anything.
-        new_ast = ast.replace_dict(self.replacements, leaf_operation=self._leaf_op)
+        Eval the ast, replacing symbols by their last value in the model.
+
+        :param ast:                 The AST to evaluate.
+        :param allow_unconstrained: When set to True, we will treat non-existent variables as unconstrained variables
+                                    and will use arbitrary concrete values for them during evaluation. Otherwise, raise
+                                    KeyErrors for non-existent variables.
+        """
+
+        if allow_unconstrained:
+            new_ast = ast.replace_dict(self.replacements, leaf_operation=self._leaf_op)
+        else:
+            new_ast = ast.replace_dict(self.constraint_only_replacements, leaf_operation=self._leaf_op_existonly)
         return backends.concrete.eval(new_ast, 1)[0]
 
     def eval_constraints(self, constraints):
@@ -68,8 +88,19 @@ class ModelCache:
         except errors.ClaripyZeroDivisionError:
             return False
 
-    def eval_list(self, asts):
-        return tuple(self.eval_ast(c) for c in asts)
+    def eval_list(self, asts, allow_unconstrained: bool=True) -> Tuple:
+        """
+        Evaluate a list of ASTs.
+
+        :param asts:                A list of ASTs to evaluate.
+        :param allow_unconstrained: When set to True, we will treat non-existent variables as unconstrained variables
+                                    and will use arbitrary concrete values for them during evaluation. Otherwise, raise
+                                    KeyErrors for non-existent variables.
+        :return:                    A tuple of evaluated results, one element per AST.
+        """
+
+        return tuple(self.eval_ast(c, allow_unconstrained=allow_unconstrained) for c in asts)
+
 
 class ModelCacheMixin:
     def __init__(self, *args, **kwargs):
@@ -221,30 +252,31 @@ class ModelCacheMixin:
         # Z3 might give us solutions for variables that we did not ask for. so we create a new dict with solutions for
         # only the variables that are under the solver's control
         m_ = dict((k, v) for k, v in m.items() if k in self.variables)
-        model = ModelCache(m_)
-        self._models.add(model)
+        if m_:
+            model = ModelCache(m_)
+            self._models.add(model)
 
     def _get_models(self, extra_constraints=()):
         for m in self._models:
             if m.eval_constraints(extra_constraints):
                 yield m
 
-    def _get_batch_solutions(self, asts, n=None, extra_constraints=()):
+    def _get_batch_solutions(self, asts, n=None, extra_constraints=(), allow_unconstrained=True):
         results = set()
 
         for m in self._get_models(extra_constraints):
             try:
-                results.add(m.eval_list(asts))
-            except ZeroDivisionError:
+                results.add(m.eval_list(asts, allow_unconstrained=allow_unconstrained))
+            except (ZeroDivisionError, KeyError):
                 continue
             if len(results) == n:
                 break
 
         return results
 
-    def _get_solutions(self, e, n=None, extra_constraints=()):
+    def _get_solutions(self, e, n=None, extra_constraints=(), allow_unconstrained=True):
         return tuple(v[0] for v in self._get_batch_solutions(
-            [e], n=n, extra_constraints=extra_constraints
+            [e], n=n, extra_constraints=extra_constraints, allow_unconstrained=allow_unconstrained,
         ))
 
 
@@ -283,7 +315,11 @@ class ModelCacheMixin:
                 raise
 
         if len(extra_constraints) == 0 and len(results) < n:
-            self._eval_exhausted.update(e.cache_key for e in asts)
+            for e in asts:
+                # only mark an AST as eval-exhausted if e.variables is a subset of variables that the current solver
+                # knows about (from its constraints)
+                if self.variables.issuperset(e.variables):
+                    self._eval_exhausted.add(e.cache_key)
 
         return results
 
@@ -293,7 +329,9 @@ class ModelCacheMixin:
     def min(self, e, extra_constraints=(), signed=False, **kwargs):
         cached = [ ]
         if e.cache_key in self._eval_exhausted or e.cache_key in self._min_exhausted:
-            cached = self._get_solutions(e, extra_constraints=extra_constraints)
+            # we set allow_unconstrained to False because we expect all returned values for e are returned by Z3,
+            # instead of some arbitrarily assigned concrete values.
+            cached = self._get_solutions(e, extra_constraints=extra_constraints, allow_unconstrained=False)
 
         if len(cached) > 0:
             signed_key = lambda v: v if v < 2**(len(e)-1) else v - 2**len(e)
@@ -307,7 +345,7 @@ class ModelCacheMixin:
     def max(self, e, extra_constraints=(), signed=False, **kwargs):
         cached = [ ]
         if e.cache_key in self._eval_exhausted or e.cache_key in self._max_exhausted:
-            cached = self._get_solutions(e, extra_constraints=extra_constraints)
+            cached = self._get_solutions(e, extra_constraints=extra_constraints, allow_unconstrained=False)
 
         if len(cached) > 0:
             signed_key = lambda v: v if v < 2**(len(e)-1) else v - 2**len(e)


Test patch:
diff --git a/tests/test_expression.py b/tests/test_expression.py
index 509c7d36..6627f900 100644
--- a/tests/test_expression.py
+++ b/tests/test_expression.py
@@ -246,6 +246,9 @@ class TestExpression(unittest.TestCase):
 
         ite = claripy.ite_dict(x, {1: 11, 2: 22, 3: 33, 4: 44, 5: 55, 6: 66, 7: 77, 8: 88, 9: 99}, claripy.BVV(0, 32))
         self.assertEqual(sorted(s.eval(ite, 100)), [0, 11, 22, 33, 44, 55, 66, 77, 88, 99])
+        # calling it a second time should not fail. it was failing on master due to we incorrectly putting the AST
+        # into the eval-exhausted set.
+        self.assertEqual(sorted(s.eval(ite, 100)), [0, 11, 22, 33, 44, 55, 66, 77, 88, 99])
 
         ss = s.branch()
         ss.add(ite == 88)
diff --git a/tests/test_solver.py b/tests/test_solver.py
index bdc269f4..54aee73f 100644
--- a/tests/test_solver.py
+++ b/tests/test_solver.py
@@ -608,6 +608,22 @@ class StandardTests(TestCase):
         print(s.max(x, extra_constraints=[x <= 18]))
         assert s.max(x) == 19
 
+    def test_cached_max(self):
+        s = claripy.Solver()
+        x = claripy.BVS("x", 32)
+        assert not s.constraints
+        assert s.max(x) == 0xffffffff
+        assert len(s.constraints) == 1  # ConstraintExpansionMixin will add a new constraint
+        assert s.max(x) == 0xffffffff  # calling it the second time, the cache should not give a different result
+
+        s = claripy.Solver()
+        y = claripy.BVS("y", 32)
+        s.add(y == 8)
+        assert s.eval(y, n=1)[0] == 8
+        assert len(s.constraints) == 1
+        assert s.max(x) == 0xffffffff
+        assert s.max(x) == 0xffffffff
+
 #
 # Multi-Solver test base classes
 #