feat: add push/pop for assumptions; add check_sat_propagation_only

this gives the user the possibility of controlling search by mostly
using assumptions (and the underlying assumption stack) as well as
boolean+theory propagation. The result is **not** complete, but can
still help as a beefed-up contextual simplifier.

src/core/Sidekick_core.ml

@@ -1180,6 +1180,44 @@ module type SOLVER = sig
         must stop (returning [Unknown]), [false] if solving can proceed.
       @param on_exit functions to be run before this returns *)
 
+  val push_assumption : t -> lit -> unit
+  (** Pushes an assumption onto the assumption stack. It will remain
+      there until it's pop'd by {!pop_assumptions}. *)
+
+  val pop_assumptions : t -> int -> unit
+  (** [pop_assumptions solver n] removes [n] assumptions from the stack.
+      It removes the assumptions that were the most
+      recently added via {!push_assumptions}.
+      Note that {!check_sat_propagations_only} can call this if it meets
+      a conflict. *)
+
+  type propagation_result =
+    | PR_sat
+    | PR_conflict of {
+        backtracked: int;
+      }
+    | PR_unsat of {
+        unsat_core: unit -> lit Iter.t;
+      }
+
+  val check_sat_propagations_only :
+    assumptions:lit list ->
+    t ->
+    propagation_result
+  (** [check_sat_propagations_only solver] uses assumptions (including
+      the [assumptions] parameter, and atoms previously added via {!push_assumptions})
+      and boolean+theory propagation to quickly assess satisfiability.
+      It is not complete; calling {!solve} is required to get an accurate
+      result.
+      @returns one of:
+
+      - [PR_sat] if the current state seems satisfiable
+      - [PR_conflict {backtracked=n}] if a conflict was found and resolved,
+      leading to backtracking [n] levels of assumptions
+      - [PR_unsat …] if the assumptions were found to be unsatisfiable, with
+        the given core.
+  *)
+
   (* TODO: allow on_progress to return a bool to know whether to stop? *)
 
   val pp_stats : t CCFormat.printer

src/sat/Solver.ml

@@ -2020,46 +2020,57 @@ module Make(Plugin : PLUGIN)
     Log.debugf 3 (fun k->k "(@[sat.gc.done :collected %d@])" n_collected);
     ()
 
-  (* Decide on a new literal, and enqueue it into the trail *)
+  (* Decide on a new literal, and enqueue it into the trail.
-  let rec pick_branch_aux self atom : unit =
+     Return [true] if a decision was made.
+     @param full if true, do decisions;
+     if false, only pick from [self.next_decisions]
+     and [self.assumptions] *)
+  let pick_branch_lit ~full self : bool =
+    let rec pick_lit () =
+      match self.next_decisions with
+      | atom :: tl ->
+        self.next_decisions <- tl;
+        pick_with_given_atom atom
+      | [] when decision_level self < AVec.size self.assumptions ->
+        (* use an assumption *)
+        let a = AVec.get self.assumptions (decision_level self) in
+        if Atom.is_true self.store a then (
+          new_decision_level self; (* pseudo decision level, [a] is already true *)
+          pick_lit ()
+        ) else if Atom.is_false self.store a then (
+          (* root conflict, find unsat core *)
+          let core = analyze_final self a in
+          raise (E_unsat (US_local {first=a; core}))
+        ) else (
+          pick_with_given_atom a
+        )
+      | [] when not full -> false
+      | [] ->
+        begin match H.remove_min self.order with
+          | v ->
+            pick_with_given_atom
+              (if Var.default_pol self.store v then Atom.pa v else Atom.na v)
+          | exception Not_found ->
+            false
+        end
+
+    (* pick a decision, trying [atom] first if it's not assigned yet. *)
+    and pick_with_given_atom (atom:atom) : bool =
       let v = Atom.var atom in
       if Var.level self.store v >= 0 then (
         assert (Atom.is_true self.store (Atom.pa v) ||
                 Atom.is_true self.store (Atom.na v));
-      pick_branch_lit self
+        pick_lit ()
       ) else (
         new_decision_level self;
         let current_level = decision_level self in
         enqueue_bool self atom ~level:current_level Decision;
         Stat.incr self.n_decisions;
         (match self.on_decision with Some f -> f self (Atom.lit self.store atom) | None -> ());
+        true
       )
-
+    in
-  and pick_branch_lit self : unit =
+    pick_lit()
-    match self.next_decisions with
-    | atom :: tl ->
-      self.next_decisions <- tl;
-      pick_branch_aux self atom
-    | [] when decision_level self < AVec.size self.assumptions ->
-      (* use an assumption *)
-      let a = AVec.get self.assumptions (decision_level self) in
-      if Atom.is_true self.store a then (
-        new_decision_level self; (* pseudo decision level, [a] is already true *)
-        pick_branch_lit self
-      ) else if Atom.is_false self.store a then (
-        (* root conflict, find unsat core *)
-        let core = analyze_final self a in
-        raise (E_unsat (US_local {first=a; core}))
-      ) else (
-        pick_branch_aux self a
-      )
-    | [] ->
-      begin match H.remove_min self.order with
-        | v ->
-          pick_branch_aux self
-            (if Var.default_pol self.store v then Atom.pa v else Atom.na v)
-        | exception Not_found -> raise_notrace E_sat
-      end
 
   (* do some amount of search, until the number of conflicts or clause learnt
      reaches the given parameters *)
@@ -2105,7 +2116,8 @@ module Make(Plugin : PLUGIN)
           reduce_clause_db st;
         );
 
