From d5fe9aa11d04334955c0653adfec15435dd2c8e9 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Mon, 23 Mar 2020 06:18:40 -0700
Subject: [PATCH] [sledge] Define Map as a functor over the Tree underlying
 Core.Map

Summary:
This diff defines Map as a functorover the underlying implementation
of Core.Map. This results in map values that are just trees, with no
comparison function closures, and with the same interface (almost) and
underlying data structure implementation as Core.Map.

Reviewed By: ngorogiannis

Differential Revision: D20583758

fbshipit-source-id: 5a4997b51
---
 sledge/lib/equality.ml        |   4 +-
 sledge/lib/import/map.ml      | 131 ++++++++--------------------------
 sledge/lib/import/map.mli     |   6 +-
 sledge/lib/import/map_intf.ml |  59 ++++-----------
 sledge/lib/term.ml            |   4 +-
 sledge/lib/term.mli           |   6 +-
 6 files changed, 54 insertions(+), 156 deletions(-)

diff --git a/sledge/lib/equality.ml b/sledge/lib/equality.ml
index c17fc2df5..9f322c2c1 100644
--- a/sledge/lib/equality.ml
+++ b/sledge/lib/equality.ml
@@ -58,7 +58,7 @@ module Subst : sig
 end = struct
   type t = Term.t Term.Map.t [@@deriving compare, equal, sexp_of]
 
-  let t_of_sexp = Term.Map.t_of_sexp Term.t_of_sexp Term.t_of_sexp
+  let t_of_sexp = Term.Map.t_of_sexp Term.t_of_sexp
   let pp = Term.Map.pp Term.pp Term.pp
 
   let pp_diff =
@@ -72,7 +72,7 @@ end = struct
   let fold = Term.Map.fold
   let iteri = Term.Map.iteri
   let for_alli = Term.Map.for_alli
-  let to_alist = Term.Map.to_alist
+  let to_alist = Term.Map.to_alist ~key_order:`Increasing
 
   (** look up a term in a substitution *)
   let apply s a = Term.Map.find s a |> Option.value ~default:a
diff --git a/sledge/lib/import/map.ml b/sledge/lib/import/map.ml
index e7c6cbb63..ea84a599b 100644
--- a/sledge/lib/import/map.ml
+++ b/sledge/lib/import/map.ml
@@ -5,39 +5,45 @@
  * LICENSE file in the root directory of this source tree.
  *)
 
-open Import0
 include Map_intf
 
-module Make (Key : OrderedType) : S with type key = Key.t = struct
-  module M = Stdlib.Map.Make (Key)
+module Make (Key : sig
+  type t [@@deriving compare, sexp_of]
+end) : S with type key = Key.t = struct
+  module KeyMap = Core.Map.Make_plain (Key)
+  module Key = KeyMap.Key
 
   type key = Key.t
-  type 'a t = 'a M.t
 
-  let compare = M.compare
-  let equal = M.equal
+  include KeyMap.Tree
 
-  let sexp_of_t sexp_of_val m =
-    List.sexp_of_t
-      (Sexplib.Conv.sexp_of_pair Key.sexp_of_t sexp_of_val)
-      (M.bindings m)
+  let compare = compare_direct
 
-  let t_of_sexp key_of_sexp val_of_sexp sexp =
-    List.fold_left
-      ~f:(fun m (k, v) -> M.add k v m)
-      ~init:M.empty
-      (List.t_of_sexp
-         (Sexplib.Conv.pair_of_sexp key_of_sexp val_of_sexp)
-         sexp)
+  let to_map t =
+    Base.Map.Using_comparator.of_tree ~comparator:Key.comparator t
+
+  let of_map m = Base.Map.Using_comparator.to_tree m
+
+  let merge_skewed x y ~combine =
+    of_map (Core.Map.merge_skewed (to_map x) (to_map y) ~combine)
+
+  let find_and_remove m k =
+    let found = ref None in
+    let m =
+      change m k ~f:(fun v ->
+          found := v ;
+          None )
+    in
+    Option.map ~f:(fun v -> (v, m)) !found
 
   let pp pp_k pp_v fs m =
     Format.fprintf fs "@[<1>[%a]@]"
       (List.pp ",@ " (fun fs (k, v) ->
            Format.fprintf fs "@[%a @<2>↦ %a@]" pp_k k pp_v v ))
-      (M.bindings m)
+      (to_alist m)
 
   let pp_diff ~data_equal pp_key pp_val pp_diff_val fs (x, y) =
-    let pp_diff_val fs = function
+    let pp_diff_elt fs = function
       | k, `Left v ->
           Format.fprintf fs "-- [@[%a@ @<2>↦ %a@]]" pp_key k pp_val v
       | k, `Right v ->
@@ -45,88 +51,7 @@ module Make (Key : OrderedType) : S with type key = Key.t = struct
       | k, `Unequal vv ->
           Format.fprintf fs "[@[%a@ @<2>↦ %a@]]" pp_key k pp_diff_val vv
     in
