
(** {1 Knuth's algorithm for prime numbers}

   The Art of Computer Programming, vol 1, page 147.

   The following code computes a table of the first `m` prime numbers.
   Though there are more efficient ways of computing prime numbers,
   the nice thing about this code is that it requires not less than
   Bertrand's postulate (for `n > 1`, there is always a prime `p` such that
   `n < p < 2n`) to be proved correct.

   This program had been proved correct using Coq and an early version of
   Why back in 2003, by Laurent Théry (INRIA Sophia-Antipolis): Laurent Théry,
   Proving Pearl: Knuth's Algorithm for Prime Numbers, TPHOLs 2003.
   Truly a tour de force, this proof includes the full proof of Bertrand's
   postulate in Coq. Here, we simply focus on the program verification part,
   posing Bertrand's postulate as a lemma that we do not prove.

   Note: Knuth's code is entering the loop where a new prime number is
   added, thus resulting into the immediate addition of 3 as `prime[1]`.
   It allows subsequent division tests to start at `prime[1]=3`, thus
   saving the division by `prime[0]=2`. We did something similar in the
   code below, though the use of a while loop looks like we did a
   special case for 3 as well.  *)

module PrimeNumbers

  use int.Int
  use int.ComputerDivision
  use number.Parity
  use number.Divisibility
  use number.Prime
  use ref.Refint
  use map.Map

  predicate no_prime_in (l u: int) =
    forall x: int. l < x < u -> not (prime x)

  (** `p[0]..p[u-1]` are the first u prime numbers *)
  predicate first_primes (p: int -> int) (u: int) =
    p[0] = 2 /\
    (* sorted *)
    (forall i j: int. 0 <= i < j < u -> p[i] < p[j]) /\
    (* only primes *)
    (forall i: int. 0 <= i < u -> prime p[i]) /\
    (* all primes *)
    (forall i: int. 0 <= i < u-1 -> no_prime_in p[i] p[i+1])

  lemma exists_prime:
    forall p: int -> int, u: int. 1 <= u -> first_primes p u ->
    forall d: int. 2 <= d <= p[u-1] -> prime d ->
    exists i: int. 0 <= i < u /\ d = p[i]

  (** Bertrand's postulate, admitted as an axiom
   (the label is there to suppress the warning issued by Why3) *)
  axiom Bertrand_postulate [@W:non_conservative_extension:N] :
    forall p: int. prime p -> not (no_prime_in p (2*p))

  use array.Array

  (** `prime_numbers m` returns an array containing the first `m`
  prime numbers *)
  let prime_numbers (m: int) : array int
    requires { m >= 2 }
    ensures  { result.length = m }
    ensures  { first_primes result.elts m }
  = let p = make m 0 in
    p[0] <- 2;
    p[1] <- 3;
    let n = ref 5 in (* candidate for next prime *)
    for j = 2 to m - 1 do
      invariant { first_primes p.elts j }
      invariant { p[j-1] < !n < 2*p[j-1] }
      invariant { odd !n }
      invariant { no_prime_in p[j-1] !n }
      let rec test (k: int) variant { 2*p[j-1] - !n, j - k }
        requires { 1 <= k < j }
        requires { first_primes p.elts j }
        requires { p[j-1] < !n < 2*p[j-1] }
        requires { odd !n }
        requires { no_prime_in p[j-1] !n }
        requires { forall i: int. 0 <= i < k -> not (divides p[i] !n) }
        ensures  { p[j-1] < !n /\ prime !n /\ no_prime_in p[j-1] !n }
      = if mod !n p[k] = 0 then begin
          assert { not (prime !n) }; n += 2; test 1
        end else if div !n p[k] > p[k] then
          test (k + 1)
        else
          assert { prime !n }
      in
      test 1;
      p[j] <- !n;
      n += 2
    done;
    p

  exception BenchFailure

  let bench () raises { BenchFailure -> true } =
     let t = prime_numbers 100 in
     if t[0] <> 2 then raise BenchFailure;
     if t[1] <> 3 then raise BenchFailure;
     if t[2] <> 5 then raise BenchFailure;
     if t[3] <> 7 then raise BenchFailure;
     if t[4] <> 11 then raise BenchFailure;
     if t[5] <> 13 then raise BenchFailure;
     if t[6] <> 17 then raise BenchFailure;
     if t[7] <> 19 then raise BenchFailure;
     if t[8] <> 23 then raise BenchFailure;
     if t[9] <> 29 then raise BenchFailure;
     t

end

