Stack.pop_until, a handy ocaml function

open Core

module Stack = struct
  include Stack

  let rec pop_until pred stack =
    pop_until' pred stack 0

  and pop_until' pred stack n =
    let a = Stack.pop stack in
    if pred a then
      (a, n)
    else
      pop_until' pred stack (n + 1)

  let rec pop_until_exn pred stack =
    pop_until_exn' pred stack 0

  and pop_until_exn' pred stack n =
    let a = Stack.pop_exn stack in
    if pred a then
      (a, n)
    else
      pop_until_exn' pred stack (n + 1)

end

let%expect_test _ =
  let stack = Stack.create () in
    Stack.push stack 10;
    Stack.push stack 20;
    Stack.push stack 30;

    printf "initial top: %d\n" (Stack.top_exn stack);
    let v, m = Stack.pop_until_exn (fun i -> i <= 19) stack in
      printf "popped count: %d\n" m;
      printf "popped value: %d\n" v;
      printf "empty?: %b\n" (Stack.is_empty stack);

      [%expect {|
        initial top: 30
        popped count: 2
        popped value: 10
        empty?: true
      |}]

let%expect_test _ =
  try
    let stack = Stack.create () in
    ignore (Stack.pop_until_exn (fun i -> i < 0) stack);
  with
    e -> printf "error: %s\n" (Exn.to_string e)
  ;

  [%expect {|error: "Stack.pop of empty stack"|}]