diff --git a/sequence.ml b/sequence.ml index 6c6bd3b..2fefe5f 100644 --- a/sequence.ml +++ b/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) diff --git a/sequence.mli b/sequence.mli index 677f79c..ce6c318 100644 --- a/sequence.mli +++ b/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. *) diff --git a/tests/test_sequence.ml b/tests/test_sequence.ml index 30f0f1c..8fd1b17 100644 --- a/tests/test_sequence.ml +++ b/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)