From fddb2fa991aed0c9fa8320aab1fd1b0553b27e31 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Thu, 16 Apr 2020 03:37:58 -0700
Subject: [PATCH] [sledge] Add Map.map_endo

Reviewed By: jvillard

Differential Revision: D20863525

fbshipit-source-id: 992ad494c
---
 sledge/lib/import/iArray.ml   | 10 +---------
 sledge/lib/import/import0.ml  | 22 ++++++++++++++++++++++
 sledge/lib/import/list.ml     | 22 ----------------------
 sledge/lib/import/map.ml      |  2 ++
 sledge/lib/import/map_intf.ml |  4 ++++
 5 files changed, 29 insertions(+), 31 deletions(-)

diff --git a/sledge/lib/import/iArray.ml b/sledge/lib/import/iArray.ml
index e45041efe..82346f257 100644
--- a/sledge/lib/import/iArray.ml
+++ b/sledge/lib/import/iArray.ml
@@ -51,15 +51,7 @@ let fold_map_until xs ~init ~f ~finish =
              | Continue x -> x
              | Stop x -> return x )) )
 
-let map_endo xs ~f =
-  let change = ref false in
-  let xs' =
-    map xs ~f:(fun x ->
-        let x' = f x in
-        if not (x' == x) then change := true ;
-        x' )
-  in
-  if !change then xs' else xs
+let map_endo xs ~f = map_endo map xs ~f
 
 let combine_adjacent ~f xs =
   let xs = i2a xs in
diff --git a/sledge/lib/import/import0.ml b/sledge/lib/import/import0.ml
index d4b6bd125..807b93a42 100644
--- a/sledge/lib/import/import0.ml
+++ b/sledge/lib/import/import0.ml
@@ -49,3 +49,25 @@ module type Monad_syntax = sig
   val ( let* ) : 'a t -> ('a -> 'b t) -> 'b t
   val ( and* ) : 'a t -> 'b t -> ('a * 'b) t
 end
+
+let map_endo map t ~f =
+  let change = ref false in
+  let t' =
+    map t ~f:(fun x ->
+        let x' = f x in
+        if x' != x then change := true ;
+        x' )
+  in
+  if !change then t' else t
+
+let filter_map_endo filter_map t ~f =
+  let change = ref false in
+  let t' =
+    filter_map t ~f:(fun x ->
+        let x'_opt = f x in
+        ( match x'_opt with
+        | Some x' when x' == x -> ()
+        | _ -> change := true ) ;
+        x'_opt )
+  in
+  if !change then t' else t
diff --git a/sledge/lib/import/list.ml b/sledge/lib/import/list.ml
index a62f6f279..57f8962e5 100644
--- a/sledge/lib/import/list.ml
+++ b/sledge/lib/import/list.ml
@@ -39,28 +39,6 @@ let fold_option xs ~init ~f =
     (fold xs ~init ~f:(fun acc elt ->
          match f acc elt with Some res -> res | None -> return None ))
 
-let map_endo map t ~f =
-  let change = ref false in
-  let t' =
-    map t ~f:(fun x ->
-        let x' = f x in
-        if not (x' == x) then change := true ;
-        x' )
-  in
-  if !change then t' else t
-
-let filter_map_endo filter_map t ~f =
-  let change = ref false in
-  let t' =
-    filter_map t ~f:(fun x ->
-        let x'_opt = f x in
-        ( match x'_opt with
-        | Some x' when x' == x -> ()
-        | _ -> change := true ) ;
-        x'_opt )
-  in
-  if !change then t' else t
-
 let filter_map_endo t ~f = filter_map_endo filter_map t ~f
 let map_endo t ~f = map_endo map t ~f
 
diff --git a/sledge/lib/import/map.ml b/sledge/lib/import/map.ml
index 876d74ca2..65117e08c 100644
--- a/sledge/lib/import/map.ml
+++ b/sledge/lib/import/map.ml
@@ -28,6 +28,8 @@ end) : S with type key = Key.t = struct
   let merge_skewed x y ~combine =
     of_map (Core.Map.merge_skewed (to_map x) (to_map y) ~combine)
 
+  let map_endo t ~f = map_endo map t ~f
+
   let fold_until m ~init ~f ~finish =
     let fold m ~init ~f =
       let f ~key ~data s = f s (key, data) in
diff --git a/sledge/lib/import/map_intf.ml b/sledge/lib/import/map_intf.ml
index e6e68c06f..e2166afe7 100644
--- a/sledge/lib/import/map_intf.ml
+++ b/sledge/lib/import/map_intf.ml
@@ -20,6 +20,10 @@ module type S = sig
 
   val compare : ('a -> 'a -> int) -> 'a t -> 'a t -> int
 
+  val map_endo : 'a t -> f:('a -> 'a) -> 'a t
+  (** Like map, but specialized to require [f] to be an endofunction, which
+      enables preserving [==] if [f] preserves [==] of every element. *)
+
   val merge_skewed :
     'a t -> 'a t -> combine:(key:key -> 'a -> 'a -> 'a) -> 'a t