refactor: new API for combination, with theories claiming terms

interface variables are terms claimed by >= 2 theories. Theories now
have a unique ID attributed at their creation.

src/base/Sidekick_base.ml

@@ -33,6 +33,7 @@ module LRA_term = LRA_term
 module Th_data = Th_data
 module Th_bool = Th_bool
 module Th_lra = Th_lra
+module Th_uf = Th_uf
 
 let k_th_bool_config = Th_bool.k_config
 let th_bool = Th_bool.theory
@@ -40,3 +41,4 @@ let th_bool_dyn : Solver.theory = Th_bool.theory_dyn
 let th_bool_static : Solver.theory = Th_bool.theory_static
 let th_data : Solver.theory = Th_data.theory
 let th_lra : Solver.theory = Th_lra.theory
+let th_uf : Solver.theory = Th_uf.theory

src/base/th_uf.ml

@@ -0,0 +1,24 @@
+(** Theory of uninterpreted functions *)
+
+open Sidekick_core
+open Sidekick_smt_solver
+
+open struct
+  module SI = Solver_internal
+
+  let on_is_subterm ~th_id (solver : SI.t) (_, _, t) : _ list =
+    let f, args = Term.unfold_app t in
+    (match Term.view f, args with
+    | Term.E_const { Const.c_view = Uconst.Uconst _; _ }, _ :: _ ->
+      SI.claim_term solver ~th_id t
+    | _ -> ());
+
+    []
+end
+
+let theory : Theory.t =
+  Theory.make ~name:"uf"
+    ~create_and_setup:(fun ~id:th_id solver ->
+      SI.on_cc_is_subterm solver (on_is_subterm ~th_id solver);
+      ())
+    ()

src/main/main.ml

@@ -177,7 +177,7 @@ let main_smt ~config () : _ result =
       Log.debugf 1 (fun k ->
           k "(@[main.th-bool.pick@ %S@])"
             (Sidekick_smt_solver.Theory.name th_bool));
-      [ th_bool; Process.th_data; Process.th_lra ]
+      [ th_bool; Process.th_uf; Process.th_data; Process.th_lra ]
     in
     Process.Solver.create_default ~proof ~theories tst
   in

src/smt/Sidekick_smt_solver.ml

@@ -13,6 +13,7 @@ module Registry = Registry
 module Solver_internal = Solver_internal
 module Solver = Solver
 module Theory = Theory
+module Theory_id = Theory_id
 
 type theory = Theory.t
 type solver = Solver.t

src/smt/solver.ml

