From ecb1bce470bcf6454f06dc5cb46f5a967489f10f Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Sun, 21 Feb 2021 13:17:23 -0800
Subject: [PATCH] [sledge] Adapt NSMap to Comparer interface

Reviewed By: jvillard

Differential Revision: D26250530

fbshipit-source-id: f28355d3b
---
 sledge/nonstdlib/NSMap.ml      | 58 ++++++++++++++++++++--------------
 sledge/nonstdlib/NSMap.mli     | 16 +++++++++-
 sledge/nonstdlib/NSMap_intf.ml |  6 +++-
 sledge/nonstdlib/dune          |  2 ++
 sledge/nonstdlib/multiset.ml   |  2 +-
 sledge/nonstdlib/multiset.mli  |  2 +-
 sledge/src/control.ml          |  8 +++--
 7 files changed, 64 insertions(+), 30 deletions(-)

diff --git a/sledge/nonstdlib/NSMap.ml b/sledge/nonstdlib/NSMap.ml
index b25093270..50f2366fa 100644
--- a/sledge/nonstdlib/NSMap.ml
+++ b/sledge/nonstdlib/NSMap.ml
@@ -8,32 +8,33 @@
 open! NS0
 include NSMap_intf
 
-module Make (Key : sig
-  type t [@@deriving compare, sexp_of]
-end) : S with type key = Key.t = struct
-  module M = Stdlib.Map.Make [@inlined] (Key)
+type ('key, +'a, 'compare_key) t = ('key, 'a, 'compare_key) Map.t
+[@@deriving compare, equal, sexp]
+
+type ('compare_key, 'compare_a) compare =
+  ('compare_key, 'compare_a) Map.compare
+[@@deriving compare, equal, sexp]
+
+module Make_from_Comparer (Key : sig
+  type t [@@deriving equal, sexp_of]
+
+  include Comparer.S with type t := t
+end) =
+struct
+  module M = Map.Make [@inlined] (Key)
 
   type key = Key.t
-  type 'a t = 'a M.t [@@deriving compare, equal]
+  type compare_key = Key.compare
+  type 'a t = 'a M.t [@@deriving compare]
 
-  let to_list = M.bindings
-  let of_list l = List.fold_left l M.empty ~f:(fun m (k, v) -> M.add k v m)
+  type 'compare_a compare = 'compare_a M.compare
+  [@@deriving compare, equal, sexp]
+
+  let comparer = M.comparer
 
-  let sexp_of_t sexp_of_data m =
-    to_list m
-    |> Sexplib.Conv.sexp_of_list
-         (Sexplib.Conv.sexp_of_pair Key.sexp_of_t sexp_of_data)
-
-  module Provide_of_sexp (Key : sig
-    type t = key [@@deriving of_sexp]
-  end) =
-  struct
-    let t_of_sexp data_of_sexp s =
-      s
-      |> Sexplib.Conv.list_of_sexp
-           (Sexplib.Conv.pair_of_sexp Key.t_of_sexp data_of_sexp)
-      |> of_list
-  end
+  include M.Provide_equal (Key)
+  include M.Provide_sexp_of (Key)
+  module Provide_of_sexp = M.Provide_of_sexp
 
   let empty = M.empty
   let singleton = M.singleton
@@ -210,10 +211,12 @@ end) : S with type key = Key.t = struct
   let existsi m ~f = M.exists (fun key data -> f ~key ~data) m
   let for_alli m ~f = M.for_all (fun key data -> f ~key ~data) m
   let fold m s ~f = M.fold (fun key data acc -> f ~key ~data acc) m s
-  let to_iter m = Iter.from_iter (fun f -> M.iter (fun k v -> f (k, v)) m)
   let keys m = Iter.from_iter (fun f -> M.iter (fun k _ -> f k) m)
   let values m = Iter.from_iter (fun f -> M.iter (fun _ v -> f v) m)
+  let to_iter m = Iter.from_iter (fun f -> M.iter (fun k v -> f (k, v)) m)
   let of_iter s = Iter.fold s M.empty ~f:(fun (k, v) m -> M.add k v m)
+  let to_list = M.bindings
+  let of_list l = List.fold_left l M.empty ~f:(fun m (k, v) -> M.add k v m)
 
   let symmetric_diff l r ~eq =
     let seq = ref Iter.empty in
@@ -250,3 +253,12 @@ end) : S with type key = Key.t = struct
     List.pp ~pre ~suf sep pp_diff_elt fs sd
 end
 [@@inline]
+
+module Make (Key : sig
+  type t [@@deriving compare, equal, sexp_of]
+end) =
+Make_from_Comparer (struct
+  include Key
+  include Comparer.Make (Key)
+end)
+[@@inline]
diff --git a/sledge/nonstdlib/NSMap.mli b/sledge/nonstdlib/NSMap.mli
index b94fd4b8e..7541df323 100644
--- a/sledge/nonstdlib/NSMap.mli
+++ b/sledge/nonstdlib/NSMap.mli
@@ -7,6 +7,20 @@
 
 include module type of NSMap_intf
 
