An implementation of the Sieve of Eratosthenes in OCaml

sieve.ml

type primality = Prime | Composite | Unknown
type sieve = primality array
exception Out_of_bounds

let make n =
  if n < 2 then invalid_arg "Sieve.make"
  else
    let sieve = Array.make n Unknown in
    sieve.(0) <- Composite;
    sieve.(1) <- Composite;
    sieve

let primality i sieve =
  sieve.(i)

let size sieve =
  Array.length sieve

let mark_prime i sieve =
  let rec mark_prime_aux i inc =
    if i < size sieve then begin
      sieve.(i) <- Composite;
      mark_prime_aux (i + inc) inc
    end in
  if sieve.(i) = Unknown then begin
    sieve.(i) <- Prime;
    mark_prime_aux (2 * i) i
  end

let find_prime n sieve =
  let rec find_prime_aux i pi =
    if i >= size sieve then raise Out_of_bounds
    else
      match sieve.(i) with
        Prime | Unknown -> begin
          (* it's a prime *)
          mark_prime i sieve;
          if pi = n then i
          else find_prime_aux (succ i) (succ pi)
        end
      | Composite -> find_prime_aux (succ i) pi in
  find_prime_aux 2 1

let iter f sieve =
  let rec iter_aux i =
    if i >= size sieve then ()
    else
      match sieve.(i) with
        Prime | Unknown -> begin
          mark_prime i sieve;
          f i;
          iter_aux (succ i)
        end
      | Composite -> iter_aux (succ i) in
  iter_aux 2

let fold f init sieve =
  let rec fold_aux i acc =
    if i >= size sieve then acc
    else
      match sieve.(i) with
        Prime | Unknown -> begin
          mark_prime i sieve;
          fold_aux (succ i) (f i acc)
        end
      | Composite -> fold_aux (succ i) acc in
  fold_aux 2 init

sieve.mli

type sieve
type primality = Prime | Composite | Unknown
val make : int -> sieve
val size : sieve -> int
val primality : int -> sieve -> primality
val mark_prime : int -> sieve -> unit
val find_prime : int -> sieve -> int
val iter : (int -> unit) -> sieve -> unit
(** iterates over all the primes in the sieve *)
val fold : (int -> 'a -> 'a) -> 'a -> sieve -> 'a
(** equivalent to a List.fold_left with a list of all primes in the sieve *)
exception Out_of_bounds