// Let's focus on writing proofs of functions.

// Consider a function that computes the sum of the numbers
// up to and including n. (Such a sum is known as a "triangle number".)
function TriangleNumber(n: int): int
  requires 0 <= n
{
  if n == 0 then 0 else TriangleNumber(n - 1) + n
}

// In a very boring way, let's write a method that does the same
// thing as this function, with a postcondition that says it does indeed
// return the same value as the function.
method ComputeTriangleNumber(n: int) returns (s: int)
  requires 0 <= n  // omit this line; what happens?
  ensures s == TriangleNumber(n)
{
  if n == 0 {
    return 0;
  } else {
    s := ComputeTriangleNumber(n - 1);
    return s + n;
  }
}
// For this method to be correct, every control path leading to the end of
// the method body must be shown to establish the postcondition.

// You may know, as Gauss did, a closed form of triangle numbers. We can
// add that to the postcondition:
method ComputeTriangleNumber2(n: int) returns (s: int)
  requires 0 <= n
  ensures s == TriangleNumber(n)
  ensures s == n * (n + 1) / 2
{
  if n == 0 {
    return 0;
  } else {
    s := ComputeTriangleNumber2(n - 1);
    return s + n;
  }
}
// For this method to be correct, every control path leading to the end of
// the method body must be shown to establish both of these postconditions.

// A consequence of the two postconditions is that TriangleNumber(n) == n * (n + 1) / 2.
// Let's change the postcondition to that:
method ComputeTriangleNumber3(n: int) returns (s: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
  if n == 0 {
    return 0;
  } else {
    s := ComputeTriangleNumber3(n - 1);
    return s + n;
  }
}

// The out-parameter "s" is no longer mentioned in the postcondition. Let's just
// remove it altogether.
method ComputeTriangleNumber4(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
  if n == 0 {
  } else {
    ComputeTriangleNumber4(n - 1);
  }
}
// As before, correctness of this method means that every control path leading to
// the end establishes the postcondition. But why would anyone want to call such
// a method, since it does not return anything? By calling the method, the caller
// learns the postcondition. This is called a lemma!

// A lemma in Dafny is like a method. The postcondition of the lemma is the statement
// of the lemma -- the proof goal of the body of the lemma. The precondition of the
// lemma is the antecedent of the lemma -- every caller must show the precondition
// in order to be allowed to call the lemma and obtain the information in the proof
// goal.
// Here is the same method but declared as a lemma and renamed "Gauss".
lemma Gauss(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
  if n == 0 {
  } else {
    Gauss(n - 1);
  }
}

