feat(UF): Store domains inside the union-find (#1119)

This patch adds the ability to store domains (as used in the bit-vector,
enum and ADT theories) directly in the union-find structure, which is
responsible for updating them automatically as class representatives are
updated.

The advantages of doing this are twofold: we no longer need to manually
keep track of substitutions (avoiding the issue that some of them are
not propagated from the union-find to the relations), and having the
domains stored in a central place rather than inside each relations
module opens the path to more communication between the domains (e.g.
accessing integer domains from the bit-vector module).

Note that constraints (notably for the bit-vector theory) are *not*
integrated into the union-find and still need to be kept track of
manually.

src/lib/reasoners/arrays_rel.ml

@@ -91,15 +91,15 @@ type t =
   }
 
 
-let empty _ =
+let empty uf =
   {gets  = G.empty;
    tbset = TBS.empty;
    split = LRset.empty;
    conseq   = LRmap.empty;
    seen  = Tmap.empty;
    new_terms = E.Set.empty;
    size_splits = Numbers.Q.one;
-  }
+  }, Uf.domains uf
 
 (*BISECT-IGNORE-BEGIN*)
 module Debug = struct
@@ -425,10 +425,10 @@ let assume env uf la =
     Conseq.fold (fun (a,ex) l ->
         ((Literal.LTerm a, ex, Th_util.Other)::l)) atoms []
   in
-  env, { Sig_rel.assume = l; remove = [] }
+  env, Uf.domains uf, { Sig_rel.assume = l; remove = [] }
 
 let query _ _ _ = None
-let add env _ _ _ = env, []
+let add env uf _ _ = env, Uf.domains uf, []
 
 let new_terms env = env.new_terms
 let instantiate ~do_syntactic_matching:_ _ env _ _ = env, []

src/lib/reasoners/bitv_rel.ml

@@ -30,6 +30,7 @@ module Ex = Explanation
 module Sy = Symbols
 module X = Shostak.Combine
 module SX = Shostak.SXH
+module MX = Shostak.MXH
 module HX = Shostak.HX
 module L = Xliteral
 
@@ -80,6 +81,8 @@ module Domain : Rel_utils.Domain with type t = Bitlist.t = struct
 
   include Bitlist
 
+  let filter_ty = is_bv_ty
+
   let fold_signed f { Bitv.value; negated } bl acc =
     let bl = if negated then lognot bl else bl in
     f value bl acc
@@ -477,22 +480,23 @@ let propagate eqs bcs dom =
 
 type t =
   { delayed : Rel_utils.Delayed.t
-  ; domain : Domains.t
   ; constraints : Constraints.t
   ; size_splits : Q.t }
 
-let empty _ =
+let empty uf =
   { delayed = Rel_utils.Delayed.create ~is_ready:X.is_constant dispatch
-  ; domain = Domains.empty
   ; constraints = Constraints.empty
-  ; size_splits = Q.one }
+  ; size_splits = Q.one },
+  Uf.GlobalDomains.add (module Domains) Domains.empty (Uf.domains uf)
 
 let assume env uf la =
+  let ds = Uf.domains uf in
+  let domain = Uf.GlobalDomains.find (module Domains) ds in
   let delayed, result = Rel_utils.Delayed.assume env.delayed uf la in
   let (domain, constraints, eqs, size_splits) =
     try
