// Recursive Fibonacci, a classic for us to play with using loops
function Fib(n: nat): nat {
  if n < 2 then n else Fib(n - 2) + Fib(n - 1)
}

// Dafny accepts the following, modulo a complaint about a missing loop body, so
// we're off to a good start.
method ComputeFib(n: nat) returns (x: nat)
  ensures x == Fib(n)
{
  x := 0;
  var i := 0;
  while i != n
    invariant 0 <= i <= n
    invariant x == Fib(i)
}

// However, a strengthened invariant helps us fill in an efficient loop body.
method ComputeFib'(n: nat) returns (x: nat)
  ensures x == Fib(n)
{
  x := 0;
  var i, y := 0, 1;
  while i != n
    invariant 0 <= i <= n
    invariant x == Fib(i) && y == Fib(i + 1)
  {
    assert i != n && 0 <= i <= n && x == Fib(i) && y == Fib(i + 1);
    // Let's again jump to the end of the body and reason backwards.

    assert 0 <= i+1 <= n && y == Fib(i + 1) && x + y == Fib(i + 2);
    x, y := y, x + y;

    assert 0 <= i+1 <= n && x == Fib(i + 1) && y == Fib(i + 2);
    i := i + 1;

    assert 0 <= i <= n && x == Fib(i) && y == Fib(i + 1);
  }
}

// OK, let's get more ambitious and compute the squares of Fibonacci numbers.
method SquareFib(N: nat) returns (x: nat)
  ensures x == Fib(N) * Fib(N)
{
  x := 0;
  var n := 0;
  while n != N
    invariant 0 <= n <= N
    invariant x == Fib(n) * Fib(n)
  {
    assert n != N && 0 <= n <= N && x == Fib(n) * Fib(n);

    // Hm, no easy and efficient way to proceed from here, since we don't have
    // Fib(n + 1) precomputed.
    
    assert 0 <= n + 1 <= N && x == Fib(n + 1) * Fib(n + 1);
    n := n + 1;

    assert 0 <= n <= N && x == Fib(n) * Fib(n);
  }
}

// But let's tackle the problem with a technique the book calls "programming by
// wishing."
method SquareFib'(N: nat) returns (x: nat)
  ensures x == Fib(N) * Fib(N)
{
  x := 0;
  var n, y := 0, 1;
  while n != N
    invariant 0 <= n <= N
    invariant x == Fib(n) * Fib(n)
    invariant y == Fib(n + 1) * Fib(n + 1)
  {
    // Let's use calc to figure out how to adjust y efficiently.
    calc {
      Fib(n + 2) * Fib(n + 2);
      == (Fib(n) + Fib(n + 1)) * (Fib(n) + Fib(n + 1));
      == Fib(n) * Fib(n) + 2 * Fib(n) * Fib(n + 1) + Fib(n + 1) * Fib(n + 1);
      == x + 2 * Fib(n) * Fib(n + 1) + y;
    }
    // OK, so if we just had Fib(n) * Fib(n + 1) precomputed, we'd be all set.

    assert 0 <= n + 1 <= N && y == Fib(n + 1) * Fib(n + 1) && y == Fib(n + 2) * Fib(n + 2);
    x := y;
    assert 0 <= n + 1 <= N && x == Fib(n + 1) * Fib(n + 1) && y == Fib(n + 2) * Fib(n + 2);
    n := n + 1;

    assert 0 <= n <= N && x == Fib(n) * Fib(n) && y == Fib(n + 1) * Fib(n + 1);
  }
}

method SquareFib''(N: nat) returns (x: nat)
  ensures x == Fib(N) * Fib(N)
{
  x := 0;
  var n, y, k := 0, 1, 0;
  while n != N
    invariant 0 <= n <= N
    invariant x == Fib(n) * Fib(n)
    invariant y == Fib(n + 1) * Fib(n + 1)
    invariant k == 2 * Fib(n) * Fib(n + 1)
  {
    // Repeated calculation of how to adjust y efficiently
    calc {
      Fib(n + 2) * Fib(n + 2);
      == (Fib(n) + Fib(n + 1)) * (Fib(n) + Fib(n + 1));
      == Fib(n) * Fib(n) + 2 * Fib(n) * Fib(n + 1) + Fib(n + 1) * Fib(n + 1);
      == x + k + y;
    }

    // Ditto for k.
    calc {
      2 * Fib(n + 1) * Fib(n + 2);
      == 2 * Fib(n + 1) * (Fib(n) + Fib(n + 1));
      == 2 * Fib(n + 1) * Fib(n) + 2 * Fib(n + 1) * Fib(n + 1);
      == k + y + y;
    }

    x, y, k := y, x + k + y, k + y + y;
    n := n + 1;
  }
}

