[sledge] Minor Map interface simplifications

Summary:
Generalize the existing find_or_add and find_and_remove operations to
find_update. Slightly simplify the interfaces of change, update,
find_or_add, and find_and_remove, reducing the gap to the natural
underlying functionality.

Reviewed By: jvillard

Differential Revision: D26451305

fbshipit-source-id: 89f67c84d

sledge/nonstdlib/NSMap.ml

@@ -116,35 +116,28 @@ struct
   let find_multi k m =
     match M.find_opt k m with None -> [] | Some vs -> vs
 
-  let find_and_remove k m =
+  let find_update x m ~f =
     let found = ref None in
     let m =
-      M.update k
-        (fun v ->
-          found := v ;
-          None )
+      M.update x
+        (function
+          | None -> f None
+          | some_v ->
+              found := some_v ;
+              f some_v )
         m
     in
-    Option.map ~f:(fun v -> (v, m)) !found
+    (!found, m)
+
+  let find_and_remove = find_update ~f:(fun _ -> None)
 
   let find_or_add k v m =
-    let found = ref None in
-    let m =
-      M.update k
-        (function
-          | None -> Some v
-          | v ->
-              found := v ;
-              v )
-        m
-    in
-    match !found with Some v -> `Found v | None -> `Added m
+    find_update k ~f:(function None -> Some v | some_v -> some_v) m
 
   let pop_min_binding m =
     min_binding m |> Option.map ~f:(fun (k, v) -> (k, v, remove k m))
 
-  let change k m ~f = M.update k f m
-  let update k m ~f = M.update k (fun v -> Some (f v)) m
+  let update k m ~f = M.update k f m
   let map m ~f = M.map f m
   let mapi m ~f = M.mapi (fun key data -> f ~key ~data) m
   let map_endo t ~f = map_endo map t ~f

sledge/nonstdlib/NSMap_intf.ml

@@ -85,10 +85,16 @@ module type S = sig
   val find_exn : key -> 'a t -> 'a
   val find_multi : key -> 'a list t -> 'a list
 
