// RUN: %dafny /compile:0 /dprint:"%t.dprint" "%s" > "%t"
// RUN: %diff "%s.expect" "%t"

predicate IsPrime(n: int)
{
  2 <= n && forall m :: 2 <= m < n ==> n % m != 0
}

// The following theorem shows that there is an infinite number of primes
lemma AlwaysMorePrimes(k: int)
  ensures exists p :: k <= p && IsPrime(p);
{
  var j, s := 0, {};
  while true
    invariant AllPrimes(s, j);
    decreases k - j;
  {
    var p := GetLargerPrime(s, j);
    if k <= p { return; }
    j, s := p, set x | 2 <= x <= p && IsPrime(x);
  }
}

// Here is an alternative formulation of the theorem
lemma NoFiniteSetContainsAllPrimes(s: set<int>)
  ensures exists p :: IsPrime(p) && p !in s;
{
  AlwaysMorePrimes(if s == {} then 0 else PickLargest(s) + 1);
}

// ------------------------- lemmas and auxiliary definitions

predicate AllPrimes(s: set<int>, bound: int)
{
  // s contains only primes
  (forall x :: x in s ==> IsPrime(x)) &&
  // every prime up to "bound" is included in s
  (forall p :: IsPrime(p) && p <= bound ==> p in s)
}

lemma GetLargerPrime(s: set<int>, bound: int) returns (p: int)
  requires AllPrimes(s, bound);
  ensures bound < p && IsPrime(p);
{
  var q := product(s);
  if exists p :: bound < p <= q && IsPrime(p) {
    p :| bound < p <= q && IsPrime(p);
  } else {
    ProductPlusOneIsPrime(s, q);
    p := q+1;
    if p <= bound {  // by contradction, establish bound < p
      assert p in s;
      product_property(s);
      assert false;
    }
  }
}

function product(s: set<int>): int
{
  if s == {} then 1 else
  var a := PickLargest(s); a * product(s - {a})
}

lemma product_property(s: set<int>)
  requires forall x :: x in s ==> 1 <= x;
  ensures 1 <= product(s) && forall x :: x in s ==> x <= product(s);
{
}

lemma ProductPlusOneIsPrime(s: set<int>, q: int)
  requires AllPrimes(s, q) && q == product(s);
  ensures IsPrime(q+1);
{
  var p := q+1;
  calc {
    true;
    { product_property(s); }
    2 <= p;
  }

  forall m | 2 <= m <= q && IsPrime(m)
    ensures p % m != 0;
  {
    assert m in s;  // because AllPrimes(s, q) && m <= q && IsPrime(m)
    RemoveFactor(m, s);
    var l := product(s-{m});
    assert m*l == q;
    MulDivMod(m, l, q, 1);
  }
  assert IsPrime_Alt(q+1);
  AltPrimeDefinition(q+1);
}

// The following lemma is essentially just associativity and commutativity of multiplication.
// To get this proof through, it is necessary to know that if x!=y and y==Pick...(s), then
// also y==Pick...(s - {x}).  It is for this reason that we use PickLargest, instead of
// picking an arbitrary element from s.
lemma RemoveFactor(x: int, s: set<int>)
  requires x in s;
  ensures product(s) == x * product(s - {x});
{
  var y := PickLargest(s);
  if x != y {
    calc {
      product(s);
      y * product(s - {y});
      { RemoveFactor(x, s - {y}); }
      y * x * product(s - {y} - {x});
      x * y * product(s - {y} - {x});
      { assert s - {y} - {x} == s - {x} - {y}; }
      x * y * product(s - {x} - {y});
      { assert y == PickLargest(s - {x}); }
      x * product(s - {x});
    }
  }
}

// This definition is like IsPrime above, except that the quantification is only over primes.
predicate IsPrime_Alt(n: int)
{
  2 <= n && forall m :: 2 <= m < n && IsPrime(m) ==> n % m != 0
}

// To show that n is prime, it suffices to prove that it satisfies the alternate definition
lemma AltPrimeDefinition(n: int)
  requires IsPrime_Alt(n);
  ensures IsPrime(n);
{
  forall m | 2 <= m < n
    ensures n % m != 0;
  {
    if !IsPrime(m) {
      var a, b := Composite(m);
      if n % m == 0 {  // proof by contradiction
        var k := n / m;
        calc {
          true;
          k == n / m;
          m * k == n;
          a * b * k == n;
        ==>  { MulDivMod(a, b*k, n, 0); }
          n % a == 0;
        ==>  // IsPrime_Alt
          !(2 <= a < n && IsPrime(a));
          { assert 2 <= a < m < n; }
          !IsPrime(a);
          false;
        }
      }
    }
  }
}

lemma Composite(c: int) returns (a: int, b: int)
  requires 2 <= c && !IsPrime(c);
  ensures 2 <= a < c && 2 <= b && a * b == c;
  ensures IsPrime(a);
{
  calc {
    true;
    !IsPrime(c);
    !(2 <= c && forall m :: 2 <= m < c ==> c % m != 0);
    exists m :: 2 <= m < c && c % m == 0;
  }
  a :| 2 <= a < c && c % a == 0;
  b := c / a;
  assert 2 <= a < c && 2 <= b && a * b == c;
  if !IsPrime(a) {
    var x, y := Composite(a);
    a, b := x, y*b;
  }
}

function PickLargest(s: set<int>): int
  requires s != {};
{
  LargestElementExists(s);
  var x :| x in s && forall y :: y in s ==> y <= x;
  x
}

lemma LargestElementExists(s: set<int>)
  requires s != {};
  ensures exists x :: x in s && forall y :: y in s ==> y <= x;
{
  var s' := s;
  while true
    invariant s' != {} && s' <= s;
    invariant forall x,y :: x in s' && y in s - s' ==> y <= x;
    decreases s';
  {
    var x :| x in s';  // pick something
    if forall y :: y in s' ==> y <= x {
      // good pick
      return;
    } else {
      // constrain the pick further
      var y :| y in s' && x < y;
      s' := set z | z in s && x < z;
      assert y in s';
    }
  }
}

// This axiom about % is needed.  Unfortunately, Z3 seems incapable of proving it.
lemma MulDivMod(a: nat, b: nat, c: nat, j: nat)
  requires a * b == c && j < a;
  ensures (c+j) % a == j;