// Another recursive function to experiment with
ghost function Power(n: nat): nat {
  if n == 0 then 1 else 2 * Power(n - 1)
}

// We can work through a loop-based implementation backward, by calculating from
// assertions.
method ComputePower(n: nat) returns (p: nat)
  ensures p == Power(n)
{
  p := 1;
  var i := 0;
  while i != n
    invariant 0 <= i <= n
    invariant p == Power(i)
  {
    assert i != n && 0 <= i <= n && p == Power(i);

    assert 0 <= i < n && 2 * p == Power(i + 1);
    p := 2 * p;
    assert 0 <= i < n && p == Power(i + 1);
    i := i + 1;
    
    assert 0 <= i <= n && p == Power(i);
  }
}

// Here's an interesting alternative loop invariant.
method ComputePower'(n: nat) returns (p: nat)
  ensures p == Power(n)
{
  p := 1;
  var i := 0;
  while i != n
    invariant 0 <= i <= n
    invariant p * Power(n - i) == Power(n)
  {
    assert i != n && 0 <= i <= n && p * Power(n - i) == Power(n);

    assert 0 <= i < n && 2 * p * Power(n - i - 1) == Power(n);
    p := 2 * p;
    assert 0 <= i < n && p * Power(n - i - 1) == Power(n);
    i := i + 1;

    assert 0 <= i <= n && p * Power(n - i) == Power(n);
  }
}
// We used an important general technique above: write a loop invariant
// explaining how current variable values need to be transformed to produce
// final answers.

// Let's now reason about sums over an unknown function.
function F(x: int): int

// Here are two alternative ways to describe a summation, given that addition is
// associative and commutative.
  
ghost function SumUp(lo: int, hi: int): int
  requires lo <= hi
  decreases hi - lo
{
  if lo == hi then 0 else F(lo) + SumUp(lo + 1, hi)
}

ghost function SumDown(lo: int, hi: int): int
  requires lo <= hi
  decreases hi - lo
{
  if lo == hi then 0 else SumDown(lo, hi - 1) + F(hi - 1)
}

// They really are equivalent.
lemma SameSums(lo: int, hi: int)
  requires lo <= hi
  ensures SumUp(lo, hi) == SumDown(lo, hi)
  decreases hi - lo
{
  if lo != hi {
    PrependSumDown(lo, hi);
  }
}

lemma PrependSumDown(lo: int, hi: int)
  requires lo < hi
  ensures F(lo) + SumDown(lo + 1, hi) == SumDown(lo, hi)
  decreases hi - lo
{
  if lo + 1 == hi {
  } else {
    calc {
      F(lo) + SumDown(lo + 1, hi);
      == F(lo) + SumDown(lo + 1, hi - 1) + F(hi - 1);
      == { PrependSumDown(lo, hi - 1); } SumDown(lo, hi - 1) + F(hi - 1);
      == SumDown(lo, hi);
    }
  }
}

// Let's try to compute SumUp with a loop.
method LoopUp(lo: int, hi: int) returns (s: int)
  requires lo <= hi
  ensures s == SumUp(lo, hi)
{
  s := 0;
  var i := lo;
  while i != hi
    invariant lo <= i <= hi
    invariant s == SumUp(lo, i)
  {
    assert s == SumUp(lo, i);

    AppendSumUp(lo, i);
    // At this point, we realize the need for a lemma, invoked just above.
    assert s + F(i) == SumUp(lo, i + 1);
    s := s + F(i);
    assert s == SumUp(lo, i + 1);
    i := i + 1;
    assert s == SumUp(lo, i);
  }
}