-    let sd =
-      M.merge
-        (fun _ v1o v2o ->
-          match (v1o, v2o) with
-          | Some v1, Some v2 when not (data_equal v1 v2) ->
-              Some (`Unequal (v1, v2))
-          | Some v1, None -> Some (`Left v1)
-          | None, Some v2 -> Some (`Right v2)
-          | _ -> None )
-        x y
-    in
-    if not (M.is_empty sd) then
-      Format.fprintf fs "[@[<hv>%a@]];@ "
-        (List.pp ";@ " pp_diff_val)
-        (M.bindings sd)
-
-  let empty = M.empty
-  let set m ~key ~data = M.add key data m
-
-  let add_exn m ~key ~data =
-    M.update key
-      (function None -> Some data | Some _ -> raise Duplicate)
-      m
-
-  let add_multi m ~key ~data =
-    M.update key
-      (function None -> Some [data] | Some vs -> Some (data :: vs))
-      m
-
-  let remove m k = M.remove k m
-  let update m k ~f = M.update k (fun vo -> Some (f vo)) m
-
-  let merge m n ~f =
-    M.merge
-      (fun k v1o v2o ->
-        match (v1o, v2o) with
-        | Some v1, Some v2 -> f ~key:k (`Both (v1, v2))
-        | Some v1, None -> f ~key:k (`Left v1)
-        | None, Some v2 -> f ~key:k (`Right v2)
-        | None, None -> None )
-      m n
-
-  let merge_skewed m n ~combine =
-    M.merge
-      (fun k v1o v2o ->
-        match (v1o, v2o) with
-        | Some v1, Some v2 -> Some (combine ~key:k v1 v2)
-        | Some _, None -> v1o
-        | None, Some _ -> v2o
-        | None, None -> None )
-      m n
-
-  let map m ~f = M.map f m
-  let filter_keys m ~f = M.filter (fun k _ -> f k) m
-
-  let filter_mapi m ~f =
-    M.fold
-      (fun k v m ->
-        match f ~key:k ~data:v with Some v' -> M.add k v' m | None -> m )
-      m M.empty
-
-  let is_empty = M.is_empty
-  let length = M.cardinal
-  let mem m k = M.mem k m
-  let find m k = M.find_opt k m
-
-  let find_and_remove m k =
-    let found = ref None in
-    let m =
-      M.update k
-        (fun v ->
-          found := v ;
-          None )
-        m
-    in
-    Option.map ~f:(fun v -> (v, m)) !found
-
-  let find_multi m k = try M.find k m with Not_found -> []
-  let data m = M.fold (fun _ v s -> v :: s) m []
-  let to_alist = M.bindings
-  let iter m ~f = M.iter (fun _ v -> f v) m
-  let iteri m ~f = M.iter (fun k v -> f ~key:k ~data:v) m
-  let for_alli m ~f = M.for_all (fun key data -> f ~key ~data) m
-  let fold m ~init ~f = M.fold (fun key data s -> f ~key ~data s) m init
+    let sd = Core.Sequence.to_list (symmetric_diff ~data_equal x y) in
+    if not (List.is_empty sd) then
+      Format.fprintf fs "[@[<hv>%a@]];@ " (List.pp ";@ " pp_diff_elt) sd
 end
diff --git a/sledge/lib/import/map.mli b/sledge/lib/import/map.mli
index bc9dceeb7..01f652af4 100644
--- a/sledge/lib/import/map.mli
+++ b/sledge/lib/import/map.mli
@@ -5,6 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  *)
 
-open Import0
 include module type of Map_intf
-module Make (Key : OrderedType) : S with type key = Key.t
+
+module Make (Key : sig
+  type t [@@deriving compare, sexp_of]
+end) : S with type key = Key.t
diff --git a/sledge/lib/import/map_intf.ml b/sledge/lib/import/map_intf.ml
index 850b9eb8e..ccc9c7bdc 100644
--- a/sledge/lib/import/map_intf.ml
+++ b/sledge/lib/import/map_intf.ml
@@ -9,60 +9,27 @@ open Import0
 
 module type S = sig
   type key
-  type +'a t
 
