faster List.flat_map

Makefile

@@ -14,12 +14,6 @@ clean:
 doc:
 	dune build @doc
 
-BENCH_TARGETS=run_benchs.exe run_bench_hash.exe
-
-benchs:
-	dune build $(addprefix benchs/, $(BENCH_TARGETS)) --profile=release
-	@for i in $(BENCH_TARGETS) ; do ln -sf _build/default/benchs/$$i ; done
-
 examples:
 	dune build examples/id_sexp.exe
 
@@ -39,4 +33,4 @@ reindent:
 	@find src '(' -name '*.ml' -or -name '*.mli' ')' -type f -print0 | xargs -0 echo "reindenting: "
 	@find src '(' -name '*.ml' -or -name '*.mli' ')' -type f -print0 | xargs -0 ocp-indent -i
 
-.PHONY: all benchs test clean build doc update_next_tag watch examples
+.PHONY: all test clean build doc update_next_tag watch examples

benchs/run_benchs.ml

@@ -93,18 +93,42 @@ module L = struct
     else
       Sek.Persistent.of_list 0 [ x; x + 1; x + 2; x + 3 ]
 
+  let flat_map_kont f l =
+    let rec aux f l kont =
+      match l with
+      | [] -> kont []
+      | x :: l' ->
+        let y = f x in
+        let kont' tail =
+          match y with
+          | [] -> kont tail
+          | [ x ] -> kont (x :: tail)
+          | [ x; y ] -> kont (x :: y :: tail)
+          | l -> kont (CCList.append l tail)
+        in
+        aux f l' kont'
+    in
+    aux f l (fun l -> l)
+
   let bench_flat_map ?(time = 2) n =
     let l = CCList.(1 -- n) in
     let ral = CCRAL.of_list l in
     let sek = Sek.Persistent.of_list 0 l in
-    let flatten_map_ l () = ignore @@ List.flatten (CCList.map f_ l)
-    and flatmap l () = ignore @@ CCList.flat_map f_ l
-    and flatten_ccmap_ l () = ignore @@ List.flatten (List.map f_ l)
-    and flatmap_ral_ l () = ignore @@ CCRAL.flat_map f_ral_ l
-    and flatmap_sek s () = ignore @@ Sek.Persistent.flatten_map 0 f_sek_ s in
+    let flatten_map_ l () =
+      ignore @@ Sys.opaque_identity @@ List.flatten (CCList.map f_ l)
+    and flatmap_kont l () = ignore @@ Sys.opaque_identity @@ flat_map_kont f_ l
+    and flatmap l () = ignore @@ Sys.opaque_identity @@ CCList.flat_map f_ l
+    and flatten_ccmap_ l () =
+      ignore @@ Sys.opaque_identity @@ List.flatten (List.map f_ l)
+    and flatmap_ral_ l () =
+      ignore @@ Sys.opaque_identity @@ CCRAL.flat_map f_ral_ l
+    and flatmap_sek s () =
+      ignore @@ Sys.opaque_identity @@ Sek.Persistent.flatten_map 0 f_sek_ s
+    in
     B.throughputN time ~repeat
       [
         "flat_map", flatmap l, ();
+        "flat_map_kont", flatmap_kont l, ();
         "flatten o CCList.map", flatten_ccmap_ l, ();
         "flatten o map", flatten_map_ l, ();
         "ral_flatmap", flatmap_ral_ ral, ();
@@ -155,7 +179,11 @@ module L = struct
     and sek_flatten s () =
       opaque_ignore (Sek.Persistent.flatten s : _ Sek.Persistent.t)
     and funvec_flatten v () =
-      opaque_ignore (CCFun_vec.fold_rev ~x:CCFun_vec.empty ~f:(fun acc x -> CCFun_vec.append x acc) v : _ CCFun_vec.t)
+      opaque_ignore
+        (CCFun_vec.fold_rev ~x:CCFun_vec.empty
+           ~f:(fun acc x -> CCFun_vec.append x acc)
+           v
+          : _ CCFun_vec.t)
     in
     let l =
       CCList.mapi (fun i x -> CCList.(x -- (x + min i 100))) CCList.(1 -- n)
@@ -170,7 +198,7 @@ module L = struct
         "CCList.flatten", (fun () -> ignore (CCList.flatten l)), ();
         "List.flatten", (fun () -> ignore (List.flatten l)), ();
         "fold_right append", fold_right_append_ l, ();
-        "funvec.(fold_right append)", (funvec_flatten v), ();
+        "funvec.(fold_right append)", funvec_flatten v, ();
         "CCList.(fold_right append)", cc_fold_right_append_ l, ();
         "Sek.flatten", sek_flatten sek, ();
       ]
