diff --git a/src/core/CCList.ml b/src/core/CCList.ml index e00e1e32..d5d791ac 100644 --- a/src/core/CCList.ml +++ b/src/core/CCList.ml @@ -281,11 +281,6 @@ let fold_flat_map_i f acc l = in aux f acc 0 [] l -let rec unfold f seed = - match f seed with - | None -> [] - | Some (v, next) -> v :: unfold f next - [@@@iflt 5.1] (* keep this because it's tailrec for < 5.1 *) @@ -314,6 +309,32 @@ let init len f = else direct_ 0 +let rec unfold_kont f seed k = + match f seed with + | None -> k [] + | Some (v, next) -> + let k' tl = v :: tl in + unfold_kont f next k' + +let[@inline] unfold f seed = + let rec direct i f seed = + if i = 0 then + unfold_kont f seed (fun x -> x) + else ( + match f seed with + | None -> [] + | Some (v, next) -> v :: direct (i - 1) f next + ) + in + direct 100 f seed + +[@@@else_] + +let[@tail_mod_cons] rec unfold f seed = + match f seed with + | None -> [] + | Some (v, next) -> v :: unfold f next + [@@@endif] let rec compare f l1 l2 = diff --git a/tests/core/t_list.ml b/tests/core/t_list.ml index 4d2026ba..b32aef65 100644 --- a/tests/core/t_list.ml +++ b/tests/core/t_list.ml @@ -141,6 +141,20 @@ in unfold f 0 = [ 0; 2; 4; 6; 8; 10 ] ;; +t @@ fun () -> +let l = + unfold + (fun n -> + if n < 1_000_000 then + Some (n, n + 1) + else + None) + 0 +in +assert_equal ~printer:Q.Print.int (List.length l) 1_000_000; +true +;; + t @@ fun () -> init 0 (fun _ -> 0) = [];; t @@ fun () -> init 1 (fun x -> x) = [ 0 ];; t @@ fun () -> init 1000 (fun x -> x) = 0 -- 999;;