+type ('key, +'a, 'compare_key) t [@@deriving compare, equal, sexp]
+type ('compare_key, 'compare_a) compare [@@deriving compare, equal, sexp]
+
 module Make (Key : sig
-  type t [@@deriving compare, sexp_of]
+  type t [@@deriving compare, equal, sexp_of]
 end) : S with type key = Key.t
+
+module Make_from_Comparer (Key : sig
+  type t [@@deriving equal, sexp_of]
+
+  include Comparer.S with type t := t
+end) :
+  S
+    with type key = Key.t
+    with type compare_key = Key.compare
+    with type 'compare_a compare = (Key.compare, 'compare_a) compare
+    with type 'a t = (Key.t, 'a, Key.compare) t
diff --git a/sledge/nonstdlib/NSMap_intf.ml b/sledge/nonstdlib/NSMap_intf.ml
index 856bd8c0c..b44f7e3e1 100644
--- a/sledge/nonstdlib/NSMap_intf.ml
+++ b/sledge/nonstdlib/NSMap_intf.ml
@@ -9,13 +9,17 @@ open! NS0
 
 module type S = sig
   type key
+  type compare_key
   type +'a t [@@deriving compare, equal, sexp_of]
 
+  include Comparer.S1 with type 'a t := 'a t
+
   module Provide_of_sexp (_ : sig
     type t = key [@@deriving of_sexp]
   end) : sig
-    val t_of_sexp : (Sexp.t -> 'a) -> Sexp.t -> 'a t
+    type 'a t [@@deriving of_sexp]
   end
+  with type 'a t := 'a t
 
   (** {1 Construct} *)
 
diff --git a/sledge/nonstdlib/dune b/sledge/nonstdlib/dune
index 757e97581..d09d84d7d 100644
--- a/sledge/nonstdlib/dune
+++ b/sledge/nonstdlib/dune
@@ -11,3 +11,5 @@
  (preprocess
   (pps ppx_sledge ppx_trace))
  (inline_tests))
+
+(include_subdirs unqualified)
diff --git a/sledge/nonstdlib/multiset.ml b/sledge/nonstdlib/multiset.ml
index 7d5636bfb..d1a940bee 100644
--- a/sledge/nonstdlib/multiset.ml
+++ b/sledge/nonstdlib/multiset.ml
@@ -11,7 +11,7 @@ open! NS0
 include Multiset_intf
 
 module Make (Elt : sig
-  type t [@@deriving compare, sexp_of]
+  type t [@@deriving compare, equal, sexp_of]
 end)
 (Mul : MULTIPLICITY) =
 struct
diff --git a/sledge/nonstdlib/multiset.mli b/sledge/nonstdlib/multiset.mli
index cca280839..b378f89b6 100644
--- a/sledge/nonstdlib/multiset.mli
+++ b/sledge/nonstdlib/multiset.mli
@@ -10,6 +10,6 @@
 include module type of Multiset_intf
 
 module Make (Elt : sig
-  type t [@@deriving compare, sexp_of]
+  type t [@@deriving compare, equal, sexp_of]
 end)
 (Mul : MULTIPLICITY) : S with type mul = Mul.t with type elt = Elt.t
diff --git a/sledge/src/control.ml b/sledge/src/control.ml
index eba954a91..1a197214f 100644
--- a/sledge/src/control.ml
+++ b/sledge/src/control.ml
@@ -15,7 +15,7 @@ module Make (Opts : Domain_intf.Opts) (Dom : Domain_intf.Dom) = struct
 
     val pp : t pp
 
-    type as_inlined_location = t [@@deriving compare, sexp_of]
+    type as_inlined_location = t [@@deriving compare, equal, sexp_of]
 
     val empty : t
     val push_call : Llair.func Llair.call -> Dom.from_call -> t -> t option
@@ -81,6 +81,8 @@ module Make (Opts : Domain_intf.Opts) (Dom : Domain_intf.Dom) = struct
         | _, Throw _ -> 1
         | Empty, Empty -> 0
 
+    let equal_as_inlined_location = [%compare.equal: as_inlined_location]
+
     let invariant s =
       let@ () = Invariant.invariant [%here] s [%sexp_of: t] in
       match s with
@@ -160,7 +162,7 @@ module Make (Opts : Domain_intf.Opts) (Dom : Domain_intf.Dom) = struct
           { dst: Llair.Block.t
           ; src: Llair.Block.t option
           ; stk: Stack.as_inlined_location }
-        [@@deriving compare, sexp_of]
+        [@@deriving compare, equal, sexp_of]
       end
 
       include T
@@ -175,7 +177,7 @@ module Make (Opts : Domain_intf.Opts) (Dom : Domain_intf.Dom) = struct
     module Depths = struct
       module M = Map.Make (Edge)
 
-      type t = int M.t [@@deriving compare, sexp_of]
+      type t = int M.t [@@deriving compare, equal, sexp_of]
 
       let empty = M.empty
       let find = M.find