feat(Theory): Disj.forall_cof (#6)

src/DisjointSet.ml

@@ -9,6 +9,10 @@ sig
   val test_and_union : key -> key -> t -> bool * t
 
   val merge : t -> t -> t
+
+  type finger
+  val finger : key -> t -> finger
+  val test_finger : key -> finger -> bool
 end
 
 module Make (O : Map.OrderedType) : S with type key = O.t =
@@ -68,4 +72,22 @@ struct
     snd @@ test_and_union x y h
 
   let merge h1 h2 = M.fold union h2.parent h1
+
+  type finger =
+    { root : key
+    ; parent : key M.t
+    }
+
+  let finger k t =
+    { root = root k t
+    ; parent = t.parent
+    }
+
+  let test_finger k {root; parent} =
+    let rec loop x =
+      match M.find_opt x parent with
+      | Some x -> (loop[@tailcall]) x
+      | None -> O.compare x root = 0
+    in
+    loop k
 end

src/DisjointSet.mli

@@ -9,6 +9,10 @@ sig
   val test_and_union : key -> key -> t -> bool * t
 
   val merge : t -> t -> t
+
+  type finger
+  val finger : key -> t -> finger
+  val test_finger : key -> finger -> bool
 end
 
 module Make (O : Map.OrderedType) : S with type key = O.t

src/Theory.ml

@@ -38,6 +38,7 @@ sig
     val envelope_alg : Alg.t -> t
     val decompose : t -> Alg.t list
     val simplify_cof : t -> cof -> cof
+    val forall_cof : t -> dim * cof -> cof
     val meet2 : t -> t -> t
   end
 end
@@ -236,6 +237,10 @@ struct
       | Var v ->
         test_var thy' v
 
+    let finger thy sym : UF.finger = UF.finger sym thy.classes
+
+    let test_finger f r = UF.test_finger r f
+
     (* XXX: this function was never profiled *)
     let meet2' thy'1 thy'2 =
       let thy' =
@@ -384,6 +389,30 @@ struct
       in
       go cof
 
+    (* XXX: this function was never profiled *)
+    let forall_cof thy (sym, cof) =
+      let fingers = List.map thy ~f:(fun (thy', _) -> Alg.finger thy' sym) in
+      let forall_eq (x, y) =
+        match
+          List.for_all fingers ~f:(fun f -> Alg.test_finger f x),
+          List.for_all fingers ~f:(fun f -> Alg.test_finger f y)
+        with
+        | true, true -> Free.top
+        | true, _ | _, true -> Free.bot
+        | _ -> Free.eq x y
+      in
+      let rec go =
+        function
+        | Cof Eq (x, y) ->
+          forall_eq (x, y)
+        | Var v -> Var v
+        | Cof Join phis ->
+          B.join @@ List.map ~f:go phis
+        | Cof Meet phis ->
+          B.meet @@ List.map ~f:go phis
+      in
+      go cof
+
     (* XXX: this function was never profiled *)
     let meet2 (thy1 : t) (thy2 : t) : t =
       (* a correct but unoptimized theory *)

src/Theory.mli

@@ -91,6 +91,14 @@ sig
         Also, the simplification is an expensive process and should be used wisely. *)
     val simplify_cof : t -> cof -> cof
 
+    (** [forall_cof thy (r, cof)] computes [forall r. cof] with respect to the equations in the theory [thy],
+        using the syntactic quantifier elimination and potentially other simplification procedures
+        used in {!val:eq}, {!val:join}, and {!val:meet}. This is slower than {!val:Builder.Echo.S.forall}
+        which does not take equations into consideration.
+
+        Note: this is experimental and might be removed if it turns out that we do not need it. *)
+    val forall_cof : t -> dim * cof -> cof
+
     (** [meet2 thy1 thy2] computes the conjunction of the two theories [thy1] and [thy2].
         This is useful for supporting compilation units with top-level cofibration declarations. *)
     val meet2 : t -> t -> t