lemma AppendSumUp(lo: int, hi: int)
  requires lo <= hi
  ensures SumUp(lo, hi + 1) == SumUp(lo, hi) + F(hi)
  decreases hi - lo
{
  if lo < hi {
    AppendSumUp(lo+1, hi);
  }
}

// Here's another approach: LoopUp's invariant is naturally written in terms of
// SumDown, not SumUp!
method LoopUp'(lo: int, hi: int) returns (s: int)
  requires lo <= hi
  ensures s == SumUp(lo, hi)
{
  s := 0;
  var i := lo;
  while i != hi
    invariant lo <= i <= hi
    invariant s == SumDown(lo, i)
  {
    s := s + F(i);
    i := i + 1;
  }

  // At the end, we recall equivalence between SumUp and SumDown.
  SameSums(lo, hi);
}

// Finally, yet another approach, using the pattern of characterizing what work
// is left to be done, to transform intermediate values into final answers.
method LoopUp''(lo: int, hi: int) returns (s: int)
  requires lo <= hi
  ensures s == SumUp(lo, hi)
{
  s := 0;
  var i := lo;
  while i != hi
    invariant lo <= i <= hi
    invariant s + SumUp(i, hi) == SumUp(lo, hi)
  {
    s := s + F(i);
    i := i + 1;
  }
}

// Let's see how this same pattern works with inductive datatypes.
datatype List<T> = Nil | Cons(hd: T, tl: List<T>)

function Length<T>(ls: List<T>): nat {
  match ls
  case Nil => 0
  case Cons(_, ls') => 1 + Length(ls')
}

method LengthLoop<T>(ls: List<T>) returns (sum: nat)
  ensures sum == Length(ls)
{
  var l := ls;
  sum := 0;
  while l.Cons?
    invariant sum + Length(l) == Length(ls)
    // Following the pattern of "explain what additional work needs to be done,
    // to produce the final answer"
    decreases Length(l)
    // Interesting use of the target function to show termination, too!
  {
    sum := sum + 1;
    l := l.tl;
  }
}

predicate Member<T(==)>(x: T, ls: List<T>) {
  match ls
  case Nil => false
  case Cons(y, ls') => y == x || Member(x, ls')
}

method MemberLoop<T(==)>(x: T, ls: List<T>) returns (b: bool)
  ensures b == Member(x, ls)
{
  var l := ls;
  while l.Cons?
    invariant Member(x, l) == Member(x, ls)
    // The invariant is simpler, because the transformed current problem
    // instance is exactly equivalent to the original instance.
    decreases Length(l)
  {
    if l.hd == x {
      return true;
    } else {
      l := l.tl;
    }
  }

  return false;
}

// So lists weren't so bad to work with.  How about binary trees?

datatype Tree = Leaf | Node(left: Tree, data: int, right: Tree)

predicate TreeMember(x: int, t: Tree) {
  match t
  case Leaf => false
  case Node(left, data, right) => TreeMember(x, left) || data == x || TreeMember(x, right)
}

// What makes a tree a valid binary search tree?
ghost predicate IsBst(t: Tree) {
  match t
  case Leaf => true
  case Node(left, data, right) =>
    IsBst(left)
    && IsBst(right)
    && (forall n :: TreeMember(n, left) ==> n < data)
    && (forall n :: TreeMember(n, right) ==> n > data)
}

// Given that property, we can prove an efficient way to evaluate membership
// checks.
function BstLookup(x: int, t: Tree): (b: bool)
  requires IsBst(t)
  ensures b == TreeMember(x, t)
{
  match t
  case Leaf => false
  case Node(left, data, right) =>
    if x == data then
      true
    else if x < data then
      BstLookup(x, left)
    else
      BstLookup(x, right)
}

// And it's quite straightforward to redo that last function as a loop.
method BstLookupLoop(x: int, t: Tree) returns (b: bool)
  requires IsBst(t)
  ensures b == TreeMember(x, t)
{
  var tr := t;

  while tr.Node?
    invariant IsBst(tr)
    invariant BstLookup(x, tr) == TreeMember(x, t)
    // Again, the reduced instance of the problem is known to have the same
    // answer as the original instance.
    decreases tr
    // Note: we could have written "decreases l" in the prior examples, too.
  {
    if x == tr.data {
      return true;
    } else if x < tr.data {
      tr := tr.left;
    } else {
      tr := tr.right;
    }
  }

  return false;
}

