Sequence.persistent_lazy added, for caching without upfront cost

sequence.ml

@@ -148,13 +148,16 @@ module MList = struct
     | Nil
     | Cons of 'a array * int ref * 'a node ref
 
-  let of_seq seq =
+  (* build and call callback on every element *)
+  let of_seq_with seq k =
     let start = ref Nil in
     let chunk_size = ref 8 in
     (* fill the list. prev: tail-reference from previous node *)
     let prev, cur = ref start, ref Nil in
     seq
-      (fun x -> match !cur with
+      (fun x ->
+        k x;  (* callback *)
+        match !cur with
         | Nil ->
           let n = !chunk_size in
           if n < 4096 then chunk_size := 2 * !chunk_size;
@@ -172,6 +175,9 @@ module MList = struct
     !prev := !cur;
     !start
 
+  let of_seq seq =
+    of_seq_with seq (fun _ -> ())
+
   let is_empty = function
     | Nil -> true
     | Cons _ -> false
@@ -237,6 +243,18 @@ let persistent seq =
   let l = MList.of_seq seq in
   MList.to_seq l
 
+type 'a lazy_seq = [`Suspend | `Cached of 'a t] ref
+
+let persistent_lazy (seq:'a t) =
+  let (r:'a lazy_seq) = ref `Suspend in
+  fun k ->
+    match !r with
+    | `Cached seq' -> seq' k
+    | `Suspend ->
+        (* here if this traversal is interruted, no caching occurs *)
+        let seq' = MList.of_seq_with seq k in
+        r := `Cached (MList.to_seq seq')
+
 (** Sort the sequence. Eager, O(n) ram and O(n ln(n)) time. *)
 let sort ?(cmp=Pervasives.compare) seq =
   (* use an intermediate list, then sort the list *)

sequence.mli

@@ -157,6 +157,17 @@ val persistent : 'a t -> 'a t
       {b Note}: calling persistent on an already persistent sequence
       will still make a new copy of the sequence! *)
 
+val persistent_lazy : 'a t -> 'a t
+  (** Lazy version of {!persistent}. When calling [persistent_lazy s],
+      a new sequence [s'] is immediately returned (without actually consuming
+      [s]) in constant time; the first time [s'] is iterated on,
+      it also consumes [s] and caches its content into a inner data
+      structure that will back [s'] for future iterations.
+
+      {b warning}: on the first traversal of [s'], if the traversal
+      is interrupted prematurely ({!take}, etc.) then [s'] will not be
+      memorized, and the next call to [s'] will traverse [s] again. *)
+
 val sort : ?cmp:('a -> 'a -> int) -> 'a t -> 'a t
   (** Sort the sequence. Eager, O(n) ram and O(n ln(n)) time.
       It iterates on elements of the argument sequence immediately,