[sledge] Add Map.map_endo

Reviewed By: jvillard

Differential Revision: D20863525

fbshipit-source-id: 992ad494c

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

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

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
 

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

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