// Let's assume unknown predicates for evenness and oddness, where we assume
// lemmas to explain what the predicates mean.  This section of the code is
// heavily inspired by Amazon's training materials at:
//   https://dafny.org/teaching-material/Lectures/2-3-Logic-ProvingByExplaining.html

ghost predicate Even(n: nat)
ghost predicate Odd(n: nat)

lemma EvenP()
  ensures forall n: nat :: Even(n) == exists m: nat :: n == 2 * m
lemma OddP()
  ensures forall n: nat :: Odd(n) == !Even(n)

// These definitions will be handy to use to practice proving lemmas with
// interesting usage of logical connectives and quantifiers.

lemma SuccOddIsEven1()
  ensures forall n: nat :: Odd(n) ==> Even(n + 1)
{
  // Note how we call these two assumed lemmas, just to make the "definitions"
  // of the two predicates visible to Dafny.
  EvenP();
  OddP();

  // Then we can step through stating some "obvious" facts, to guide the
  // automated prover toward our goal.
  assert forall n: nat :: n % 2 == 0 ==> n == 2 * (n / 2);
  assert forall n: nat :: n == 2 * (n / 2) ==> Even(n);
  assert forall n: nat :: Odd(n) ==> n % 2 == 1;
  assert forall n: nat :: n % 2 == 1 ==> (n + 1) % 2 == 0;
}

// That was a lot of foralls we had to write!  It's natural to make more
// variables function parameters instead of explicitly quantified.
lemma SuccOddIsEven2(n: nat)
  ensures Odd(n) ==> Even(n + 1)
{
  EvenP();
  OddP();

  assert forall n: nat :: n % 2 == 0 ==> n == 2 * (n / 2);
  assert forall n: nat :: n == 2 * (n / 2) ==> Even(n);
  assert Odd(n) ==> n % 2 == 1;
  assert Odd(n) ==> (n + 1) % 2 == 0;
}

// We also prefer not to write explicit implications, instead using
// preconditions.
lemma SuccOddIsEven3(n: nat)
  requires Odd(n)
  ensures Even(n + 1)
{
  EvenP();
  OddP();
  
  assert forall n: nat :: n % 2 == 0 ==> n == 2 * (n / 2);
  assert forall n: nat :: n == 2 * (n / 2) ==> Even(n);
  assert n % 2 == 1;
  assert (n + 1) % 2 == 0;
}

// The proof gets nicer-looking by splitting out an auxiliary lemma.
lemma DivBy2IsEven()
  ensures forall n: nat :: n % 2 == 0 ==> Even(n)
{
  EvenP(); OddP();
  assert forall n: nat :: n % 2 == 0 ==> n == 2 * (n / 2);
}

lemma SuccOddIsEven4(n: nat)
  requires Odd(n)
  ensures Even(n + 1)
{
  OddP();
  DivBy2IsEven();
  
  assert n % 2 == 1;
  assert (n + 1) % 2 == 0;
}

// Another lemma helps demonstrate support for several of the
// connectives/quantifiers at once.  Note that we can write proofs like this one
// *incrementally* by putting an `assume false` marker at points where the proof
// is unfinished.  Then we can go back and fill them in incrementally.
// However, we're skipping to the complete proof here.
lemma ProveForallImplies()
  ensures forall n: nat :: Even(n) ==> Even(n + 2)
{
  EvenP();
  
  // A `forall` statement is for proving a `forall` fact by introducing a local
  // variable standing for what is quantified over.  We use `ensures` to state a
  // property in terms of that variable, which must be proved by the body of the
  // `forall` statement.
  forall n: nat
    ensures Even(n) ==> Even(n + 2)
  {
    // An implication is proved with an `if` statement.
    if Even(n) {
      // This next tricky little construct is for declaring a variable
      // satisfying a certain logical property, so long as we have convinced
      // Dafny that at least one such value exists.  The traditional name for
      // this operation is *Skolemization*, after Norwegian mathematician
      // Thoralf Skolem.
      var a: nat :| n == 2 * a;

      // Now we need to prove an `exists` fact for the lemma's conclusion.
      // Perhaps surprisingly, just reminding Dafny of the following
      // expression is enough for it to try plugging that value into the
      // `exists`, at which point the proof is complete.
      var b: nat := 2 * (a + 1);
    }
  }
}

