add assert_and_track to optimize for #2116

Signed-off-by: Nikolaj Bjorner <nbjorner@microsoft.com>

src/api/api_opt.cpp

@@ -79,6 +79,16 @@ extern "C" {
         Z3_CATCH;
     }
 
+    void Z3_API Z3_optimize_assert_and_track(Z3_context c, Z3_optimize o, Z3_ast a, Z3_ast t) {
+        Z3_TRY;
+        LOG_Z3_optimize_assert_and_track(c, o, a, t);
+        RESET_ERROR_CODE();
+        CHECK_FORMULA(a,);        
+        CHECK_FORMULA(t,);        
+        to_optimize_ptr(o)->add_hard_constraint(to_expr(a), to_expr(t));
+        Z3_CATCH;
+    }
+
     unsigned Z3_API Z3_optimize_assert_soft(Z3_context c, Z3_optimize o, Z3_ast a, Z3_string weight, Z3_symbol id) {
         Z3_TRY;
         LOG_Z3_optimize_assert_soft(c, o, a, weight, id);

src/api/c++/z3++.h

@@ -2646,6 +2646,11 @@ namespace z3 {
             strm << weight;
             return handle(Z3_optimize_assert_soft(ctx(), m_opt, e, strm.str().c_str(), 0));
         }
+        void add(expr const& e, expr const& t) {
+            assert(e.is_bool());
+            Z3_optimize_assert_and_track(ctx(), m_opt, e, t);
+        }
+
         handle add(expr const& e, char const* weight) {
             assert(e.is_bool());
             return handle(Z3_optimize_assert_soft(ctx(), m_opt, e, weight, 0));

src/api/python/z3/z3.py

@@ -7336,6 +7336,36 @@ class Optimize(Z3PPObject):
         self.add(fml)
         return self
 
+    def assert_and_track(self, a, p):
+        """Assert constraint `a` and track it in the unsat core using the Boolean constant `p`.
+
+        If `p` is a string, it will be automatically converted into a Boolean constant.
+
+        >>> x = Int('x')
+        >>> p3 = Bool('p3')
+        >>> s = Optimize()
+        >>> s.set(unsat_core=True)
+        >>> s.assert_and_track(x > 0,  'p1')
+        >>> s.assert_and_track(x != 1, 'p2')
+        >>> s.assert_and_track(x < 0,  p3)
+        >>> print(s.check())
+        unsat
+        >>> c = s.unsat_core()
+        >>> len(c)
+        2
+        >>> Bool('p1') in c
+        True
+        >>> Bool('p2') in c
+        False
+        >>> p3 in c
+        True
+        """
+        if isinstance(p, str):
+            p = Bool(p, self.ctx)
+        _z3_assert(isinstance(a, BoolRef), "Boolean expression expected")
+        _z3_assert(isinstance(p, BoolRef) and is_const(p), "Boolean expression expected")
+        Z3_optimize_assert_and_track(self.ctx.ref(), self.optimize, a.as_ast(), p.as_ast())
+
     def add_soft(self, arg, weight = "1", id = None):
         """Add soft constraint with optional weight and optional identifier.
            If no weight is supplied, then the penalty for violating the soft constraint

src/api/z3_optimization.h

@@ -56,11 +56,23 @@ extern "C" {
        \brief Assert hard constraint to the optimization context.
 
        \sa Z3_optimize_assert_soft
+       \sa Z3_optimize_assert_and_track
 
        def_API('Z3_optimize_assert', VOID, (_in(CONTEXT), _in(OPTIMIZE), _in(AST)))
     */
     void Z3_API Z3_optimize_assert(Z3_context c, Z3_optimize o, Z3_ast a);
 
+
+    /**
+       \brief Assert tracked hard constraint to the optimization context.
+
+       \sa Z3_optimize_assert
+       \sa Z3_optimize_assert_soft
+
+       def_API('Z3_optimize_assert_and_track', VOID, (_in(CONTEXT), _in(OPTIMIZE), _in(AST), _in(AST)))
+    */
+    void Z3_API Z3_optimize_assert_and_track(Z3_context c, Z3_optimize o, Z3_ast a, Z3_ast t);
+
     /**
        \brief Assert soft constraint to the optimization context.
        \param c - context

src/sat/sat_drat.cpp

@@ -76,13 +76,6 @@ namespace sat {
             return;
         }
 
-#if 0
-            if (st == status::deleted) (*m_out) << "d ";
-            for (unsigned i = 0; i < n; ++i) (*m_out) << c[i] << " ";
-            (*m_out) << "0\n";
-	    return;
-#endif
-
         char buffer[10000];
         char digits[20];     // enough for storing unsigned
         char* lastd = digits + sizeof(digits);
@@ -93,7 +86,7 @@ namespace sat {
             buffer[1] = ' ';
             len = 2;
         }
-        for (unsigned i = 0; i < n && len < sizeof(buffer); ++i) {
+        for (unsigned i = 0; i < n; ++i) {
             literal lit = c[i];
             unsigned v = lit.var();            
             if (lit.sign()) buffer[len++] = '-';
@@ -104,34 +97,18 @@ namespace sat {
                 v /= 10;
                 SASSERT(d > digits);
             }
-            if (len + lastd - d < sizeof(buffer)) {
-                memcpy(buffer + len, d, lastd - d);
-                len += static_cast<unsigned>(lastd - d);
+	    SASSERT(len + lastd - d < sizeof(buffer));
+	    memcpy(buffer + len, d, lastd - d);
+	    len += static_cast<unsigned>(lastd - d);            
+	    buffer[len++] = ' ';
+	    if (len + 50 > sizeof(buffer)) {
+	        m_out->write(buffer, len);
+	        len = 0;
             }
-            else {
-                len = sizeof(buffer) + 1;
-            }
-            if (len < sizeof(buffer)) {
-                buffer[len++] = ' ';
-            }
-        }
-        
-        if (len < sizeof(buffer) + 2) {
-            buffer[len++] = '0';
-            buffer[len++] = '\n';
-        }
-	else {
-            len = sizeof(buffer) + 1;
-	}
-        if (len <= sizeof(buffer)) {
-            m_out->write(buffer, len);
-        }
-        else {
-            if (st == status::deleted) (*m_out) << "d ";
-            for (unsigned i = 0; i < n; ++i) (*m_out) << c[i] << " ";
-            (*m_out) << "0\n";
-        }
-        
+        }        
+	buffer[len++] = '0';
+	buffer[len++] = '\n';
+	m_out->write(buffer, len);               
     }
 
     void drat::bdump(unsigned n, literal const* c, status st) {