primes.ml

let label = "iitalics"
let time_limit = 5.
let threads = 1
let tags = "algorithm=base,faithful=yes,bits=1"

let report iters time =
  Printf.printf "%s;%d;%f;%d;%s\n"
    label iters time threads tags

(* optionally runs afterwards to check solution against a reference *)
let verify is_prime =
  let errors = ref [] in
  try
    let f = open_in "primes.bin" in
    let rec loop i =
      let check expected =
        let output = is_prime i in
        if expected <> output && List.length !errors < 8 then
          errors := (i,output) :: !errors
      in
      match input_char f with
      | 'P' -> (check true;  loop (i + 1))
      | 'c' -> (check false; loop (i + 1))
      | _   -> loop i
    in
    loop 1
  with
  | Sys_error _ ->
     print_endline "verification skipped"
  | End_of_file ->
     match List.rev !errors with
     | [] -> print_endline "verification ok"
     | errs ->
        ( print_endline "verification failed:"
        ; errs |> List.iteri (fun i (x, out) ->
                      Printf.printf "%d. is_prime(%d) = %s\n"
                        (i + 1) x (Bool.to_string out)) )

module Storage =
  struct
    type cell = {v : int} [@@ocaml.unboxed]
    let cell_of_int : int -> cell = Obj.magic
    let cell_of_char : char -> cell = Obj.magic
    let char_of_cell : cell -> char = Obj.magic

    let bits = 8
    let _ = assert (Sys.int_size >= bits)
    let[@inline always] index n = n / bits
    let[@inline always] mask n = 1 lsl (n mod bits)
    let[@inline always] set {v} mask = {v = v lor mask}
    let[@inline always] test {v} mask = v land mask = 0

    type t = {a : bytes} [@@ocaml.unboxed]
    let get {a} i = cell_of_char (Bytes.unsafe_get a i)
    let put {a} i c = Bytes.unsafe_set a i (char_of_cell c)
    let create mask0 mask1 ~count =
      let len = (count + 1) / bits + 1 in
      let t = {a = Bytes.make len (cell_of_int mask1 |> char_of_cell)} in
      (put t 0 (cell_of_int mask0); t)

    module O = struct
      let ( |! ) = set
      let ( |? ) = test
      let ( .|() ) = get
      let ( .|()<- ) = put
    end
  end

module Sieve =
  struct
    type t =
      { count : int
      ; flags : Storage.t }

    open Storage.O

    (** [is_prime sv n] uses sieve [sv] to determine the primeness of [n], returning
     *  [true] iff it is prime. *)
    let is_prime sv n =
      if n < 1 || n > sv.count then
        invalid_arg "out of range of sieve"
      else
        let i, m = Storage.(index n, mask n) in
        sv.flags.|(i) |? m

    (** [create count] returns a fully populated sieve handling primes between
     *  1 and [count]. *)
    let create count =
      let flags =
        Storage.create ~count
          0x52 (* 2 is prime, 1 is not *)
          0x55 (* other evens are prime *)
      in

      (* naive fill routine *)
      let rec fill inc n =
        if n <= count then
          let i, m = Storage.(index n, mask n) in
          ( flags.|(i) <- flags.|(i) |! m
          ; fill inc (n + inc) )
      in

      (* invariant: (n mod 8 = 1) && (inc mod 8 = 2) *)
      let rec fill_2r inc n =
        if n + inc*3 <= count then
          let i0 = Storage.index n in
          let i1 = Storage.index (n + inc) in
          let i2 = Storage.index (n + inc*2) in
          let i3 = Storage.index (n + inc*3) in
          ( flags.|(i0) <- flags.|(i0) |! (1 lsl 1)
          ; flags.|(i1) <- flags.|(i1) |! (1 lsl 3)
          ; flags.|(i2) <- flags.|(i2) |! (1 lsl 5)
          ; flags.|(i3) <- flags.|(i3) |! (1 lsl 7)
          ; fill_2r inc (n + inc*4) )
        else
          fill inc n
      in

      (* invariant: (n mod 8 = 1) && (inc mod 8 = 6) *)
      let rec fill_6r inc n =
        if n + inc*3 <= count then
          let i0 = Storage.index n in
          let i1 = Storage.index (n + inc) in
          let i2 = Storage.index (n + inc*2) in
          let i3 = Storage.index (n + inc*3) in
          ( flags.|(i0) <- flags.|(i0) |! (1 lsl 1)
          ; flags.|(i1) <- flags.|(i1) |! (1 lsl 7)
          ; flags.|(i2) <- flags.|(i2) |! (1 lsl 5)
          ; flags.|(i3) <- flags.|(i3) |! (1 lsl 3)
          ; fill_6r inc (n + inc*4) )
        else
          fill inc n
      in

      (* runs naive algorithm until chance is found to switch to one of the specialized
         unrolling above *)
      let rec fill inc n =
        if n mod 8 = 1 then
          match inc mod 8 with
          | 2         -> fill_2r inc n
          | _ (* 6 *) -> fill_6r inc n
        else if n <= count then
          let i, m = Storage.(index n, mask n) in
          ( flags.|(i) <- flags.|(i) |! m
          ; fill inc (n + inc) )
      in

      let limit = sqrt (float_of_int count) |> int_of_float in
      let rec loop n =
        if n <= limit then
          let i, m = Storage.(index n, mask n) in
          ( (if flags.|(i) |? m then
               fill (n * 2) (n * n))
          ; loop (n + 2) )
        else
          {count; flags}
      in
      loop 3
  end


let t0 = Unix.gettimeofday ()
let t1 = t0 +. time_limit

let rec timer iters =
  let sieve =
    Sys.opaque_identity @@
      Sieve.create 1_000_000
  in
  let t = Unix.gettimeofday () in
  if t > t1 then
    ( report (iters + 1) (t -. t0)
    ; verify (Sieve.is_prime sieve) )
  else
    timer (iters + 1)

;; timer 0