@@ -239,8 +267,8 @@ module L = struct
         opaque_ignore (CCRAL.set l i (-i))
       done
       (* TODO: implement set
-    and bench_funvec l () =
-      for _i = 0 to n-1 do opaque_ignore (CCFun_vec.set (* TODO *)) done
+         and bench_funvec l () =
+           for _i = 0 to n-1 do opaque_ignore (CCFun_vec.set (* TODO *)) done
       *)
     and bench_batvec l () =
       for i = 0 to n - 1 do
@@ -375,6 +403,8 @@ module L = struct
              "flat_map"
              @>> B.Tree.concat
                    [
+                     app_int (bench_flat_map ~time:2) 2;
+                     app_int (bench_flat_map ~time:2) 30;
                      app_int (bench_flat_map ~time:2) 100;
                      app_int (bench_flat_map ~time:2) 10_000;
                      app_int (bench_flat_map ~time:4) 100_000;

run_bench_hash.sh

@@ -0,0 +1,4 @@
+#!/bin/sh
+
+OPTS="--profile=release --display=quiet"
+exec dune exec $OPTS -- benchs/run_benchs_hash.exe $@

run_benchs.sh

@@ -0,0 +1,4 @@
+#!/bin/sh
+
+OPTS="--profile=release --display=quiet"
+exec dune exec $OPTS -- benchs/run_benchs.exe $@

src/core/CCList.ml

@@ -104,23 +104,28 @@ let map f l =
   in
   direct f direct_depth_default_ l
 
-let direct_depth_append_ = 10_000
-
 let append l1 l2 =
+  let[@inline] safe l1 l2 = List.rev_append (List.rev l1) l2 in
   let rec direct i l1 l2 =
     match l1 with
     | [] -> l2
+    | [ x ] -> x :: l2
     | _ when i = 0 -> safe l1 l2
-    | x :: l1' -> x :: direct (i - 1) l1' l2
-  and safe l1 l2 = List.rev_append (List.rev l1) l2 in
-  match l1 with
-  | [] -> l2
-  | [ x ] -> x :: l2
-  | [ x; y ] -> x :: y :: l2
-  | _ -> direct direct_depth_append_ l1 l2
+    | x :: y :: tl1 -> x :: y :: direct (i - 1) tl1 l2
+  in
+  direct 1000 l1 l2
 
 [@@@endif]
 
+(* Wrapper around [append] to optimize for the case of short [l1],
+   and for the case of [l2 = []] (saves the whole copy of [l1]!) *)
+let[@inline] append l1 l2 =
+  match l1, l2 with
+  | [], _ -> l2
+  | _, [] -> l1
+  | [ x ], _ -> x :: l2
+  | x :: y :: tl1, _ -> x :: y :: append tl1 l2
+
 let ( @ ) = append
 let[@inline] cons' l x = x :: l
 
@@ -324,22 +329,43 @@ let rec equal f l1 l2 =
   | [], _ | _, [] -> false
   | x1 :: l1', x2 :: l2' -> f x1 x2 && equal f l1' l2'
 
+let rec flat_map_kont f l kont =
+  match l with
+  | [] -> kont []
+  | [ x ] ->
+    let x = f x in
+    kont x
+  | x :: l' ->
+    let x = f x in
+    let kont' tail = kont (append x tail) in
+    flat_map_kont f l' kont'
+
+[@@@iflt 5.1]
+
+let[@inline] flat_map f l =
+  match l with
+  | [] -> []
+  | [ x ] -> f x
+  | x :: tl -> flat_map_kont f l Fun.id
+
+[@@@else_]
+
 let flat_map f l =
-  let rec aux f l kont =
+  let rec direct i f l =
     match l with
-    | [] -> kont []
-    | x :: l' ->
-      let y = f x in
-      let kont' tail =
-        match y with
-        | [] -> kont tail
-        | [ x ] -> kont (x :: tail)
-        | [ x; y ] -> kont (x :: y :: tail)
-        | l -> kont (append l tail)
-      in
-      aux f l' kont'
+    | [] -> []
+    | [ x ] -> f x
+    | [ x; y ] -> append (f x) (f y)
+    | _ when i = 0 -> flat_map_kont f l Fun.id
+    | x :: y :: tl ->
+      let x = f x in
+      let y = f y in
+      let tl = direct (i - 1) f tl in
+      append x (append y tl)
   in
-  aux f l (fun l -> l)
+  direct 1000 f l
+
+[@@@endif]
 
 let flat_map_i f l =
   let rec aux f i l kont =