// This mysterious pair of functions is the heart of the termination argument
// for our next example.

ghost function Size(t: Tree): nat {
  match t
  case Leaf => 0
  case Node(left, _, right) => Size(left) + 2 + Size(right)
}

ghost function Sizes(ls: List<Tree>): nat {
  match ls
  case Nil => 0
  case Cons(hd, tl) => Size(hd) + 1 + Sizes(tl)
}

// We also define what it means for a data value to be found within a list of
// trees.
ghost predicate TreesMember(x: int, trs: List<Tree>) {
  exists t :: Member(t, trs) && TreeMember(x, t)
}

// Now we see why it was worth defining all of the above.  When we *don't* know
// that a tree is a binary search tree, we need a less "linear" membership
// check, which we handle through explicitly representing a *stack* of pending
// trees to search.
method TreeMemberLoop(x: int, t: Tree) returns (b: bool)
  ensures b == TreeMember(x, t)
{
  var trs := Cons(t, Nil);
  // Here is our stack, which mirrors the call stack in traditional
  // implementations of recursive functions.  For instance, a Python interpreter
  // internally maintains a call stack but hides it from the programmer.
  // Initially, we have a single "call" on the stack, on the original tree.

  TreesMemberBase(x, t);
  // This lemma invocation shows that the loop invariant holds on entry.

  while trs.Cons?
    invariant TreeMember(x, t) == TreesMember(x, trs)
    // Another pleasingly simple invariant, as the current problem instance has
    // the same answer as the original.
    decreases Sizes(trs)
    // And here is where we use our counterintuitive size function.  See below
    // for why it's the right choice.
  {
    ghost var old_trs := trs;
    var tr := trs.hd;
    trs := trs.tl;

    // So we pop tr off of the stack and need to handle it properly.
    // If it is a Leaf, there is nothing more to do.  However....
    if tr.Node? {
      // OK, the node we're handling is nonempty.
      if x == tr.data {
        // It happens to contain x at its root.  Done, after adding some
        // assertions to help Dafny use the induction hypothesis by proving the
        // "exists" inside the definition of TreesMember.
        assert Member(tr, Cons(tr, trs));
        assert TreesMember(x, Cons(tr, trs));
        return true;
      } else {
        // Didn't find x yet?  Queue up "recursive calls" to the left and right
        // children.  That is, push them onto the stack.
        trs := Cons(tr.left, Cons(tr.right, trs));

        assert Sizes(trs) < Sizes(old_trs);
        // Here is where we see why we picked a good definition for Sizes.
        // Every element of the outer list adds 1 to the total, where every Node
        // within a tree adds 2.  Why the asymmetry?  Because, in this case, we:
        // - *Increase* the length of the outer list by one.
        // - *Decrease* the total number of Nodes by one.
        // We need these two to cancel out and produce a net lower size.

        // Finally, a little calculation explains how the new problem instance
        // has the same answer as the previous one.
        calc {
          TreesMember(x, trs);
          == TreesMember(x, Cons(tr.left, Cons(tr.right, old_trs.tl)));
          == { TreesMemberCons(x, tr.left, Cons(tr.right, old_trs.tl));
               TreesMemberCons(x, tr.right, old_trs.tl); }
               TreeMember(x, tr.left) || TreeMember(x, tr.right) || TreesMember(x, old_trs.tl);
          == TreeMember(x, old_trs.hd) || TreesMember(x, old_trs.tl);
          == TreesMember(x, Cons(old_trs.hd, old_trs.tl));
          == TreesMember(x, old_trs);
        }
      }
    }
  }

  return false;
}

lemma TreesMemberBase(x: int, tr: Tree)
  ensures TreesMember(x, Cons(tr, Nil)) == TreeMember(x, tr)
{
}

lemma TreesMemberCons(x: int, tr: Tree, trs: List<Tree>)
  ensures TreesMember(x, Cons(tr, trs)) == TreeMember(x, tr) || TreesMember(x, trs) 
{  
}