// Other classic proof methods can be expressed in programming syntax, too.
// Consider proof by contradiction, where we prove a formula is false by using
// it as an `if` test and then deriving `false`.
lemma ProofByContradiction(n: nat)
  requires Even(n)
  ensures !exists m: nat :: n == 2 * m + 1
{
  EvenP();

  // The following is legal because `n` is known to be `Even`.
  var m1: nat :| n == 2 * m1;

  // Now assume for a contradiction that there does exist `m`.
  if exists m: nat :: n == 2 * m + 1 {
    // Then we can produce a particular value `m2.
    var m2: nat :| n == 2 * m2 + 1;
    // And its existence is contradictory!
    assert false;
  }
}

// We've done some proof by case analysis so far, with both `if` and `match`.
// Generally, when we write a program that exhaustively considers cases, the
// lemma is considered proved when the proof goes through in each case.
lemma ProofByCaseAnalysis(n: nat)
  ensures Even(n) || Odd(n)
{
  EvenP();
  OddP();

  if n % 2 == 0 {
    assert n == 2 * (n / 2);
  } else {
    assert n % 2 == 1;
    assert n == (2 * (n / 2)) + 1;
  }
  // Dafny could actually figure out this particular lemma without the case
  // structure, but it should still be helpful as an example to inspire harder
  // proofs.
}

// Induction proofs can use the same constructs.
lemma {:induction false} ProofByInduction(n: nat)
  requires n % 2 == 0
  ensures Even(n)
{
  EvenP();
  
  if n == 0 {
    assert n == 2 * 0;
  } else {
    assert (n - 2) % 2 == 0;
    ProofByInduction(n - 2);
    var m: nat :| n - 2 == 2 * m;
    assert n == 2 * m + 2;
    assert n == 2 * (m + 1);
  }
}


// Let's explore some of the pragmatics of the `:|` operator a bit more.
// We haven't talked up to this point about *running* Dafny programs, and it is
// not clear at first how to "run" code that materializes values out of thin
// air, given just the mathematical assurance that a value exists!  However,
// that pattern is useful for certain computations on sets.  For instance, let's
// sum up sets of integers
ghost function SetSum_fails(s: set<int>): int {
  if |s| == 0 then
    0
  else
    var n :| n in s;
    // Without the `ghost` keyword above, Dafny complains that this pick isn't
    // uniquely determined.  It certainly isn't, in general!  Dafny reasonably
    // forces executable functions to be deterministic, and `ghost` says this
    // function is for specification purposes only, so nondeterminism is
    // allowed.
    
    n + SetSum_fails(s - {n})
}

// Why do we care so much about determinism?  Check out this failed proof attempt.
lemma SetSumUnion_fails(a: set<int>, b: set<int>)
  ensures SetSum_fails(a + b) == SetSum_fails(a) + SetSum_fails(b) - SetSum_fails(a * b)
{
  if |a + b| == 0 {
  } else {
    // Let's start by mirroring the structure of `SetSum`.
    var n :| n in a + b;
    // So far so good: Dafny accepts that `n` exists.

    assert SetSum_fails(a + b) == n + SetSum_fails(a + b - {n});
    // Wait, why can't Dafny prove this equation, which looks like it's just
    // directly unfolding the definition of `SetSum`?  Remember that `:|`
    // includes no promise of which value it picks, if there are multiple
    // options!  Our `n` here may be different than the `n` in the definition of
    // `SetSum`.
  }
}

// OK, so let's make the behavior very obviously deterministic.
// First, we prove that every (finite) set has a minimum.
lemma SetHasMinimum(s: set<int>)
  requires |s| != 0
  ensures exists min :: min in s && forall n :: n in s ==> min <= n
{
  var n :| n in s;
  if |s - {n}| == 0 {
  } else {
    SetHasMinimum(s - {n});
    assert forall m :: m in s ==> m == n || m in s - {n};
    // This last assertion helps Dafny infer useful case splits.
  }
}