@@ -40,6 +40,7 @@ type t = {
   mutable last_res: res option;
   stat: Stat.t;
   proof: P.t;
+  theory_id_gen: Theory_id.state;
   n_clause_input: int Stat.counter;
   n_clause_internal: int Stat.counter;
   n_solve: int Stat.counter; (* config: Config.t *)
@@ -53,8 +54,11 @@ let mk_theory = Theory.make
 
 let add_theory_p (type a) (self : t) (th : a Theory.p) : a =
   let (module Th) = th in
-  Log.debugf 2 (fun k -> k "(@[smt-solver.add-theory@ :name %S@])" Th.name);
-  let st = Th.create_and_setup self.si in
+  let th_id = Theory_id.fresh self.theory_id_gen in
+  Log.debugf 2 (fun k ->
+      k "(@[smt-solver.add-theory@ :id %a@ :name %S@])" Theory_id.pp th_id
+        Th.name);
+  let st = Th.create_and_setup ~id:th_id self.si in
   (* add push/pop to the internal solver *)
   Solver_internal.add_theory_state self.si ~st ~push_level:Th.push_level
     ~pop_levels:Th.pop_levels;
@@ -77,6 +81,7 @@ let create arg ?(stat = Stat.global) ?size ~proof ~theories tst () : t =
       last_res = None;
       solver = Sat_solver.create ~proof ?size ~stat (SI.to_sat_plugin si);
       stat;
+      theory_id_gen = Theory_id.create ();
       n_clause_input = Stat.mk_int stat "smt.solver.add-clause.input";
       n_clause_internal = Stat.mk_int stat "smt.solver.add-clause.internal";
       n_solve = Stat.mk_int stat "smt.solver.solve";

src/smt/solver.mli

@@ -18,7 +18,7 @@ type theory = Theory.t
 
 val mk_theory :
   name:string ->
-  create_and_setup:(Solver_internal.t -> 'th) ->
+  create_and_setup:(id:Theory_id.t -> Solver_internal.t -> 'th) ->
   ?push_level:('th -> unit) ->
   ?pop_levels:('th -> int -> unit) ->
   unit ->

src/smt/solver_internal.ml

@@ -21,7 +21,6 @@ module type PREPROCESS_ACTS = sig
   val mk_lit : ?sign:bool -> term -> lit
   val add_clause : lit list -> step_id -> unit
   val add_lit : ?default_pol:bool -> lit -> unit
-  val add_term_needing_combination : term -> unit
 end
 
 type preprocess_actions = (module PREPROCESS_ACTS)
@@ -83,9 +82,7 @@ let[@inline] has_delayed_actions self =
   not (Queue.is_empty self.delayed_actions)
 
 let on_preprocess self f = self.preprocess <- f :: self.preprocess
-
-let add_term_needing_combination self t =
-  Th_combination.add_term_needing_combination self.th_comb t
+let claim_term self ~th_id t = Th_combination.claim_term self.th_comb ~th_id t
 
 let on_model ?ask ?complete self =
   Option.iter (fun f -> self.model_ask <- f :: self.model_ask) ask;
@@ -130,9 +127,6 @@ let preprocess_term_ (self : t) (t0 : term) : unit =
     let mk_lit ?sign t : Lit.t = Lit.atom ?sign self.tst t
     let add_lit ?default_pol lit : unit = delayed_add_lit self ?default_pol lit
     let add_clause c pr : unit = delayed_add_clause self ~keep:true c pr
-
-    let add_term_needing_combination t =
-      Th_combination.add_term_needing_combination self.th_comb t
   end in
   let acts = (module A : PREPROCESS_ACTS) in
 

src/smt/solver_internal.mli

@@ -73,10 +73,6 @@ module type PREPROCESS_ACTS = sig
 
   val add_lit : ?default_pol:bool -> lit -> unit
   (** Ensure the literal will be decided/handled by the SAT solver. *)
-
-  val add_term_needing_combination : term -> unit
-  (** Declare this term as being a foreign variable in the theory,
-    which means it needs to go through theory combination. *)
 end
 
 type preprocess_actions = (module PREPROCESS_ACTS)
@@ -102,9 +98,9 @@ val preprocess_clause_array : t -> lit array -> step_id -> lit array * step_id
 val simplify_and_preproc_lit : t -> lit -> lit * step_id option
 (** Simplify literal then preprocess it *)
 
-val add_term_needing_combination : t -> term -> unit
-(** Declare this term as being a foreign variable in the theory,
-    which means it needs to go through theory combination. *)
+val claim_term : t -> th_id:Theory_id.t -> term -> unit
+(** Claim a term, for a theory that might decide or merge it with another
+    term. This is useful for theory combination. *)
 
 (** {3 hooks for the theory} *)
 

src/smt/th_combination.ml

@@ -6,6 +6,7 @@ type t = {
   processed: T.Set.t T.Tbl.t;  (** type -> set of terms *)
   unprocessed: T.t Vec.t;
   new_lits: Lit.t Vec.t;
+  claims: Theory_id.Set.t T.Tbl.t;  (** term -> theories claiming it *)
   n_terms: int Stat.counter;
   n_lits: int Stat.counter;
 }
@@ -15,6 +16,7 @@ let create ?(stat = Stat.global) tst : t =
     tst;
     processed = T.Tbl.create 8;
     unprocessed = Vec.create ();
+    claims = T.Tbl.create 8;
     new_lits = Vec.create ();
     n_terms = Stat.mk_int stat "smt.thcomb.terms";
     n_lits = Stat.mk_int stat "smt.thcomb.intf-lits";
@@ -28,10 +30,31 @@ let processed_ (self : t) t : bool =
 
 let add_term_needing_combination (self : t) (t : T.t) : unit =
   if not (processed_ self t) then (
-    Log.debugf 50 (fun k -> k "(@[th.comb.add-term-needing-comb@ %a@])" T.pp t);
+    Log.debugf 50 (fun k ->
+        k "(@[th.comb.add-term-needing-comb@ `%a`@ :ty `%a`@])" T.pp t T.pp
+          (T.ty t));
     Vec.push self.unprocessed t
   )
 
+let claim_term (self : t) ~th_id (t : T.t) : unit =
+  (* booleans don't need theory combination *)
+  if T.is_bool (T.ty t) then
+    ()
+  else (
+    Log.debugf 50 (fun k ->
+        k "(@[th.comb.claim :th-id %a@ `%a`@])" Theory_id.pp th_id T.pp t);
+    let set =
+      try T.Tbl.find self.claims t with Not_found -> Theory_id.Set.empty
+    in
+    let set' = Theory_id.Set.add th_id set in
+    if Theory_id.Set.(not (equal set set')) then (
+      T.Tbl.replace self.claims t set';
+      (* first time we have 2 theories, means we need combination *)
+      if Theory_id.Set.cardinal set' = 2 then
+        add_term_needing_combination self t
+    )
+  )
+
 let pop_new_lits (self : t) : Lit.t list =
   (* first, process new terms, if any *)
   while not (Vec.is_empty self.unprocessed) do

src/smt/th_combination.mli

@@ -6,11 +6,17 @@ type t
 
 val create : ?stat:Stat.t -> Term.store -> t
 
-val add_term_needing_combination : t -> Term.t -> unit
-(** [add_term_needing_combination self t] means that [t] occurs as a foreign
-  variable in another term, so it is important that its theory, and the
-  theory in which it occurs, agree on it being equal to other
-  foreign terms. *)
+val claim_term : t -> th_id:Theory_id.t -> Term.t -> unit
+(** [claim_term self ~th_id t] means that theory with ID [th_id]
+    claims the term [t].
+
+    This means it might assert [t = u] or [t ≠ u] for some other term [u],
+    or it might assign a value to [t] in the model in case of a SAT answer.
+    That means it has to agree with other theories on what [t] is equal to.
+
+    If a term is claimed by several theories, it will be eligible for theory
+    combination.
+*)
 
 val pop_new_lits : t -> Lit.t list
 (** Get the new literals that the solver needs to decide, so that the

src/smt/theory.ml

@@ -19,7 +19,7 @@ module type S = sig
   type t
 
   val name : string
-  val create_and_setup : Solver_internal.t -> t
+  val create_and_setup : id:Theory_id.t -> Solver_internal.t -> t
   val push_level : t -> unit
   val pop_levels : t -> int -> unit
 end

src/smt/theory_id.ml

@@ -0,0 +1,12 @@
+include CCInt
+
+type state = int ref
+
+let create () = ref 1
+
+let fresh (self : state) =
+  let n = !self in
+  incr self;
+  n
+
+module Set = Util.Int_set

src/smt/theory_id.mli

@@ -0,0 +1,10 @@
+type t = private int
+
+include Sidekick_sigs.EQ_ORD_HASH_PRINT with type t := t
+
+type state
+
+val create : unit -> state
+val fresh : state -> t
+
+module Set : CCSet.S with type elt = t

src/smtlib/Process.ml

@@ -60,7 +60,7 @@ module Check_cc = struct
 
   let theory =
     Solver.mk_theory ~name:"cc-check"
-      ~create_and_setup:(fun si ->
+      ~create_and_setup:(fun ~id:_ si ->
         let n_calls = Stat.mk_int (SI.stats si) "check-cc.call" in
         SI.on_cc_conflict si (fun { cc; th; c } ->
             if not th then (
@@ -335,12 +335,11 @@ let process_stmt ?gc ?restarts ?(pp_cnf = false) ?proof_file ?pp_model
   | Statement.Stmt_data _ -> E.return ()
   | Statement.Stmt_define _ -> Error.errorf "cannot deal with definitions yet"
 
-module Th_data = Sidekick_base.Th_data
-module Th_bool = Sidekick_base.Th_bool
-module Th_lra = Sidekick_base.Th_lra
+open Sidekick_base
 
 let th_bool = Th_bool.theory
 let th_bool_dyn : Solver.theory = Th_bool.theory_dyn
 let th_bool_static : Solver.theory = Th_bool.theory_static
 let th_data : Solver.theory = Th_data.theory
 let th_lra : Solver.theory = Th_lra.theory
+let th_uf = Th_uf.theory

src/smtlib/Process.mli

@@ -8,6 +8,7 @@ val th_bool_static : Solver.theory
 val th_bool : Config.t -> Solver.theory
 val th_data : Solver.theory
 val th_lra : Solver.theory
+val th_uf : Solver.theory
 
 type 'a or_error = ('a, string) CCResult.t
 

src/th-bool-dyn/Sidekick_th_bool_dyn.ml

@@ -352,7 +352,7 @@ end = struct
   let final_check (self : state) solver acts (lits : Lit.t Iter.t) =
     check_ ~final:true self solver acts lits
 
-  let create_and_setup (solver : SI.t) : state =
+  let create_and_setup ~id:_ (solver : SI.t) : state =
     let tst = SI.tst solver in
     let stat = SI.stats solver in
     let self =

src/th-bool-static/Sidekick_th_bool_static.ml

@@ -287,7 +287,7 @@ end = struct
     | B_atom _ -> ());
     ()
 
-  let create_and_setup si =
+  let create_and_setup ~id:_ si =
     Log.debug 2 "(th-bool.setup)";
     let st = create ~stat:(SI.stats si) (SI.tst si) in
     SI.add_simplifier si (simplify st);

src/th-cstor/Sidekick_th_cstor.ml

@@ -80,7 +80,7 @@ end = struct
   let pop_levels ((module P) : t) n = P.pop_levels n
   let n_levels ((module P) : t) = P.n_levels ()
 
-  let create_and_setup (si : SI.t) : t =
+  let create_and_setup ~id:_ (si : SI.t) : t =
     Log.debug 1 "(setup :th-cstor)";
     let self = ST.create_and_setup ~size:32 (SI.cc si) in
     self

src/th-data/Sidekick_th_data.ml

@@ -296,6 +296,7 @@ end = struct
   module N_tbl = Backtrackable_tbl.Make (E_node)
 
   type t = {
+    th_id: Sidekick_smt_solver.Theory_id.t;
     tst: Term.store;
     proof: Proof_trace.t;
     cstors: ST_cstors.t; (* repr -> cstor for the class *)
@@ -463,6 +464,10 @@ end = struct
         [])
     | T_cstor _ | T_other _ -> []
 
+  let on_is_subterm (self : t) (si : SI.t) (_cc, _repr, t) : _ list =
+    if is_data_ty (Term.ty t) then SI.claim_term si ~th_id:self.th_id t;
+    []
+
   let cstors_of_ty (ty : ty) : A.Cstor.t list =
     match A.as_datatype ty with
     | Ty_data { cstors } -> cstors
@@ -783,9 +788,10 @@ end = struct
         Some (c, args))
     | None -> None
 
-  let create_and_setup (solver : SI.t) : t =
+  let create_and_setup ~id:th_id (solver : SI.t) : t =
     let self =
       {
+        th_id;
         tst = SI.tst solver;
         proof = SI.proof solver;
         cstors = ST_cstors.create_and_setup ~size:32 (SI.cc solver);
@@ -801,6 +807,7 @@ end = struct
     Log.debugf 1 (fun k -> k "(setup :%s)" name);
     SI.on_preprocess solver (preprocess self);
     SI.on_cc_new_term solver (on_new_term self);
+    SI.on_cc_is_subterm solver (on_is_subterm self solver);
     (* note: this needs to happen before we modify the plugin data *)
     SI.on_cc_pre_merge solver (on_pre_merge self);
     SI.on_partial_check solver (on_partial_check self);

src/th-lra/sidekick_th_lra.ml

@@ -124,6 +124,7 @@ module Make (A : ARG) = (* : S with module A = A *) struct
   module ST_exprs = Sidekick_cc.Plugin.Make (Monoid_exprs)
 
   type state = {
+    th_id: Sidekick_smt_solver.Theory_id.t;
     tst: Term.store;
     proof: Proof_trace.t;
     gensym: Gensym.t;
@@ -142,11 +143,12 @@ module Make (A : ARG) = (* : S with module A = A *) struct
     n_conflict: int Stat.counter;
   }
 
-  let create (si : SI.t) : state =
+  let create ~th_id (si : SI.t) : state =
     let stat = SI.stats si in
     let proof = SI.proof si in
     let tst = SI.tst si in
     {
+      th_id;
       tst;
       proof;
       in_model = Term.Tbl.create 8;
@@ -272,11 +274,6 @@ module Make (A : ARG) = (* : S with module A = A *) struct
     | Geq -> S_op.Geq
     | Gt -> S_op.Gt
 
-  (* add [t] to the theory combination system if it's not just a constant
-     of type Real *)
-  let add_lra_var_to_th_combination (si : SI.t) (t : term) : unit =
-    if not (Term.is_const t) then SI.add_term_needing_combination si t
-
   (* TODO: refactor that and {!var_encoding_comb} *)
   (* turn a linear expression into a single constant and a coeff.
      This might define a side variable in the simplex. *)
@@ -302,20 +299,13 @@ module Make (A : ARG) = (* : S with module A = A *) struct
         proxy, A.Q.one)
 
   (* look for subterms of type Real, for they will need theory combination *)
-  let on_subterm (_self : state) (si : SI.t) (t : Term.t) : unit =
+  let on_subterm (self : state) (si : SI.t) (t : Term.t) : unit =
     Log.debugf 50 (fun k -> k "(@[lra.cc-on-subterm@ %a@])" Term.pp_debug t);
     match A.view_as_lra t with
-    | LRA_other _ when not (A.has_ty_real t) ->
-      (* for a non-LRA term [f args], if any of [args] is in LRA,
-         it needs theory combination *)
-      let _, args = Term.unfold_app t in
-      List.iter
-        (fun arg ->
-          if A.has_ty_real arg then SI.add_term_needing_combination si arg)
-        args
+    | LRA_other _ when not (A.has_ty_real t) -> ()
     | LRA_pred _ | LRA_const _ -> ()
     | LRA_op _ | LRA_other _ | LRA_mult _ ->
-      SI.add_term_needing_combination si t
+      SI.claim_term si ~th_id:self.th_id t
 
   (* preprocess linear expressions away *)
   let preproc_lra (self : state) si (module PA : SI.PREPROCESS_ACTS)
@@ -374,11 +364,7 @@ module Make (A : ARG) = (* : S with module A = A *) struct
       (* obtain a single variable for the linear combination *)
       let v, c_v = le_comb_to_singleton_ self le_comb in
       declare_term_to_cc ~sub:false v;
-      LE_.Comb.iter
-        (fun v _ ->
-          declare_term_to_cc ~sub:true v;
-          add_lra_var_to_th_combination si v)
-        le_comb;
+      LE_.Comb.iter (fun v _ -> declare_term_to_cc ~sub:true v) le_comb;
 
       (* turn into simplex constraint. For example,
          [c . v <= const] becomes a direct simplex constraint [v <= const/c]
@@ -714,9 +700,9 @@ module Make (A : ARG) = (* : S with module A = A *) struct
 
   let k_state = SMT.Registry.create_key ()
 
-  let create_and_setup si =
+  let create_and_setup ~id si =
     Log.debug 2 "(th-lra.setup)";
-    let st = create si in
+    let st = create ~th_id:id si in
     SMT.Registry.set (SI.registry si) k_state st;
     SI.add_simplifier si (simplify st);
     SI.on_preprocess si (preproc_lra st);