-        pick_branch_lit st
+        let decided = pick_branch_lit ~full:true st in
+        if not decided then raise_notrace E_sat
     done
 
   let eval_level (self:t) (a:atom) =
@@ -2298,6 +2310,49 @@ module Make(Plugin : PLUGIN)
     end in
     (module M)
 
+  type propagation_result =
+    | PR_sat
+    | PR_conflict of { backtracked: int }
+    | PR_unsat of (lit,clause,proof_step) Solver_intf.unsat_state
+
+  (* decide on assumptions, and do propagations, but no other kind of decision *)
+  let propagate_under_assumptions (self:t) : propagation_result =
+    let result = ref PR_sat in
+    try
+      while true do
+        match propagate self with
+        | Some confl ->
+          (* When the theory has raised Unsat, add_boolean_conflict
+             might 'forget' the initial conflict clause, and only add the
+             analyzed backtrack clause. So in those case, we use add_clause
+             to make sure the initial conflict clause is also added. *)
+          if Clause.attached self.store confl then (
+            add_boolean_conflict self confl
+          ) else (
+            add_clause_ ~pool:self.clauses_learnt self confl
+          );
+          Stat.incr self.n_conflicts;
+
+          (* see by how much we backtracked the decision trail *)
+          let new_lvl = decision_level self in
+          assert (new_lvl < AVec.size self.assumptions);
+          let backtracked = AVec.size self.assumptions - new_lvl in
+          result := PR_conflict {backtracked};
+          AVec.shrink self.assumptions new_lvl;
+          raise_notrace Exit
+
+        | None -> (* No Conflict *)
+
+          let decided = pick_branch_lit self ~full:false in
+          if not decided then (
+            result := PR_sat;
+            raise Exit
+          )
+      done;
+      assert false
+    with Exit ->
+      !result
+
   let add_clause_atoms_ self ~pool ~removable (c:atom array) (pr:proof_step) : unit =
     try
       let c = Clause.make_a self.store ~removable c pr in
@@ -2342,26 +2397,55 @@ module Make(Plugin : PLUGIN)
     Vec.push self.clause_pools p;
     Clause_pool_id._unsafe_of_int id
 
-  let solve
+  (* run [f()] with additional assumptions *)
-      ?(on_progress=fun _ -> ())
+  let with_local_assumptions_ (self:t) (assumptions:lit list) f =
-      ?(assumptions=[]) (self:t) : res =
+    let old_assm_lvl = AVec.size self.assumptions in
-    cancel_until self 0;
-    AVec.clear self.assumptions;
     List.iter
       (fun lit ->
          let a = make_atom_ self lit in
          AVec.push self.assumptions a)
       assumptions;
+    try
+      let x = f() in
+      AVec.shrink self.assumptions old_assm_lvl;
+      x
+    with e ->
+      AVec.shrink self.assumptions old_assm_lvl;
+      raise e
+
+  let solve
+      ?(on_progress=fun _ -> ())
+      ?(assumptions=[]) (self:t) : res =
+    cancel_until self 0; (* make sure we are at level 0 *)
+    with_local_assumptions_ self assumptions @@ fun () ->
     try
       solve_ ~on_progress self;
       Sat (mk_sat self)
-    with E_unsat us ->
+    with
-      (* FIXME
+    | E_unsat us ->
-      (* emit empty clause *)
-      Proof.with_proof self.proof (Proof.emit_redundant_clause Iter.empty);
-         *)
       Unsat (mk_unsat self us)
 
