fold_while impl for sequences

sequence.ml

@@ -347,6 +347,21 @@ let take_while p seq k =
     seq (fun x -> if p x then k x else raise ExitTakeWhile)
   with ExitTakeWhile -> ()
 
+exception ExitFoldWhile
+
+let fold_while f s seq =
+  let state = ref s in
+  let consume x =
+      let acc, cont = f (!state) x in
+      state := acc;
+      match cont with
+      | `Stop -> raise ExitFoldWhile
+      | `Continue -> ()
+  in
+  try
+    seq consume; !state
+  with ExitFoldWhile -> !state
+
 let drop n seq k =
   let count = ref 0 in
   seq (fun x -> if !count >= n then k x else incr count)

sequence.mli

@@ -119,7 +119,7 @@ val iter : ('a -> unit) -> 'a t -> unit
 val iteri : (int -> 'a -> unit) -> 'a t -> unit
   (** Iterate on elements and their index in the sequence *)
 
-val fold : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b
+val fold : ('a -> 'b -> 'a) -> 'a -> 'b t -> 'a
   (** Fold over elements of the sequence, consuming it *)
 
 val foldi : ('b -> int -> 'a -> 'b) -> 'b -> 'a t -> 'b
@@ -272,6 +272,11 @@ val take_while : ('a -> bool) -> 'a t -> 'a t
       Will work on an infinite sequence [s] if the predicate is false for at
       least one element of [s]. *)
 
+val fold_while : ('a -> 'b -> 'a * [`Stop | `Continue]) -> 'a -> 'b t -> 'a
+  (** Folds over elements of the sequence, stopping early if the accumulator
+      returns [('a, `Stop)]
+      @since NEXT_RELEASE *)
+
 val drop : int -> 'a t -> 'a t
   (** Drop the [n] first elements of the sequence. Lazy. *)
 

tests/test_sequence.ml

@@ -42,6 +42,12 @@ let test_foldi () =
   OUnit.assert_equal [1, "world"; 0, "hello"] l;
   ()
 
+let test_fold_while () =
+  let n = S.of_list [true;true;false;true]
+    |> S.fold_while (fun acc b -> if b then acc+1, `Continue else acc, `Stop) 0 in
+  OUnit.assert_equal 2 n;
+  ()
+
 let test_exists () =
   S.(1 -- 100)
     |> S.exists (fun x -> x = 59)