[sledge] Add Map.partition_map

Reviewed By: jvillard

Differential Revision: D24532356

fbshipit-source-id: 3f9c1a56c

sledge/nonstdlib/map.ml

@@ -64,6 +64,15 @@ end) : S with type key = Key.t = struct
 
   let union x y ~f = M.union f x y
   let partition m ~f = M.partition f m
+
+  let partition_map m ~f =
+    M.fold
+      (fun k v (l, r) ->
+        match (f k v : _ Either.t) with
+        | Left a -> (M.add k a l, r)
+        | Right b -> (l, M.add k b r) )
+      m (empty, empty)
+
   let is_empty = M.is_empty
 
   let root_key m =

sledge/nonstdlib/map_intf.ml

@@ -50,6 +50,9 @@ module type S = sig
   val union : 'a t -> 'a t -> f:(key -> 'a -> 'a -> 'a option) -> 'a t
   val partition : 'a t -> f:(key -> 'a -> bool) -> 'a t * 'a t
 
+  val partition_map :
+    'a t -> f:(key -> 'a -> ('b, 'c) Either.t) -> 'b t * 'c t
+
   (** {1 Query} *)
 
   val is_empty : 'a t -> bool

sledge/nonstdlib/multiset.ml

@@ -64,6 +64,7 @@ struct
       | `Right j -> Some (Mul.neg j) )
 
   let partition = M.partition
+  let partition_map = M.partition_map
 
   let map m ~f =
     let m' = empty in

sledge/nonstdlib/multiset_intf.ml

@@ -71,6 +71,7 @@ module type S = sig
       [Mul.equal m m'] for all elements. *)
 
   val partition : t -> f:(elt -> mul -> bool) -> t * t
+  val partition_map : t -> f:(elt -> mul -> (mul, mul) Either.t) -> t * t
 
   (* queries *)