Better iterative optimization for expressions

manticore/core/smtlib/solver.py

@@ -39,6 +39,9 @@
 )
 consts.add("z3_bin", default="z3", description="Z3 binary to use")
 consts.add("defaultunsat", default=True, description="Consider solver timeouts as unsat core")
+consts.add(
+    "optimize", default=True, description="Use smtlib command optimize to find min/max if available"
+)
 
 
 # Regular expressions used by the solver
@@ -464,18 +467,18 @@ def get_all_values(self, constraints, expression, maxcnt=None, silent=False):
                     raise SolverError("Timeout")
             return result
 
-    def optimize(self, constraints: ConstraintSet, x: BitVec, goal: str, M=10000):
+    def optimize(self, constraints: ConstraintSet, x: BitVec, goal: str, max_iter=10000):
         """
         Iteratively finds the maximum or minimum value for the operation
         (Normally Operators.UGT or Operators.ULT)
 
         :param constraints: constraints to take into account
         :param x: a symbol or expression
         :param goal: goal to achieve, either 'maximize' or 'minimize'
-        :param M: maximum number of iterations allowed
+        :param max_iter: maximum number of iterations allowed
         """
+        # TODO: consider adding a mode to return best known value on timeout
         assert goal in ("maximize", "minimize")
-        assert isinstance(x, BitVec)
         operation = {"maximize": Operators.UGE, "minimize": Operators.ULE}[goal]
 
         with constraints as temp_cs:
@@ -486,7 +489,7 @@ def optimize(self, constraints: ConstraintSet, x: BitVec, goal: str, M=10000):
             self._send(aux.declaration)
 
             start = time.time()
-            if getattr(self, f"support_{goal}"):
+            if consts.optimize and getattr(self, f"support_{goal}", False):
                 self._push()
                 try:
                     self._assert(operation(X, aux))
@@ -498,9 +501,9 @@ def optimize(self, constraints: ConstraintSet, x: BitVec, goal: str, M=10000):
                         # This will be a line like NAME |-> VALUE
                         maybe_sat = self._recv()
                         if maybe_sat == "sat":
-                            m = RE_MIN_MAX_OBJECTIVE_EXPR_VALUE.match(_status)
-                            if m:
-                                expr, value = m.group("expr"), m.group("value")
+                            match = RE_MIN_MAX_OBJECTIVE_EXPR_VALUE.match(_status)
+                            if match:
+                                expr, value = match.group("expr"), match.group("value")
                                 assert expr == aux.name
                                 return int(value)
                             else:
@@ -510,9 +513,9 @@ def optimize(self, constraints: ConstraintSet, x: BitVec, goal: str, M=10000):
                         if not (ret.startswith("(") and ret.endswith(")")):
                             raise SolverError("bad output on max, z3 may have been killed")
 
-                        m = RE_OBJECTIVES_EXPR_VALUE.match(ret)
-                        if m:
-                            expr, value = m.group("expr"), m.group("value")
+                        match = RE_OBJECTIVES_EXPR_VALUE.match(ret)
+                        if match:
+                            expr, value = match.group("expr"), match.group("value")
                             assert expr == aux.name
                             return int(value)
                         else:
@@ -522,15 +525,56 @@ def optimize(self, constraints: ConstraintSet, x: BitVec, goal: str, M=10000):
                     self._reset(temp_cs)
                     self._send(aux.declaration)
 
-            operation = {"maximize": Operators.UGT, "minimize": Operators.ULT}[goal]
+            operation = {"maximize": Operators.UGE, "minimize": Operators.ULE}[goal]
             self._assert(aux == X)
+
+            # Find one value and use it as currently known min/Max
+            if not self._is_sat():
+                raise SolverException("UNSAT")
+            last_value = self._getvalue(aux)
+            self._assert(operation(aux, last_value))
+
+            # This uses a binary search to find a suitable range for aux
+            # Use known solution as min or max depending on the goal
+            if goal == "maximize":
+                m, M = last_value, (1 << x.size) - 1
+            else:
+                m, M = 0, last_value
+
+            # Iteratively divide the range
+            L = None
+            while L not in (M, m):
+                L = (m + M) // 2
+                self._push()
+                try:
+                    self._assert(operation(aux, L))
+                    sat = self._is_sat()
+                finally:
+                    self._pop()
+
+                # depending on the goal move one of the extremes
+                if goal == "maximize" and sat or goal == "minimize" and not sat:
+                    m = L
+                else:
+                    M = L
+
+                if time.time() - start > consts.timeout:
+                    raise SolverError("Timeout")
+
+            # At this point we know aux is inside [m,M]
+            # Lets constrain it to that range
+            self._assert(Operators.UGE(aux, m))
+            self._assert(Operators.ULE(aux, M))
+
+            # And now check all remaining possible extremes
             last_value = None
             i = 0
             while self._is_sat():
                 last_value = self._getvalue(aux)
                 self._assert(operation(aux, last_value))
+                self._assert(aux != last_value)
                 i = i + 1
-                if i > M:
+                if i > max_iter:
                     raise SolverError("Optimizing error, maximum number of iterations was reached")
                 if time.time() - start > consts.timeout:
                     raise SolverError("Timeout")

tests/other/test_smtlibv2.py

@@ -429,6 +429,25 @@ def testBitvector_max(self):
         cs.add(a >= 100)
         self.assertTrue(self.solver.check(cs))
         self.assertEqual(self.solver.minmax(cs, a), (100, 200))
+        from manticore import config
+
+        consts = config.get_group("smt")
+        consts.optimize = False
+        cs = ConstraintSet()
+        a = cs.new_bitvec(32)
+        cs.add(a <= 200)
+        cs.add(a >= 100)
+        self.assertTrue(self.solver.check(cs))
+        self.assertEqual(self.solver.minmax(cs, a), (100, 200))
+        consts.optimize = True
+
+    def testBitvector_max_noop(self):
+        from manticore import config
+
+        consts = config.get_group("smt")
+        consts.optimize = False
+        self.testBitvector_max()
+        consts.optimize = True
 
     def testBitvector_max1(self):
         cs = ConstraintSet()
@@ -438,6 +457,14 @@ def testBitvector_max1(self):
         self.assertTrue(self.solver.check(cs))
         self.assertEqual(self.solver.minmax(cs, a), (101, 199))
 
+    def testBitvector_max1_noop(self):
+        from manticore import config
+
+        consts = config.get_group("smt")
+        consts.optimize = False
+        self.testBitvector_max1()
+        consts.optimize = True
+
     def testBool_nonzero(self):
         self.assertTrue(BoolConstant(True).__bool__())
         self.assertFalse(BoolConstant(False).__bool__())