// The proof has the same structure as you may have encountered in previous classes
// or in high-school geometry. It is a _proof by induction_ with the base case `n == 0`
// and an induction step for `n > 0`. The induction hypothesis is obtained by calling
// the lemma recursively. Let's apply the rules from Class 3 (here, going in the
// forward direction) and add assert statements to show this in more detail.
lemma Gauss2(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
  if n == 0 {
    assert TriangleNumber(n) == 0; // by definition of TriangleNumber
    assert n * (n + 1) / 2 == 0; // by arithmetic
    // Therefore:
    assert TriangleNumber(n) == n * (n + 1) / 2;
  } else {
    var n' := n - 1;
    Gauss2(n');
    assert TriangleNumber(n') == n' * (n' + 1) / 2; // we get this information from the postcondition of the call Gauss(n')
    assert TriangleNumber(n) == TriangleNumber(n') + n; // by definition of TriangleNumber
    // Therefore:
    assert TriangleNumber(n) == n' * (n' + 1) / 2 + n;
    // We also have:
    assert n' * (n' + 1) / 2 + n == (n - 1) * n / 2 + n;
    assert (n - 1) * n / 2 + n == (n - 1) * n / 2 + 2 * n / 2;
    assert (n - 1) * n / 2 + 2 * n / 2 == (n - 1 + 2) * n / 2;
    assert (n - 1 + 2) * n / 2 == n * (n + 1) / 2;
    // And so:
    assert TriangleNumber(n) == n * (n + 1) / 2;
  }
}

// Sometimes, you may have to write out proofs in a lot of detail, maybe to figure out what's missing
// in a proof that does not go through automatically or maybe to give the verifier enough hints about
// the proof structure so that it can complete the proof. When you do, long sequences of assert
// statements (like in Gauss2 above) are hard to read.
// Here are two structuring devices you can use to make proofs easier to read, "assert by" and "calc".
lemma Gauss3(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
  if n == 0 {
    assert TriangleNumber(n) == n * (n + 1) / 2 by {
      assert TriangleNumber(n) == 0; // by definition of TriangleNumber
      assert n * (n + 1) / 2 == 0; // by arithmetic
    }
    
  } else {
    assert TriangleNumber(n - 1) == (n - 1) * n / 2 by {
      var n' := n - 1;
      Gauss3(n');
      assert TriangleNumber(n') == n' * (n' + 1) / 2;
    }

    calc {
      // Note the semicolon at the end of each line -- it can be easy to forget.
      TriangleNumber(n);
    ==  // by definition of TriangleNumber
      TriangleNumber(n - 1) + n;
    ==  // from the "assert by" above
      (n - 1) * n / 2 + n;
    ==  // arithmetic
      (n - 1) * n / 2 + 2 * n / 2;
    ==  // distribute * and +
      ((n - 1) + 2) * n / 2;
    ==  // arithmetic
      (n * (n + 1)) / 2;
    }
  }
}

// Instead of just giving natural-language comments for the steps, you can add
// a pair of curly braces and justify a step by adding some code inside.
lemma Gauss4(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
  if n == 0 {
    // easy
  } else {
    calc {
      TriangleNumber(n);
    ==  // by definition of TriangleNumber
      TriangleNumber(n - 1) + n;
    ==  { Gauss4(n - 1); }  // this recursive call to the lemma obtains the induction hypothesis
      (n - 1) * n / 2 + n;
    ==  { assert n == 2 * n / 2; }
      (n - 1) * n / 2 + 2 * n / 2;
    ==  // arithmetic
      (n * (n + 1)) / 2;
    }
  }
}

// The "calc" statement above is easy to read, but you don't have to supply all the information
// to make the verifier happy. For example, the "==" between line pairs is the default operator
// and can be omitted. Here's the same proof -- shorter, but more cryptic.
lemma Gauss5(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
  if n != 0 { // to make things shorter, the "if" condition was negated and one branch omitted
    calc {
      TriangleNumber(n);
      TriangleNumber(n - 1) + n;
      { Gauss5(n - 1); }
      (n - 1) * n / 2 + n;
      { assert n == 2 * n / 2; }
      (n - 1) * n / 2 + 2 * n / 2;
      (n * (n + 1)) / 2;
    }
  }
}

// Sometimes, many parts of the proof are done automatically, so you don't have to give much
// detail. For example, see the first Gauss lemma above. In fact, for simple proofs like this,
// Dafny applies a smidgen of automatic induction (essentially, by automatically inserting
//     if n > 0 { Gauss(n - 1); }
// into each lemma body). So, this is also a proof:
lemma Gauss6(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{
}

// But what's the fun with that? You wouldn't learn anything. You can turn off automatic
// induction and WE WILL REQUIRE YOU TO DO SO in this week's homework. Here's how you
// turn off automatic induction for a lemma:
lemma {:induction false} Gauss7(n: int)
  requires 0 <= n
  ensures TriangleNumber(n) == n * (n + 1) / 2
{ // error is reported here, saying the verifier cannot find the proof
}

// --------------------------------------------------------

// Now that we know what lemmas are, that lemmas (like methods and functions) can be called
// recursively (which for a lemma has the effect of obtaining what we usually refer to
// as the induction hypothesis), and have seen two basic proof-structuring constructs
// ("assert by" and "calc"), let's write some proofs.

// Let's prove that the square of the sum of the first n positive numbers equals the
// sum of the first n positive cubes.
// Rather than declaring the parameter to be of type "int" and using a precondition that
// says the parameter is non-negative, we can use the type "nat".
function SumCubesThrough(n: nat): nat {
  if n == 0 then 0 else SumCubesThrough(n - 1) + n * n * n
}

lemma Sums(n: nat)
  ensures TriangleNumber(n) * TriangleNumber(n) == SumCubesThrough(n)
{
  if n == 0 {
  } else {
    // For brevity, here are names for things with long names
    var tn := TriangleNumber(n);
    var tn' := TriangleNumber(n - 1);
    var s' := SumCubesThrough(n - 1);
    // Pro tip: Working with non-linear arithmetic can require a lot of patience with the verifier.
    // Sometimes, it can help to name subexpressions, so we introduce the name "nn" here.
    var nn := n * n;
    calc {
      tn * tn;
      (tn' + n) * (tn' + n);
      tn' * tn' + 2 * tn' * n + nn;
      { Sums(n - 1); }
      s' + 2 * tn' * n + nn;
      { Gauss(n - 1); }
      s' + (n - 1) * nn + nn;
      s' + n * nn;
      SumCubesThrough(n);
    }
  }
}

// --------------------------------------------------------

// Our department would lose its accreditation if we didn't also use Fibonacci
// as another example of recursion. ;-)
function Fib(n: nat): nat {
  if n < 2 then n else Fib(n - 2) + Fib(n - 1)
}

// The following example shows that a `calc` does not need to
// use `==` with every step. This one uses `>=` in two of the
// steps.
lemma {:induction false} FibGetsLarger(n: nat)
  requires n >= 5
  ensures Fib(n) >= n
{
  if n == 5 {
    assert Fib(5) == 5;
  } else if n == 6 {
    assert Fib(6) == 8;
  } else {
    calc {
      Fib(n);
      Fib(n - 2) + Fib(n - 1);
    >=  { FibGetsLarger(n - 2); FibGetsLarger(n - 1); }
      n - 2 + n - 1;
    >=  { assert n >= 3; }
      n;
    }
  }
}
// Note, the `assert n >= 3;` above is mainly for human consumption.
// The verifier will check any such assertion to be true. But if the
// verifier can complete the proof without using the assertion as a hint,
// then the asserted condition looks wrong to a human. For example,
// if you change the assertion to `assert n >= 0;` or `assert 20 > 9;`,
// then the verifier will prove the condition, but the condition itself
// is not helpful to the proof of the `calc` step.

// --------------------------------------------------------

// The next two lemmas demonstrate variations in proof style.

// Here is a function that sums the elements of a given sequence.
function Sum(s: seq<int>): int {
  if |s| == 0 then 0 else s[0] + Sum(s[1..])
}

// This function computes the element-wise differences of two
// given sequences and sums up these differences.
function SumOfDifferences(s: seq<int>, t: seq<int>): int
  requires |s| == |t|
{
  if |s| == 0 then 0 else s[0] - t[0] + SumOfDifferences(s[1..], t[1..])
}

// We prove that a `SumOfDifferences` can also be computed by
// calling `Sum`. Like many other similar lemmas, the proof of can be written
// with a `calc` statement (as we'll do in an example below).
// But to show a different way to approach a proof obligation, here is an
// alternative proof style where we start by writing down some things we know
// or believe to be true. Each such thing is written down with an assertion,
// which the verifier checks. Then, we hope that looking at these facts will give us
// an idea of how to use the induction hypothesis to compute the proof.
lemma {:induction false} SequenceMinus(s: seq<int>, t: seq<int>)
  requires |s| == |t|
  ensures SumOfDifferences(s, t) == Sum(s) - Sum(t)
{
  if |s| == 0 {
  } else {
    // Here are some things we know:
    assert SumOfDifferences(s, t) == s[0] - t[0] + SumOfDifferences(s[1..], t[1..]);
    assert Sum(s) == s[0] + Sum(s[1..]);
    assert Sum(t) == t[0] + Sum(t[1..]);
    // By calling the lemma recursively, we can obtain a property that relates
    // the `Sum(_[1..])` terms on the previous two lines.
    assert SumOfDifferences(s[1..], t[1..]) == Sum(s[1..]) - Sum(t[1..]) by {
      SequenceMinus(s[1..], t[1..]);
    }
  }
}

// We can turn the previous proof into a proof calculation.
// When the proof goal of the lemma is an equality, then a typical way to
// write a proof calculation is to start from the "more complicated"
// side of the equality. Sometimes, we may need to start from both ends
// to see if we can connect the pieces.
// Which do you find easier to read or write?
lemma {:induction false} SequenceMinus'(s: seq<int>, t: seq<int>)
  requires |s| == |t|
  ensures SumOfDifferences(s, t) == Sum(s) - Sum(t)
{
  if |s| == 0 {
  } else {
    calc {
      // Here, we start from the LHS of the proof goal
      SumOfDifferences(s, t);
      // apply the definition of `SumOfDifferences`
      s[0] - t[0] + SumOfDifferences(s[1..], t[1..]);
      // In the line above, we have something that looks like a smaller version of the proof goal,
      // so let's call the lemma recursively
      { SequenceMinus'(s[1..], t[1..]); }
      s[0] - t[0] + Sum(s[1..]) - Sum(t[1..]);
      // But now what??

      // Let's try the RHS of the proof goal. The following lines
      // are written from the last line upwards.
      // 3: oh, and that completes the proof
      s[0] + Sum(s[1..]) - t[0] - Sum(t[1..]); // 2: we do the same with `Sum(t)`
      s[0] + Sum(s[1..]) - Sum(t); // 1: then this line, which applies the definition of `Sum(s)`
      Sum(s) - Sum(t); // 0: write this line first
    }
  }
}

// Consider the following problem:
//     every amount of postage that is at least 12 cents can be made from 4-cent and 5-cent stamps
// We can try to prove that using a lemma like this:
lemma PostageStampsExists(amount: int)
  requires 12 <= amount
  ensures exists fours : nat, fives : nat :: 0 <= fours && 0 <= fives && 4 * fours + 5 * fives == amount

// The existential quantifier in the postcondition may look frightening. But there is another way
// to express the property. Very often when you have to prove the existence of something (here,
// nonnegative values for `fours` and `fives`), you end up constructing those somethings. That means
// you can write a method that computes the somethings and returns the somethings in out-parameters.
// In fact, you can do the same for lemmas, because lemmas can have out-parameter, too!
// So, here is a logically equivalent way of stating the lemma, but one that is more straightforward
// to work with.


lemma PostageStamps(amount: int) returns (fours: nat, fives: nat)
  requires 12 <= amount
  ensures 4 * fours + 5 * fives == amount
{
  if amount == 12 {
    fours, fives := 3, 0;
  } else if amount == 13 {
    fours, fives := 2, 1;
  } else if amount == 14 {
    fours, fives := 1, 2;
  } else if amount == 15 {
    fours, fives := 0, 3;
  } else {
    fours, fives := PostageStamps(amount - 4);
    fours := fours + 1;
  }
}

// We can then prove the original formulation with `exists`.
// Dafny's general heuristic for proving "exists" facts is to look for a fact
// you've already concluded that matches the *body* of the `exists` but with
// particular expressions substituted for the quantified variable.  Invoking
// `PostageStamps` brings such a fact into scope, critically along with the
// names `fours` and `fives` for the numbers that exist.
// Don't be fooled by the coincidence of variable names in the `PostageStamps`
// call vs. the function postcondition.  We could switch the body to use names
// `steve` and `henry` instead, and the proof would still go through!
lemma PostageStampsExists_implemented(amount: int)
  requires 12 <= amount
  ensures exists fours : nat, fives : nat :: 0 <= fours && 0 <= fives && 4 * fours + 5 * fives == amount
{
  var fours, fives := PostageStamps(amount);
}

// ---------------------------------------------------------------------------------

// Here's another natural one for sequences.
function Reverse(s: seq<int>): seq<int> {
  if s == [] then [] else Reverse(s[1..]) + [s[0]]
}

// We want to prove that `Reverse` is its own inverse, but an additional lemma
// turns out to be handy, so let's prove it first.

lemma ReverseShuffle(s: seq<int>, x: int)
  ensures Reverse(s + [x]) == [x] + Reverse(s)
{
  if s == [] {
    // let's write out simplifications of the various subexpressions
    assert s + [x] == [x];
    assert Reverse([x]) == [x];
    assert Reverse(s) == [];
    assert [x] + Reverse(s) == [x];
  } else {
    var s' := s + [x];
    calc {
      Reverse(s');
      // definition of Reverse
      Reverse(s'[1..]) + [s'[0]];
      { assert s'[1..] == s[1..] + [x]; }
      Reverse(s[1..] + [x]) + [s'[0]];
      { ReverseShuffle(s[1..], x); }
      [x] + Reverse(s[1..]) + [s'[0]];
      { assert s'[0] == s[0]; }
      [x] + Reverse(s[1..]) + [s[0]];
      // definition of Reverse
      [x] + Reverse(s);
    }
  }
}

lemma ReverseIsItsOwnInverse(s: seq<int>)
  ensures Reverse(Reverse(s)) == s
{
 if s == [] {
  } else {
    calc {
      Reverse(Reverse(s));
      Reverse(Reverse(s[1..]) + [s[0]]);
      { ReverseShuffle(Reverse(s[1..]), s[0]); }
      [s[0]] + Reverse(Reverse(s[1..]));
      { ReverseIsItsOwnInverse(s[1..]); }
      [s[0]] + s[1..];
      s;
    }
  }
}

// ---------------------------------------------------------------------------------

// In this example, we'll use a pair of integers `(x, y)` to represent an
// interval from `x` to, but not including, `y`. For such a pair `p`, its two components
// are denoted `p.0` and `p.1`.
function InInterval(i: int, p: (int, int)): bool {
  p.0 <= i < p.1
}

// A number is in an interval sequence if it is in any of its intervals.
function InIntervalSequence(i: int, s: seq<(int, int)>): bool {
  if |s| == 0 then
    false
  else if InInterval(i, s[0]) then
    true
  else
    InIntervalSequence(i, s[1..])
}

// Before we continue, let's consider an alternative definition of InIntervalSequence.
function InIntervalSequence'(i: int, s: seq<(int, int)>): bool {
  |s| != 0 && (InInterval(i, s[0]) || InIntervalSequence'(i, s[1..]))
}

// Are these two the same definition? Let's prove that they are.
lemma {:induction false} SameIntervalDefinitions(i: int, s: seq<(int, int)>)
  ensures InIntervalSequence(i, s) == InIntervalSequence'(i, s)
{
  // Let's divide up the proof into three cases
  if |s| == 0 {
    assert !InIntervalSequence(i, s);
    assert !InIntervalSequence'(i, s);
  } else if InInterval(i, s[0]) {
    assert InIntervalSequence(i, s);
    assert InIntervalSequence'(i, s);
  } else {
    assert InIntervalSequence(i, s) == InIntervalSequence(i, s[1..]);
    assert InIntervalSequence'(i, s) == InIntervalSequence'(i, s[1..]);
    SameIntervalDefinitions(i, s[1..]);
  }
}

// Let's consider a function that optimizes interval sequences. We won't be terribly
// inventive here, just enough to show a flavor of proving that an optimized version
// of a data structure represents the same numbers as the original.
function Optimize(s: seq<(int, int)>): seq<(int, int)> {
  if s == [] then
    []
  else if s[0].1 <= s[0].0 then
    // interval s[0] is empty, so let's omit it
    Optimize(s[1..])
  else
    [s[0]] + Optimize(s[1..])
}

lemma {:induction false} OptimizationIsCorrect(i: int, s: seq<(int, int)>)
  ensures InIntervalSequence(i, Optimize(s)) == InIntervalSequence(i, s)
{
  // We consider the same three cases as in the Optimize function.
  if s == [] {
    // easy
  } else if s[0].1 <= s[0].0 {
    assert !InInterval(i, s[0]);
    OptimizationIsCorrect(i, s[1..]);
  } else {
    // This is almost too easy.
    // Once we figure out how to apply the induction hypothesis, the automated verifier
    // does all of the heavy lifting.
    OptimizationIsCorrect(i, s[1..]);
  }
}