[sledge] Implement Qset using Import.Map

Reviewed By: ngorogiannis

Differential Revision: D20583759

fbshipit-source-id: d39b13508

sledge/lib/import/import0.ml

@@ -72,13 +72,6 @@ module type Monad_syntax = sig
   val ( and* ) : 'a t -> 'b t -> ('a * 'b) t
 end
 
-module type OrderedType = sig
-  type t
-
-  val compare : t -> t -> int
-  val sexp_of_t : t -> Sexp.t
-end
-
 exception Duplicate
 
 module Return = struct type 'r t = {return: 'a. 'r -> 'a} [@@unboxed] end

sledge/lib/import/qset.ml

@@ -10,8 +10,11 @@
 open Import0
 include Qset_intf
 
-module Make (Elt : OrderedType) = struct
-  module M = Stdlib.Map.Make (Elt)
+module Make (Elt : sig
+  type t [@@deriving compare, sexp_of]
+end) =
+struct
+  module M = Map.Make (Elt)
 
   type elt = Elt.t
   type t = Q.t M.t
@@ -21,16 +24,15 @@ module Make (Elt : OrderedType) = struct
 
   let hash_fold_t hash_fold_elt s m =
     let hash_fold_q s q = Hash.fold_int s (Hashtbl.hash q) in
-    M.fold
-      (fun key data state -> hash_fold_q (hash_fold_elt state key) data)
-      m
-      (Hash.fold_int s (M.cardinal m))
+    M.fold m
+      ~init:(Hash.fold_int s (M.length m))
+      ~f:(fun ~key ~data state -> hash_fold_q (hash_fold_elt state key) data)
 
   let sexp_of_t s =
     let sexp_of_q q = Sexp.Atom (Q.to_string q) in
     List.sexp_of_t
       (Sexplib.Conv.sexp_of_pair Elt.sexp_of_t sexp_of_q)
-      (M.bindings s)
+      (M.to_alist s)
 
   let t_of_sexp elt_of_sexp sexp =
     let q_of_sexp = function
@@ -38,50 +40,43 @@ module Make (Elt : OrderedType) = struct
       | _ -> assert false
     in
     List.fold_left
-      ~f:(fun m (k, v) -> M.add k v m)
+      ~f:(fun m (key, data) -> M.add_exn m ~key ~data)
       ~init:M.empty
       (List.t_of_sexp
          (Sexplib.Conv.pair_of_sexp elt_of_sexp q_of_sexp)
          sexp)
 
-  let pp sep pp_elt fs s = List.pp sep pp_elt fs (M.bindings s)
+  let pp sep pp_elt fs s = List.pp sep pp_elt fs (M.to_alist s)
   let empty = M.empty
   let if_nz q = if Q.equal Q.zero q then None else Some q
 
   let add m x i =
-    M.update x (function Some j -> if_nz Q.(i + j) | None -> if_nz i) m
+    M.change m x ~f:(function Some j -> if_nz Q.(i + j) | None -> if_nz i)
 
-  let remove m x = M.remove x m
+  let remove m x = M.remove m x
 
   let union m n =
-    M.merge
-      (fun _ m_q n_q ->
-        match (m_q, n_q) with
-        | Some i, Some j -> if_nz Q.(i + j)
-        | Some i, None | None, Some i -> Some i
-        | None, None -> None )
-      m n
+    M.merge m n ~f:(fun ~key:_ -> function
+      | `Both (i, j) -> if_nz Q.(i + j) | `Left i | `Right i -> Some i )
 
   let map m ~f =
-    let m' = M.empty in
+    let m' = empty in
     let m, m' =
-      M.fold
-        (fun x i (m, m') ->
+      M.fold m ~init:(m, m') ~f:(fun ~key:x ~data:i (m, m') ->
           let x', i' = f x i in
           if x' == x then
-            if Q.equal i' i then (m, m') else (M.add x i' m, m')
-          else (M.remove x m, add m' x' i') )
-        m (m, m')
+            if Q.equal i' i then (m, m') else (M.set m ~key:x ~data:i', m')
+          else (M.remove m x, add m' x' i') )
     in
-    M.fold (fun x i m -> add m x i) m' m
+    M.fold m' ~init:m ~f:(fun ~key:x ~data:i m -> add m x i)
 
-  let map_counts m ~f = M.mapi f m
-  let length m = M.cardinal m
-  let count m x = try M.find x m with Not_found -> Q.zero
-  let min_elt_exn = M.min_binding
-  let min_elt = M.min_binding_opt
-  let to_list m = M.bindings m
-  let iter m ~f = M.iter f m
-  let exists m ~f = M.exists f m
-  let fold m ~f ~init = M.fold f m init
+  let map_counts m ~f = M.mapi ~f:(fun ~key ~data -> f key data) m
+  let length m = M.length m
+  let count m x = match M.find m x with Some q -> q | None -> Q.zero
+  let min_elt_exn = M.min_elt_exn
+  let min_elt = M.min_elt
+  let to_list m = M.to_alist m
+  let iter m ~f = M.iteri ~f:(fun ~key ~data -> f key data) m
+  let exists m ~f = M.existsi ~f:(fun ~key ~data -> f key data) m
+  let fold m ~f ~init = M.fold ~f:(fun ~key ~data -> f key data) m ~init
 end

sledge/lib/import/qset.mli

@@ -7,6 +7,8 @@
 
 (** Qset - Set with (signed) rational multiplicity for each element *)
 
-open Import0
 include module type of Qset_intf
-module Make (Elt : OrderedType) : S with type elt = Elt.t
+
+module Make (Elt : sig
+  type t [@@deriving compare, sexp_of]
+end) : S with type elt = Elt.t