+  let push_assumption (self:t) (lit:lit) : unit =
+    let a = make_atom_ self lit in
+    AVec.push self.assumptions a
+
+  let pop_assumptions self n : unit =
+    let n_ass = AVec.size self.assumptions in
+    assert (n <= n_ass);
+    AVec.shrink self.assumptions (n_ass - n)
+
+  let check_sat_propagations_only ?(assumptions=[]) (self:t) : propagation_result =
+    cancel_until self 0;
+    with_local_assumptions_ self assumptions @@ fun () ->
+    try
+      check_unsat_ self;
+      flush_clauses self; (* add initial clauses *)
+      propagate_under_assumptions self
+    with
+    | E_unsat us ->
+      let us = mk_unsat self us in
+      PR_unsat us
+
   let true_at_level0 (self:t) (lit:lit) : bool =
     match find_atom_ self lit with
     | None -> false

src/sat/Solver_intf.ml

@@ -288,6 +288,8 @@ module type S = sig
      and type t = proof
   (** A module to manipulate proofs. *)
 
+  (** {2 Main Solver Type} *)
+
   type t = solver
   (** Main solver type, containing all state for solving. *)
 
@@ -375,7 +377,7 @@ module type S = sig
   val add_clause_a_in_pool : t -> pool:clause_pool_id -> lit array -> proof_step -> unit
   (** Like {!add_clause_a} but using a specific clause pool *)
 
-  (* TODO: API to push/pop/clear assumptions from an inner vector *)
+(** {2 Solving} *)
 
   val solve :
     ?on_progress:(unit -> unit) ->
@@ -388,6 +390,8 @@ module type S = sig
       @param on_progress regularly called during solving
   *)
 
+  (** {2 Evaluating and adding literals} *)
+
   val add_lit : t -> ?default_pol:bool -> lit -> unit
   (** Ensure the SAT solver handles this particular literal, ie add
       a boolean variable for it if it's not already there. *)
@@ -402,5 +406,31 @@ module type S = sig
 
   val eval_lit : t -> lit -> lbool
   (** Evaluate atom in current state *)
+
+  (** {2 Assumption stack} *)
+
+  val push_assumption : t -> lit -> unit
+  (** Pushes an assumption onto the assumption stack. It will remain
+      there until it's pop'd by {!pop_assumptions}. *)
+
+  val pop_assumptions : t -> int -> unit
+  (** [pop_assumptions solver n] removes [n] assumptions from the stack.
+      It removes the assumptions that were the most
+      recently added via {!push_assumptions}. *)
+
+  (** Result returned by {!check_sat_propagations_only} *)
+  type propagation_result =
+    | PR_sat
+    | PR_conflict of { backtracked: int }
+    | PR_unsat of (lit,clause,proof_step) unsat_state
+
+  val check_sat_propagations_only : ?assumptions:lit list -> t -> propagation_result
+  (** [check_sat_propagations_only solver] uses the added clauses
+      and local assumptions (using {!push_assumptions} and [assumptions])
+      to quickly assess whether the context is satisfiable.
+      It is not complete; calling {!solve} is required to get an accurate
+      result.
+      @returns either [Ok()] if propagation yielded no conflict, or [Error c]
+        if a conflict clause [c] was found. *)
 end
 

src/smt-solver/Sidekick_smt_solver.ml

@@ -898,7 +898,6 @@ module Make(A : ARG)
             (Util.pp_iter ppc) (CC.all_classes @@ Solver_internal.cc self.si));
 
       let _lits f = SAT.iter_trail f in
-      (* TODO: theory combination *)
       let m = mk_model self _lits in
       (* TODO: check model *)
       let _ = check in
@@ -913,6 +912,30 @@ module Make(A : ARG)
       Unsat {unsat_core; unsat_proof_step}
     | exception Should_stop -> Unknown Unknown.U_asked_to_stop
 
+  let push_assumption self a = Sat_solver.push_assumption self.solver a
+  let pop_assumptions self n = Sat_solver.pop_assumptions self.solver n
+
+  type propagation_result =
+    | PR_sat
+    | PR_conflict of {
+        backtracked: int;
+      }
+    | PR_unsat of {
+        unsat_core: unit -> lit Iter.t;
+      }
+
+  let check_sat_propagations_only ~assumptions self : propagation_result =
+    match
+      Sat_solver.check_sat_propagations_only ~assumptions self.solver
+    with
+    | Sat_solver.PR_sat -> PR_sat
+    | Sat_solver.PR_conflict {backtracked} ->
+      PR_conflict {backtracked}
+
+    | Sat_solver.PR_unsat (module UNSAT) ->
+      let unsat_core () = UNSAT.unsat_assumptions () in
+      PR_unsat {unsat_core}
+
   let mk_theory (type st)
       ~name ~create_and_setup
       ?(push_level=fun _ -> ()) ?(pop_levels=fun _ _ -> ())