-  val find_and_remove : key -> 'a t -> ('a * 'a t) option
+  val find_update :
+    key -> 'a t -> f:('a option -> 'a option) -> 'a option * 'a t
+  (** [find_update k f m] is [(found, m')] where [found] is [find k m] and
+      [find k m'] is [f found] and [find h m'] is [find h m] for all other
+      [h]. *)
+
+  val find_and_remove : key -> 'a t -> 'a option * 'a t
   (** Find and remove the binding for a key. *)
 
-  val find_or_add : key -> 'a -> 'a t -> [`Added of 'a t | `Found of 'a]
+  val find_or_add : key -> 'a -> 'a t -> 'a option * 'a t
   (** Find the value bound to the given key if there is one, or otherwise
       add a binding for the given key and value. *)
 
@@ -97,8 +103,7 @@ module type S = sig
 
   (** {1 Transform} *)
 
-  val change : key -> 'a t -> f:('a option -> 'a option) -> 'a t
-  val update : key -> 'a t -> f:('a option -> 'a) -> 'a t
+  val update : key -> 'a t -> f:('a option -> 'a option) -> 'a t
   val map : 'a t -> f:('a -> 'b) -> 'b t
   val mapi : 'a t -> f:(key:key -> data:'a -> 'b) -> 'b t
 

sledge/nonstdlib/multiset.ml

@@ -57,7 +57,7 @@ struct
   let if_nz i = if Mul.equal Mul.zero i then None else Some i
 
   let add x i m =
-    M.change x m ~f:(function
+    M.update x m ~f:(function
       | Some j -> if_nz (Mul.add i j)
       | None -> if_nz i )
 

sledge/nonstdlib/multiset_intf.ml

@@ -103,7 +103,7 @@ module type S = sig
   val classify : t -> (elt, mul) zero_one_many2
   (** Classify a set as either empty, singleton, or otherwise. *)
 
-  val find_and_remove : elt -> t -> (mul * t) option
+  val find_and_remove : elt -> t -> mul option * t
   (** Find and remove an element. *)
 
   val to_iter : t -> (elt * mul) iter

sledge/src/fol/arithmetic.ml

@@ -155,8 +155,8 @@ struct
 
     let split_const poly =
       match Sum.find_and_remove Mono.one poly with
-      | Some (c, p_c) -> (p_c, c)
-      | None -> (poly, Q.zero)
+      | Some c, p_c -> (p_c, c)
+      | None, _ -> (poly, Q.zero)
 
     let partition_sign poly =
       Sum.partition_map poly ~f:(fun _ coeff ->
@@ -443,7 +443,8 @@ struct
           let for_poly = trm for_ in
           match get_mono for_poly with
           | Some m ->
-              let* c, p = Sum.find_and_remove m a in
+              let c, p = Sum.find_and_remove m a in
+              let* c = c in
               solve_for_mono Sum.empty c m p
           | _ -> None )
   end

sledge/src/fol/context.ml

@@ -153,10 +153,11 @@ end = struct
         Trm.Map.fold s (t, ks, s) ~f:(fun ~key ~data (t, ks, s) ->
             if is_valid_eq ks key data then (t, ks, s)
             else
-              let t = Trm.Map.add ~key ~data t
-              and ks =
+              let t = Trm.Map.add ~key ~data t in
+              let ks =
                 Var.Set.diff ks (Var.Set.union (Trm.fv key) (Trm.fv data))
-              and s = Trm.Map.remove key s in
+              in
+              let s = Trm.Map.remove key s in
               (t, ks, s) )
       in
       if s' != s then partition_valid_ t' ks' s' else (t', ks', s')
@@ -474,7 +475,7 @@ let rec canon x a =
 
 let add_use_of sup sub use =
   Trm.Map.update sub use ~f:(fun u ->
-      Use.add sup (Option.value ~default:Use.empty u) )
+      Some (Use.add sup (Option.value ~default:Use.empty u)) )
 
 let add_uses_of a use =
   Iter.fold ~f:(add_use_of a) (Theory.solvable_trms a) use
@@ -493,8 +494,8 @@ let rec canon_extend_ a x =
   | InterpAtom -> (a, x)
   | NonInterpAtom -> (
     match Trm.Map.find_or_add a a x.rep with
-    | `Found a' -> (a', x)
-    | `Added rep -> (a, {x with rep}) )
+    | Some a', _ -> (a', x)
+    | None, rep -> (a, {x with rep}) )
   | InterpApp ->
       if Trm.Map.mem a x.rep then
         (* optimize: a already a rep so don't need to consider subterms *)
@@ -509,10 +510,10 @@ let rec canon_extend_ a x =
         (* norm wrt rep and add subterms *)
         let a_norm, x = Trm.fold_map ~f:canon_extend_ a x in
         match Trm.Map.find_or_add a_norm a_norm x.rep with
-        | `Found a' ->
+        | Some a', _ ->
             (* a_norm already equal to a' *)
             (a', x)
-        | `Added rep ->
+        | None, rep ->
             (* a_norm newly added *)
             let use = add_uses_of a_norm x.use in
             let x = {x with rep; use} in
@@ -531,13 +532,12 @@ let move_cls_to_rep a_cls a' rep =
   Cls.fold a_cls rep ~f:(fun e rep -> Trm.Map.add ~key:e ~data:a' rep)
 
 let find_and_move_cls noninterp ~of_:u ~to_:u' cls =
-  let u_cls, cls =
-    Trm.Map.find_and_remove u cls |> Option.value ~default:(Cls.empty, cls)
-  in
+  let u_cls, cls = Trm.Map.find_and_remove u cls in
+  let u_cls = Option.value u_cls ~default:Cls.empty in
   let u_cls = if noninterp then Cls.add u u_cls else u_cls in
   let cls =
     Trm.Map.update u' cls ~f:(fun u'_cls ->
-        Cls.union u_cls (Option.value u'_cls ~default:Cls.empty) )
+        Some (Cls.union u_cls (Option.value u'_cls ~default:Cls.empty)) )
   in
   (u_cls, cls)
 
@@ -548,7 +548,7 @@ let move_uses ~rem:f ~add:t use =
   (* remove f from use of each of its subterms not shared with t *)
   let use =
     Trm.Set.fold f_trms use ~f:(fun e use ->
-        Trm.Map.change e use ~f:(function
+        Trm.Map.update e use ~f:(function
           | Some e_use ->
               let e_use' = Use.remove f e_use in
               if Use.is_empty e_use' then None else Some e_use'
@@ -558,14 +558,14 @@ let move_uses ~rem:f ~add:t use =
   let use =
     Trm.Set.fold ft_trms use ~f:(fun e use ->
         Trm.Map.update e use ~f:(function
-          | Some e_use -> Use.add t (Use.remove f e_use)
+          | Some e_use -> Some (Use.add t (Use.remove f e_use))
           | None -> assert false ) )
   in
   (* add t to use of each of its subterms not shared with f *)
   let use =
     Trm.Set.fold t_trms use ~f:(fun e use ->
         Trm.Map.update e use ~f:(fun e_use ->
-            Use.add t (Option.value e_use ~default:Use.empty) ) )
+            Some (Use.add t (Option.value e_use ~default:Use.empty)) ) )
   in
   use
 
@@ -872,7 +872,7 @@ let trim ks x =
   let clss =
     Trm.Set.fold reps x.cls ~f:(fun rep clss ->
         Trm.Map.update rep clss ~f:(fun cls0 ->
-            Cls.add rep (Option.value cls0 ~default:Cls.empty) ) )
+            Some (Cls.add rep (Option.value cls0 ~default:Cls.empty)) ) )
   in
   (* enumerate expanded classes and update solution subst *)
   Trm.Map.fold clss x ~f:(fun ~key:a' ~data:ecls x ->

sledge/src/sh.ml

@@ -81,10 +81,10 @@ let fold_vars ?ignore_ctx ?ignore_pure fold_vars q s ~f =
 
 let rec var_strength_ xs m q =
   let add v m =
-    match Var.Map.find v m with
-    | None -> Var.Map.add ~key:v ~data:`Anonymous m
-    | Some `Anonymous -> Var.Map.add ~key:v ~data:`Existential m
-    | Some _ -> m
+    Var.Map.update v m ~f:(function
+      | None -> Some `Anonymous
+      | Some `Anonymous -> Some `Existential
+      | o -> o )
   in
   let xs = Var.Set.union xs q.xs in
   let m_stem =