// Now let's wrap minimum-picking in a function.
function SetMinimum(s: set<int>): int
  requires |s| != 0
{
  SetHasMinimum(s); // Note that we get to call our lemma,
                    // to justify existence of a minimum!
  var n :| n in s && forall m :: m in s ==> n <= m;
  n
}

// Back to computing the sum, which is now easy using our wrapper function.
function SetSum(s: set<int>): int {
  if |s| == 0 then
    0
  else
    var n := SetMinimum(s);
    n + SetSum(s - {n})
}

// Now proving our union law is only moderately strenuous.
lemma SetSumUnion(a: set<int>, b: set<int>)
  ensures SetSum(a + b) == SetSum(a) + SetSum(b) - SetSum(a * b)
{
  if |a + b| == 0 {
  } else {
    var n := SetMinimum(a + b);

    assert SetSum(a + b) == n + SetSum(a + b - {n});
    // Great: now the assertion is accepted!

    // To finish the proof, Dafny again needs help choosing useful case splits.
    // We also need to assert the relevant set-theory facts in each case.
    // Dafny is smart enough to prove them but not to realize on its own that
    // they are relevant!
    
    if n in a {
      if n in b {
        assert a + b - {n} == (a - {n}) + (b - {n});
        assert (a - {n}) * (b - {n}) == a * b - {n};
      } else {
        assert a + b - {n} == (a - {n}) + b;
        assert (a - {n}) * b == a * b;
      }
    } else {
      assert a + b - {n} == a + (b - {n});
      assert a * (b - {n}) == a * b;
    }
  }
}

// We can even extract our code to run in conventional programming languages,
// despite the Skolemization.  The following command line produces a new
// directory containing some grungy Python code:
//    dafny translate py class09.dfy --allow-warnings
// Once we get past the complications from optimizations Dafny applied
// automatically, we see it has inserted a loop over all set elements, checking
// for each one whether it is the minimum.  Not so efficient, but at least it
// finds the right answer!


// We wrap up with a final example: permutations.

// Instead of our friend List, we're going to use Dafny's sequences in a functional
// linked-list style. This provides a whole bunch of conveniences, but means we will
// need to remind Dafny of facts like [a, b, c, ...] == [a] + [b, c, ...], and we'll
// constantly be writing [0] and [1..] instead of head and tail.

// We'll also use Map:

function Map<A, B>(f: A -> B, ls: seq<A>): seq<B> {
  if ls == [] then
    []
  else
    [f(ls[0])] + Map(f, ls[1..])
}

// And FlatMap is like Map, but `f` produces any number of output elements for each input
// element, and they are collected into a single list:

function FlatMap<A, B>(f: A -> seq<B>, ls: seq<A>): seq<B> {
  if ls == [] then
    []
  else
    f(ls[0]) + FlatMap(f, ls[1..])
}

method TestFlatMap() {
  assert FlatMap((n) => [2*n, 2*n+1], [0, 1, 2]) == [0, 1, 2, 3, 4, 5];
}

// Calculate all permutations of a list.
function Permutations<T>(ls: seq<T>): seq<seq<T>> {
  if ls == [] then
    [ [] ]
  else
    FlatMap((rp) => AddAtAllPositions(rp, ls[0]), Permutations(ls[1..]))
}

// Find all ways to insert an element at some position in a list.
function AddAtAllPositions<T>(ls: seq<T>, elt: T): seq<seq<T>> {
  if ls == [] then
    [ [elt] ]
  else
    [ [elt] + ls ] + Map((l) => [ls[0]] + l, AddAtAllPositions(ls[1..], elt))
}