-  val compare : ('a -> 'a -> int) -> 'a t -> 'a t -> int
-  val equal : ('a -> 'a -> bool) -> 'a t -> 'a t -> bool
-  val sexp_of_t : ('a -> Sexp.t) -> 'a t -> Sexp.t
-  val t_of_sexp : (Sexp.t -> key) -> (Sexp.t -> 'a) -> Sexp.t -> 'a t
-  val pp : key pp -> 'a pp -> 'a t pp
-
-  val pp_diff :
-       data_equal:('a -> 'a -> bool)
-    -> key pp
-    -> 'a pp
-    -> ('a * 'a) pp
-    -> ('a t * 'a t) pp
+  module Key : sig
+    type t = key
 
-  (* initial constructors *)
-  val empty : 'a t
+    include Core.Comparator.S with type t := t
+  end
 
-  (* constructors *)
-  val set : 'a t -> key:key -> data:'a -> 'a t
-  val add_exn : 'a t -> key:key -> data:'a -> 'a t
-  val add_multi : 'a list t -> key:key -> data:'a -> 'a list t
-  val remove : 'a t -> key -> 'a t
-  val update : 'a t -> key -> f:('a option -> 'a) -> 'a t
+  include Core_kernel.Map_intf.Make_S_plain_tree(Key).S
 
-  val merge :
-       'a t
-    -> 'b t
-    -> f:
-         (   key:key
-          -> [`Both of 'a * 'b | `Left of 'a | `Right of 'b]
-          -> 'c option)
-    -> 'c t
+  val compare : ('a -> 'a -> int) -> 'a t -> 'a t -> int
 
   val merge_skewed :
     'a t -> 'a t -> combine:(key:key -> 'a -> 'a -> 'a) -> 'a t
 
-  val map : 'a t -> f:('a -> 'b) -> 'b t
-  val filter_keys : 'a t -> f:(key -> bool) -> 'a t
-  val filter_mapi : 'a t -> f:(key:key -> data:'a -> 'b option) -> 'b t
-
-  (* queries *)
-  val is_empty : 'b t -> bool
-  val length : 'b t -> int
-  val mem : 'a t -> key -> bool
-  val find : 'a t -> key -> 'a option
   val find_and_remove : 'a t -> key -> ('a * 'a t) option
-  val find_multi : 'a list t -> key -> 'a list
-  val data : 'a t -> 'a list
-  val to_alist : 'a t -> (key * 'a) list
+  val pp : key pp -> 'a pp -> 'a t pp
 
-  (* traversals *)
-  val iter : 'a t -> f:('a -> unit) -> unit
-  val iteri : 'a t -> f:(key:key -> data:'a -> unit) -> unit
-  val for_alli : 'a t -> f:(key:key -> data:'a -> bool) -> bool
-  val fold : 'a t -> init:'s -> f:(key:key -> data:'a -> 's -> 's) -> 's
+  val pp_diff :
+       data_equal:('a -> 'a -> bool)
+    -> key pp
+    -> 'a pp
+    -> ('a * 'a) pp
+    -> ('a t * 'a t) pp
 end
diff --git a/sledge/lib/term.ml b/sledge/lib/term.ml
index 77b639ba6..c0582ab1b 100644
--- a/sledge/lib/term.ml
+++ b/sledge/lib/term.ml
@@ -100,7 +100,7 @@ end = struct
 end
 
 include T
-module Map = Map.Make (T)
+module Map = struct include Map.Make (T) include Provide_of_sexp (T) end
 module Set = struct include Set.Make (T) include Provide_of_sexp (T) end
 
 let fix (f : (t -> 'a as 'f) -> 'f) (bot : 'f) (e : t) : 'a =
@@ -334,7 +334,7 @@ module Var = struct
   module Subst = struct
     type t = T.t Map.t [@@deriving compare, equal, sexp_of]
 
-    let t_of_sexp = Map.t_of_sexp T.t_of_sexp T.t_of_sexp
+    let t_of_sexp = Map.t_of_sexp T.t_of_sexp
 
     let invariant s =
       Invariant.invariant [%here] s [%sexp_of: t]
diff --git a/sledge/lib/term.mli b/sledge/lib/term.mli
index d29e3de48..a03955f5b 100644
--- a/sledge/lib/term.mli
+++ b/sledge/lib/term.mli
@@ -141,7 +141,11 @@ module Var : sig
   end
 end
 
-module Map : Map.S with type key := t
+module Map : sig
+  include Map.S with type key := t
+
+  val t_of_sexp : (Sexp.t -> 'a) -> Sexp.t -> 'a t
+end
 
 val ppx : Var.strength -> t pp
 val pp : t pp