-      let ((constraints, domain), eqs, size_splits) =
-        List.fold_left (fun ((bcs, dom), eqs, ss) (a, _root, ex, orig) ->
+      let (constraints, eqs, size_splits) =
+        List.fold_left (fun (bcs, eqs, ss) (a, _root, ex, orig) ->
             let ss =
               match orig with
               | Th_util.CS (Th_bitv, n) -> Q.(ss * n)
@@ -505,9 +509,8 @@ let assume env uf la =
             in
             match a, orig with
             | L.Eq (rr, nrr), Subst when is_bv_r rr ->
-              let dom = Domains.subst ~ex rr nrr dom in
               let bcs = Constraints.subst ~ex rr nrr bcs in
-              ((bcs, dom), eqs, ss)
+              (bcs, eqs, ss)
             | L.Distinct (false, [rr; nrr]), _ when is_1bit rr ->
               (* We don't (yet) support [distinct] in general, but we must
                  support it for case splits to avoid looping.
@@ -517,10 +520,10 @@ let assume env uf la =
               let not_nrr =
                 Shostak.Bitv.is_mine (Bitv.lognot (Shostak.Bitv.embed nrr))
               in
-              ((bcs, dom), (Uf.LX.mkv_eq rr not_nrr, ex) :: eqs, ss)
-            | _ -> ((bcs, dom), eqs, ss)
+              (bcs, (Uf.LX.mkv_eq rr not_nrr, ex) :: eqs, ss)
+            | _ -> (bcs, eqs, ss)
           )
-          ((env.constraints, env.domain), [], env.size_splits)
+          (env.constraints, [], env.size_splits)
           la
       in
       let eqs, constraints, domain = propagate eqs constraints domain in
@@ -542,14 +545,17 @@ let assume env uf la =
   let result =
     { result with assume = List.rev_append assume result.assume }
   in
-  { delayed ; constraints ; domain ; size_splits }, result
+  { delayed ; constraints ; size_splits },
+  Uf.GlobalDomains.add (module Domains) domain ds,
+  result
 
 let query _ _ _ = None
 
-let case_split env _uf ~for_model =
+let case_split env uf ~for_model =
   if not for_model && Stdlib.(env.size_splits >= Options.get_max_split ()) then
     []
   else
+    let domain = Uf.GlobalDomains.find (module Domains) (Uf.domains uf) in
     (* Look for representatives with minimal, non-fully known, domain size.
 
        We first look among the constrained variables, then if there are no
@@ -567,28 +573,28 @@ let case_split env _uf ~for_model =
           Some (nunk', SX.add r xs)
         | _ -> Some (nunk, SX.singleton r)
     in
+    let f_acc' r acc =
+      let r, _ = Uf.find_r uf r in
+      List.fold_left (fun acc { Bitv.bv; _ } ->
+          match bv with
+          | Bitv.Cte _ -> acc
+          | Other r | Ext (r, _, _, _) ->
+            let bl = Domains.get r.value domain in
+            f_acc r.value bl acc
+        ) acc (Shostak.Bitv.embed r)
+    in
     let _, candidates =
-      match
-        Constraints.fold_args (fun r acc ->
-            List.fold_left (fun acc { Bitv.bv; _ } ->
-                match bv with
-                | Bitv.Cte _ -> acc
-                | Other r | Ext (r, _, _, _) ->
-                  let bl = Domains.get r.value env.domain in
-                  f_acc r.value bl acc
-              ) acc (Shostak.Bitv.embed r)
-          ) env.constraints None
-      with
+      match Constraints.fold_args f_acc' env.constraints None with
       | Some (nunk, xs) -> nunk, xs
-      | None ->
-        match Domains.fold_leaves f_acc env.domain None with
+      | _ ->
+        match Domains.fold_leaves f_acc domain None with
         | Some (nunk, xs) -> nunk, xs
         | None -> 0, SX.empty
     in
     (* For now, just pick a value for the most significant bit. *)
     match SX.choose candidates with
     | r ->
-      let bl = Domains.get r env.domain in
+      let bl = Domains.get r domain in
       let w = Bitlist.width bl in
       let unknown = Z.extract (Z.lognot @@ Bitlist.bits_known bl) 0 w in
       let bitidx = Z.numbits unknown  - 1 in
@@ -608,19 +614,16 @@ let case_split env _uf ~for_model =
 let add env uf r t =
   let delayed, eqs = Rel_utils.Delayed.add env.delayed uf r t in
   let env, eqs =
-    match X.type_info r with
-    | Tbitv _ -> (
-        try
-          let dom = Domains.add r env.domain in
-          let bcs = extract_constraints env.constraints uf r t in
-          let eqs, bcs, dom = propagate eqs bcs dom in
-          { env with constraints = bcs ; domain = dom }, eqs
-        with Domains.Inconsistent ex ->
-          raise @@ Ex.Inconsistent (ex, Uf.cl_extract uf)
-      )
-    | _ -> env, eqs
+    if is_bv_r r then
+      try
+        let constraints = extract_constraints env.constraints uf r t in
+        { env with constraints }, eqs
+      with Domains.Inconsistent ex ->
+        raise @@ Ex.Inconsistent (ex, Uf.cl_extract uf)
+    else
+      env, eqs
   in
-  { env with delayed }, eqs
+  { env with delayed }, Uf.domains uf, eqs
 
 let optimizing_objective _env _uf _o = None
 

src/lib/reasoners/ccx.ml

@@ -105,11 +105,15 @@ module Main : S = struct
 
   type r = Shostak.Combine.r
 
-  let empty = {
-    use = Use.empty ;
-    uf = Uf.empty;
-    relation = Rel.empty [];
-  }
+  let empty =
+    let uf = Uf.empty in
+    let relation, domains = Rel.empty uf in
+    let uf = Uf.set_domains uf domains in
+    {
+      use = Use.empty;
+      uf;
+      relation;
+    }
 
   let empty_facts () =
     Sig_rel.{ equas   = Queue.create ();
@@ -462,7 +466,8 @@ module Main : S = struct
   let replay_atom env sa =
     Options.exec_thread_yield ();
     let sa = make_unique sa in
-    let relation, result = Rel.assume env.relation env.uf sa in
+    let relation, domains, result = Rel.assume env.relation env.uf sa in
+    let env = { env with uf = Uf.set_domains env.uf domains } in
     let env = { env with relation = relation } in
     let env = clean_use env result.remove in
     env, result.assume
@@ -489,7 +494,8 @@ module Main : S = struct
       let nuse = Use.up_add env.use t rt lvs in
 
       (* If finitetest is used we add the term to the relation *)
-      let rel, eqs = Rel.add env.relation nuf rt t in
+      let rel, doms, eqs = Rel.add env.relation nuf rt t in
+      let nuf = Uf.set_domains nuf doms in
       Debug.rel_add_cst t eqs;
       (* We add terms made from relations as fact *)
       List.iter (fun (a,ex) -> add_fact facts (LSem a, ex, Th_util.Other)) eqs;