Stdlib.Seq extension / 正直 Core とか Batteries とか使えば終わる話ではある

seq.ml

include Stdlib.Seq

let cons x xs = fun () -> Cons (x, xs)

let rev xs =
  let rec aux acc ys =
    match ys () with
    | Nil -> acc
    | Cons (y, ys) -> aux (cons y acc) ys
  in
  aux empty xs

let mem ?(equal=fun x y -> compare x y = 0) x xs =
  let rec aux acc xs =
    match xs () with
    | Nil -> acc
    | Cons (y, ys) -> aux (equal x y || acc) ys
  in
  aux false xs

let for_all f xs =
  let rec aux acc xs =
    match xs () with
    | Nil -> acc
    | Cons (y, ys) -> aux (f y && acc) ys
  in
  aux true xs

let rec drop count xs =
  if count > 0 then begin
    match xs () with
    | Nil -> empty
    | Cons (_, xs) -> drop (pred count) xs
  end else xs

let take count xs =
  let rec aux acc n rest =
    if n > 0 then begin
      match rest () with
      | Nil -> acc
      | Cons (y, ys) -> aux (cons y acc) (pred n) ys
    end else acc
  in
  aux empty count xs |> rev

let rec drop_while cond xs =
  match xs () with
  | Nil -> empty
  | Cons (x, xs) when cond x -> drop_while cond xs
  | _ -> xs

let take_while cond xs =
  let rec aux acc rest =
    match rest () with
    | Nil -> acc
    | Cons (y, ys) when cond y -> aux (cons y acc) ys
    | _ -> acc
  in
  aux empty xs |> rev

let unfold f init =
  let rec aux g k =
    match g () with
    | None -> k empty
    | Some (x, xs) -> aux (fun () -> f xs) (fun z -> k @@ cons x z)
  in
  aux (fun () -> f init) Fun.id

let generate f = unfold (fun x -> Some (f x, succ x)) 0

let init count f = generate f |> take count

let mapi f xs =
  let rec aux acc i ys =
    match ys () with
    | Nil -> rev acc
    | Cons (y, ys) -> aux (cons (f i y) acc) (succ i) ys
  in
  aux empty 0 xs

let iteri f xs =
  let rec aux i ys =
    match ys () with
    | Nil -> ()
    | Cons (y, ys) -> f i y; aux (succ i) ys
  in
  aux 0 xs

let of_list xs = List.to_seq xs

let to_list xs = List.of_seq xs

let of_option = function None -> empty | Some x -> return x

seq.mli

include module type of Stdlib.Seq

val cons : 'a -> 'a t -> 'a t
val rev : 'a t -> 'a t
val mem : ?equal:('a -> 'a -> bool) -> 'a -> 'a t -> bool
val for_all : ('a -> bool) -> 'a t -> bool
val drop : int -> 'a t -> 'a t
val take : int -> 'a t -> 'a t
val drop_while : ('a -> bool) -> 'a t -> 'a t
val take_while : ('a -> bool) -> 'a t -> 'a t
val unfold : ('a -> ('b * 'a) option) -> 'a -> 'b t
val generate : (int -> 'a) -> 'a t
val init : int -> (int -> 'a) -> 'a t
val mapi : (int -> 'a -> 'b) -> 'a t -> 'b t
val iteri : (int -> 'a -> unit) -> 'a t -> unit
val of_list : 'a list -> 'a t
val to_list : 'a t -> 'a list
val of_option : 'a option -> 'a t

$ ocamlc -c seq.mli seq.ml

#load "seq.cmo";;

let range (s, e) =
  let rec aux acc n =
    if n >= s then aux (n :: acc) (pred n)
    else acc
  in
  aux [] e

let () =
  let print_int_endline x = print_int x |> print_newline in
  
  range (1, 9)
  |> List.to_seq
  |> Seq.drop 3
  |> Seq.take 3
  |> Seq.iter print_int_endline; (* 4, 5, 6 *)
                                 
  print_endline "====";
  
  range (1, 9)
  |> List.to_seq
  |> Seq.drop_while (fun x -> x < 5)
  |> Seq.take_while (fun x -> x < 8)
  |> Seq.iter print_int_endline; (* 5, 6, 7 *)
;;

let is_prime x =
  if x < 2 then false else
    Seq.init x Fun.id
    |> Seq.drop 2
    |> Seq.for_all (fun n -> x mod n <> 0)

let () =
  let xs =
    Seq.init 100 succ
    |> Seq.filter is_prime
    |> Seq.map string_of_int
    |> Seq.rev
  in
  
  (* 97, 89, 83, 79, 73, 71, 67, 61, 59, 53, 47, 43, 41, 37, 31, 29, 23, 19, 17, 13, 11, 7, 5, 3, 2 *)
  Seq.iter print_endline xs