method TestPermutations() {
  assert Permutations([1]) == [ [1] ];
  assert Permutations([2, 3]) == [ [2, 3], [3, 2] ] by {
    assert forall a: int, b: int :: [a] + [b] == [a, b];
  }
  assert Permutations([1, 2, 3]) == [ [1, 2, 3], [2, 1, 3], [2, 3, 1],
                                      [1, 3, 2], [3, 1, 2], [3, 2, 1] ] by {
    assert forall a: int, b: int :: [a] + [b] == [a, b];
    assert forall a: int, b: int, c: int :: [a] + [b, c] == [a, b, c];
  }
}

// To prove correctness, we start by practicing the patterns we learned for
// forall and for implication.
lemma PermutationsCorrect<T>(ls: seq<T>)
  ensures forall perm :: perm in Permutations(ls) <==> multiset(perm) == multiset(ls)
{
  forall perm
    ensures perm in Permutations(ls) <==> multiset(perm) == multiset(ls)
  {
    if perm in Permutations(ls) {
      InPermutations_SameElements(ls, perm);
    }
    if multiset(perm) == multiset(ls) {
      SameElements_InPermutations(ls, perm);
    }
  }
}

// First one direction of the bi-implication.
lemma InPermutations_SameElements<T>(ls: seq<T>, perm: seq<T>)
  requires perm in Permutations(ls)
  ensures multiset(perm) == multiset(ls)
{
  if ls == [] {
  } else {
    assert ls == [ls[0]] + ls[1..]; // a necessary reminder to Dafny
    var f := (rp) => AddAtAllPositions(rp, ls[0]);
    assert Permutations(ls) == FlatMap(f, Permutations(ls[1..]));
    // We need a property about FlatMap to reason about its use where. We know that
    // `perm` is in the result of the FlatMap; we would like to know which permutation
    // of `ls[1..]` is comes from (by inserting `ls[0]` at some position).
    FlatMapInputWitness(f, Permutations(ls[1..]), perm);
    // The lemma tells us such an `x` exists:
    assert exists x :: x in Permutations(ls[1..]) && perm in AddAtAllPositions(x, ls[0]);
    // And therefore we can use Skolemization to get a name for it!
    var x :| x in Permutations(ls[1..]) && perm in AddAtAllPositions(x, ls[0]);

    InPermutations_SameElements(ls[1..], x);

    // Finally, we need a correctness property about the `AddAtAllPositions` helper.
    AddAtAllPositionsElements(x, ls[0], perm);
  }
}

lemma FlatMapInputWitness<A, B>(f: A -> seq<B>, ls: seq<A>, b: B)
  requires b in FlatMap(f, ls)
  ensures exists a :: a in ls && b in f(a)
{
}

// By the way, since Dafny's automated prover is easily confused by quantifiers,
// it often helps to find ways to avoid using them, like *lemmas with return
// values*, illustrated here. Automatic induction isn't as effective, unfortunately.
lemma FlatMapInputWitness'<A, B>(f: A -> seq<B>, ls: seq<A>, b: B) returns (a: A)
  requires b in FlatMap(f, ls)
  ensures a in ls && b in f(a)
{
  if b in f(ls[0]) {
    a := ls[0];
  } else {
    a := FlatMapInputWitness'(f, ls[1..], b);
  }
}

// The main lemma (for this direction) about the main helper function.
lemma AddAtAllPositionsElements<T>(ls: seq<T>, elt: T, added: seq<T>)
  requires added in AddAtAllPositions(ls, elt)
  ensures multiset(added) == multiset(ls) + multiset{elt}
{
  if ls == [] {
  } else if added == [elt] + ls {
  } else {
    // This time we need a property about Map: we know `added` is in the result,
    // and we would like to know where it came from.
    assert added in Map((l) => [ls[0]] + l, AddAtAllPositions(ls[1..], elt));
    MapInputWitness((l) => [ls[0]] + l, AddAtAllPositions(ls[1..], elt), added);
    var x :| x in AddAtAllPositions(ls[1..], elt) && added == [ls[0]] + x;
    AddAtAllPositionsElements(ls[1..], elt, x);
    // Dafny has done so much work for us, but now it needs a reminder:
    assert ls == [ls[0]] + ls[1..];
  }
}

