List.partition_map : (a -> (b, c) Either.t) -> a list -> b list * c list

Changes

@@ -166,6 +166,10 @@ Working version
   type 'a Either.t = Left of 'a | Right of 'b
   (Gabriel Scherer, review by Daniel Bünzli, Thomas Refis, Jeremy Yallop)
 
+- #9066: List.partition_map :
+    ('a -> ('b, 'c) Either.t) -> 'a list -> 'b list * 'c list
+  (Gabriel Scherer, review by Jeremy Yallop)
+
 - #9587: Arg: new Rest_all spec to get all rest arguments in a list
   (this is similar to Rest, but makes it possible to detect when there
    are no arguments (an empty list) after the rest marker)

stdlib/.depend

@@ -392,13 +392,16 @@ stdlib__lexing.cmi :
 stdlib__list.cmo : \
     stdlib__sys.cmi \
     stdlib__seq.cmi \
+    stdlib__either.cmi \
     stdlib__list.cmi
 stdlib__list.cmx : \
     stdlib__sys.cmx \
     stdlib__seq.cmx \
+    stdlib__either.cmx \
     stdlib__list.cmi
 stdlib__list.cmi : \
-    stdlib__seq.cmi
+    stdlib__seq.cmi \
+    stdlib__either.cmi
 stdlib__listLabels.cmo : \
     stdlib__list.cmi \
     stdlib__listLabels.cmi

stdlib/list.ml

@@ -283,6 +283,17 @@ let partition p l =
   | x :: l -> if p x then part (x :: yes) no l else part yes (x :: no) l in
   part [] [] l
 
+let partition_map p l =
+  let rec part left right = function
+  | [] -> (rev left, rev right)
+  | x :: l ->
+     begin match p x with
+       | Either.Left v -> part (v :: left) right l
+       | Either.Right v -> part left (v :: right) l
+     end
+  in
+  part [] [] l
+
 let rec split = function
     [] -> ([], [])
   | (x,y)::l ->

stdlib/list.mli

@@ -274,6 +274,21 @@ val partition : ('a -> bool) -> 'a list -> 'a list * 'a list
    elements of [l] that do not satisfy [p].
    The order of the elements in the input list is preserved. *)
 
+val partition_map : ('a -> ('b, 'c) Either.t) -> 'a list -> 'b list * 'c list
+(** [partition_map f l] returns a pair of lists [(l1, l2)] such that,
+    for each element [x] of the input list [l]:
+    - if [f x] is [Left y1], then [y1] is in [l1], and
+    - if [f x] is [Right y2], then [y2] is in [l2].
+
+    The output elements are included in [l1] and [l2] in the same
+    relative order as the corresponding input elements in [l].
+
+    In particular, [partition_map (fun x -> if p x then Left x else Right x) l]
+    is equivalent to [partition p l].
+
+    @since 4.12.0
+*)
+
 
 (** {1 Association lists} *)
 

testsuite/tests/lib-list/test.ml

@@ -1,12 +1,20 @@
 (* TEST
 *)
 
+let is_even x = (x mod 2 = 0)
+
 let string_of_even_opt x =
-  if x mod 2 = 0 then
+  if is_even x then
     Some (string_of_int x)
   else
     None
 
+let string_of_even_or_int x =
+  if is_even x then
+    Either.Left (string_of_int x)
+  else
+    Either.Right x
+
 (* Standard test case *)
 let () =
   let l = List.init 10 (fun x -> x) in
@@ -36,6 +44,11 @@ let () =
 
   assert (List.filteri (fun i _ -> i < 2) (List.rev l) = [9; 8]);
 
+  assert (List.partition is_even [1; 2; 3; 4; 5]
+          = ([2; 4], [1; 3; 5]));
+  assert (List.partition_map string_of_even_or_int [1; 2; 3; 4; 5]
+          = (["2"; "4"], [1; 3; 5]));
+
   assert (List.compare_lengths [] [] = 0);
   assert (List.compare_lengths [1;2] ['a';'b'] = 0);
   assert (List.compare_lengths [] [1;2] < 0);