lemma MapInputWitness<A, B>(f: A -> B, ls: seq<A>, b: B)
  requires b in Map(f, ls)
  ensures exists a :: a in ls && b == f(a)
{
}

// Now we turn to the other direction of the bi-implication.
lemma SameElements_InPermutations<T>(ls1: seq<T>, ls2: seq<T>)
  requires multiset(ls1) == multiset(ls2)
  ensures ls2 in Permutations(ls1)
{
  if ls1 == [] {
  } else {
    // We eventually want to invoke the induction hypothesis. To do that with `ls1[1..]`,
    // we will need a corresponding `ls2`-but-without-`ls1[0]`. To that end, we define
    // the `RemoveFirst` function below. In addition, we need to prove a correctness
    // property that shows the precondition of the IH is satisfied.
    assert ls1[0] in multiset(ls2); // a necessary reminder to Dafny
    RemoveFirstElements(ls2, ls1[0]); // the lemma about the elements of RemoveFirst(..)
    assert ls1 == [ls1[0]] + ls1[1..]; // and another reminder

    SameElements_InPermutations(ls1[1..], RemoveFirst(ls2, ls1[0]));

    // We have a pretty intuitive interaction between AddAtAllPositions and RemoveFirst
    // that we need to prove, showing that `ls2` is in the result of adding `ls1[0]`
    // back in everywhere after removing it.
    AddAtAllPositionsRemoveFirst(ls2, ls1[0]);
    // Finally, we invoke our second FlatMap lemma to prove `ls2` is in its result.
    FlatMapInputOutput((rp) => AddAtAllPositions(rp, ls1[0]), Permutations(ls1[1..]), RemoveFirst(ls2, ls1[0]), ls2);
  }
}

// Remove the first occurrence of a value from a list.
function RemoveFirst<T(==)>(ls: seq<T>, elt: T): seq<T> {
  if ls == [] then
    []
  else if ls[0] == elt then
    ls[1..]
  else
    [ls[0]] + RemoveFirst(ls[1..], elt)
}

// And here's the key correctness property of that function.
lemma RemoveFirstElements<T>(ls: seq<T>, elt: T)
  requires elt in ls
  ensures multiset(RemoveFirst(ls, elt)) == multiset(ls) - multiset{elt}
{
  assert ls == [ls[0]] + ls[1..];
}

lemma AddAtAllPositionsRemoveFirst<T>(ls: seq<T>, elt: T)
  requires elt in ls
  ensures ls in AddAtAllPositions(RemoveFirst(ls, elt), elt)
{
  if ls == [] {
  } else if elt == ls[0] {
    assert ls == AddAtAllPositions(ls[1..], elt)[0];
  } else {
    // We want to invoke the induction hypothesis; can we use ls[1..]?
    // Just to expand the definition:
    var f := (l) => [RemoveFirst(ls, elt)[0]] + l;
    assert AddAtAllPositions(RemoveFirst(ls, elt), elt) == [ [elt] + RemoveFirst(ls, elt) ] + Map(f, AddAtAllPositions(RemoveFirst(ls, elt)[1..], elt));
    // And we know `elt != ls[0]`.
    assert RemoveFirst(ls, elt)[1..] == RemoveFirst(ls[1..], elt);
    AddAtAllPositionsRemoveFirst(ls[1..], elt); // (would be called automatically)
    // With the result of the IH, we can invoke our second lemma about Map to prove that
    // `ls` is in its result.
    assert ls == [ls[0]] + ls[1..]; // with a necessary reminder
    MapInputOutput(f, AddAtAllPositions(RemoveFirst(ls[1..], elt), elt), ls[1..], ls);
  }
}

lemma MapInputOutput<A, B>(f: A -> B, ls: seq<A>, a: A, b: B)
  requires a in ls
  requires b == f(a)
  ensures b in Map(f, ls)
{
}

lemma FlatMapInputOutput<A, B>(f: A -> seq<B>, ls: seq<A>, a: A, b: B)
  requires a in ls
  requires b in f(a)
  ensures b in FlatMap(f, ls)
{
}