// Here's a simple container data structure: linked lists of natural numbers.
// You've seen them in the reading that was just due, so let's express past a
// quick recap.
datatype NatList = NNil | NCons(head: nat, tail: NatList)
// The 'datatype' command is for declaring an *inductive type*,
// by listing out the *constructor* functions that can be used to construct values in that type.
// Note that the 'Cons' constructor is *recursive*, mentioning the same type.
// Its name stands for "construct," and as for why that's a good name for this
// constructor, it depends on the history of functional programming, where the
// Lisp language introduced this convention!

// Here is how we build a decreasing list of the first 'n' positive numbers.
function FirstN(n: nat): NatList {
  if n == 0 then NNil else NCons(n, FirstN(n-1))
}

// We can also pull apart values of inductive types in recursive functions.
// Let's add up all the numbers in a list.
function Sum(ls: NatList): nat {
  if ls == NNil then 0 else ls.head + Sum(ls.tail)
}

// These definitions are a natural fit for lemmas proved by induction.
// Now we'll let ourselves take full advantage of Dafny's automatic induction!
lemma Sum_FirstN(n: nat)
  ensures Sum(FirstN(n)) == n * (n + 1) / 2
{
}
// Cool!  Proved automatically.

// These ideas around inductive datatypes come from a long tradition of
// *functional programming* in languages like Haskell.
// Dafny also supports (and we often prefer) a *pattern-matching* notation more
// in-line with that tradition.
function Sum'(ls: NatList): nat {
  match ls
  case NNil => 0
  case NCons(n, ls') => n + Sum'(ls')
}

lemma Sum_Sum'(ls: NatList)
  ensures Sum(ls) == Sum'(ls)
{
}

lemma Sum_Sum'_FirstN(n: nat)
  ensures Sum'(FirstN(n)) == n * (n + 1) / 2
{
}
// Yeah, Dafny can prove that one automatically, too.
// We're often doing more explicit proofs for didactic purposes only.
// Soon enough we will get to examples tough enough that Dafny requires help!

// You probably noticed that nothing about the definition of lists is so
// specialized to numbers as data.  Indeed, it is common to implement *generic*
// container types, able to accept any kinds of data inside.  Here is a generic
// version, with a *type parameter* 'T'.  (We are still in recap content from
// the last reading.)
datatype List<T> = Nil | Cons(head: T, tail: List<T>)

// Let's now tour through some classic list operations.
// First, how many elements does a list have?
function Length<T>(xs: List<T>): nat {
  match xs
  case Nil => 0
  case Cons(_, tail) => 1 + Length(tail)
}

// Next, we can use a version of 'Cons' that adds a new element at the *end*
// of a list, instead of the start.
function Snoc<T>(xs: List<T>, y: T): List<T> {
  match xs
  case Nil => Cons(y, Nil)
  case Cons(x, tail) => Cons(x, Snoc(tail, y))
}

// How do the last two functions interact?
lemma LengthSnoc<T>(xs: List<T>, y: T)
  ensures Length(Snoc(xs, y)) == Length(xs) + 1
{
}

// The next classic is appending two lists together.
function Append'<T>(xs: List<T>, ys: List<T>): List<T> {
  match xs
  case Nil => ys
  case Cons(x, tail) => Cons(x, Append'(tail, ys))
}

// It is easy to capture the interaction with 'Length'.
lemma LengthAppend<T>(xs: List<T>, ys: List<T>)
  ensures Length(Append'(xs, ys)) == Length(xs) + Length(ys)
{
}

// However, we could also build this property into the original definition as a
// postcondition!
function Append<T>(xs: List<T>, ys: List<T>): List<T>
  ensures Length(Append(xs, ys)) == Length(xs) + Length(ys)
{
  match xs
  case Nil => ys
  case Cons(x, tail) => Cons(x, Append(tail, ys))
}
// While Dafny methods become *opaque* after they are finished, and thus we only
// know their properties that were recorded in their specifications, we have
// flexibility with functions, whose definitions remain *transparent* afterward.
// The benefit of baking in properties is that Dafny is automatically "reminded"
// of them each time we call the function, saving us from needing to invoke
// lemmas manually.  The downside is that a function definition can get very
// crowded, if we need to include every relevant property!

// Let's prove some more algebraic properties of 'Append'.

lemma AppendNil<T>(xs: List<T>)
  ensures Append(xs, Nil) == xs
{
}

lemma AppendAssociative<T>(xs: List<T>, ys: List<T>, zs:List<T>)
  ensures Append(Append(xs, ys), zs) == Append(xs, Append(ys, zs))
{
}

lemma SnocAppend<T>(xs: List<T>, y: T)
  ensures Snoc(xs, y) == Append(xs, Cons(y, Nil))
{
}

// And some other classic list functions:

function Take<T>(xs: List<T>, n: nat): List<T>
  requires n <= Length(xs)
{
  if n == 0 then Nil else Cons(xs.head, Take(xs.tail, n - 1))
}

function Drop<T>(xs: List<T>, n: nat): List<T>
  requires n <= Length(xs)
{
  if n == 0 then xs else Drop(xs.tail, n - 1)
}

lemma AppendTakeDrop<T>(xs: List<T>, n: nat)
  requires n <= Length(xs)
  ensures Append(Take(xs, n), Drop(xs, n)) == xs
{
}

lemma TakeDropAppend<T>(xs: List<T>, ys: List<T>)
  ensures Take(Append(xs, ys), Length(xs)) == xs
  ensures Drop(Append(xs, ys), Length(xs)) == ys
{
}
// Note that Dafny rejects the postconditions here if we use the version of
// 'Append' without a postcondition asserting the length of the resulting list!

function At<T>(xs: List<T>, i: nat): T
  requires i < Length(xs)
{
  if i == 0 then xs.head else At(xs.tail, i - 1)
}

lemma AtDropHead<T>(xs: List<T>, i: nat)
  requires i < Length(xs)
  ensures Drop(xs, i).Cons? && At(xs, i) == Drop(xs, i).head
{
}
// Note we used a question-mark operation to confirm which list constructor was
// used to build a certain value.  Otherwise, Dafny complains that the
// postcondition is ill-formed ("may crash").

lemma AtAppend<T>(xs: List<T>, ys: List<T>, i: nat)
  requires i < Length(Append(xs, ys))
  ensures At(Append(xs, ys), i)
       == if i < Length(xs) then
            At(xs, i)
          else  
            At(ys, i - Length(xs))
{
}
          
// At what position in a list does an element appear?
// And what do we do if the element is not found?
function Find<T(==)>(xs: List<T>, y: T): nat
  ensures Find(xs, y) <= Length(xs)
{
  match xs
  case Nil => 0
  case Cons(x, tail) =>
    if x == y then 0 else 1 + Find(tail, y)
}
// Note the '(==)' annotation on the type parameter, to require it to support a
// particular operator.

// Answer to question just above: we return the list length when the element is
// not found.
lemma AtFind<T>(xs: List<T>, y: T)
  ensures Find(xs, y) == Length(xs) || At(xs, Find(xs, y)) == y
{
}

lemma BeforeFind<T>(xs: List<T>, y: T, i: nat)
  ensures i < Find(xs, y) ==> At(xs, i) != y
{
}

lemma FindAppend<T>(xs: List<T>, ys: List<T>, y: T)
  ensures Find(xs, y) == Length(xs) || Find(Append(xs, ys), y) == Find(xs, y)
{
}

lemma FindDrop<T>(xs: List<T>, y: T, i: nat)
  ensures i <= Find(xs, y) ==> Find(xs, y) == Find(Drop(xs, i), y) + i
{
}

function SlowReverse<T>(xs: List<T>): List<T> {
  match xs
  case Nil => Nil
  case Cons(x, tail) => Snoc(SlowReverse(tail), x)
}

// Surprisingly, the following isn't proved automatically!
lemma LengthSlowReverse<T>(xs: List<T>)
  ensures Length(SlowReverse(xs)) == Length(xs)
{
  match xs
  case Nil =>
  case Cons(x, tail) => LengthSnoc(SlowReverse(tail), x);
}

// There is an asymptotically faster version of reverse that appears in most
// standard libraries of functional languages.
function ReverseAux<T>(xs: List<T>, acc: List<T>): List<T>
{
  match xs
  case Nil => acc
  case Cons(x, tail) => ReverseAux(tail, Cons(x, acc))
}

function Reverse<T>(xs: List<T>): List<T> {
  ReverseAux(xs, Nil)
}

lemma ReverseAuxSlowCorrect<T>(xs: List<T>, acc: List<T>)
  ensures ReverseAux(xs, acc) == Append(SlowReverse(xs), acc)
{
  match xs
  case Nil =>
  case Cons(x, tail) =>
    calc {
      Append(SlowReverse(xs), acc);
      == Append(Snoc(SlowReverse(tail), x), acc);
      == { SnocAppend(SlowReverse(tail), x); } Append(Append(SlowReverse(tail), Cons(x, Nil)), acc);
      == { AppendAssociative(SlowReverse(tail), Cons(x, Nil), acc); } Append(SlowReverse(tail), Append(Cons(x, Nil), acc));
      == { assert Append(Cons(x, Nil), acc) == Cons(x, acc); } Append(SlowReverse(tail), Cons(x, acc));
      == { ReverseAuxSlowCorrect(tail, Cons(x, acc)); } ReverseAux(tail, Cons(x, acc));
      == ReverseAux(xs, acc);
    }
}

lemma ReverseCorrect<T>(xs: List<T>)
  ensures Reverse(xs) == SlowReverse(xs)
{
  calc {
    Reverse(xs);
    == ReverseAux(xs, Nil);
    == { ReverseAuxSlowCorrect(xs, Nil); } Append(SlowReverse(xs), Nil);
    == { AppendNil(SlowReverse(xs)); } SlowReverse(xs);
  }
}

lemma ReverseAuxCorrect<T>(xs: List<T>, acc: List<T>)
  ensures ReverseAux(xs, acc) == Append(Reverse(xs), acc)
{
  ReverseCorrect(xs);
  ReverseAuxSlowCorrect(xs, acc);
}

lemma LengthReverse<T>(xs: List<T>)
  ensures Length(Reverse(xs)) == Length(xs)
{
  ReverseCorrect(xs);
  LengthSlowReverse(xs);
}

lemma ReverseAuxAppend<T>(xs: List<T>, ys: List<T>, acc: List<T>)
  ensures ReverseAux(Append(xs, ys), acc) == Append(Reverse(ys), ReverseAux(xs, acc))
{
  match xs
  case Nil => ReverseAuxCorrect(ys, acc);
  case Cons(_, _) =>
}

// OK, now for something completely different!
// Let's do some number theory, explaining the natural numbers and their key
// operations from first principles.

datatype Unary = Zero | Suc(pred: Unary)
// That's 'Suc' for "successor" and 'pred' for "predecessor."
// These are the *Peano numbers*: every natural is either zero or the successor
// of another natural.

// We can convert back to the built-in type.
function UnaryToNat(x: Unary): nat {
  match x
  case Zero => 0
  case Suc(x') => 1 + UnaryToNat(x')
}

// And we can convert in the opposite direction.
function NatToUnary(n: nat): Unary {
  if n == 0 then Zero else Suc(NatToUnary(n-1))
}

lemma NatUnaryCorrespondence(n: nat, x: Unary)
  ensures UnaryToNat(NatToUnary(n)) == n
  ensures NatToUnary(UnaryToNat(x)) == x
{
}

predicate Less(x: Unary, y: Unary) {
  y != Zero && (x.Suc? ==> Less(x.pred, y.pred))
}

lemma LessCorrect(x: Unary, y: Unary)
  ensures Less(x, y) <==> UnaryToNat(x) < UnaryToNat(y)
{
}

lemma LessTransitive(x: Unary, y: Unary, z: Unary)
  requires Less(x, y) && Less(y, z)
  ensures Less(x, z)
{
}

function Add(x: Unary, y: Unary): Unary {
  match y
  case Zero => x
  case Suc(y') => Suc(Add(x, y'))
}

lemma AddCorrect(x: Unary, y: Unary)
  ensures UnaryToNat(Add(x, y)) == UnaryToNat(x) + UnaryToNat(y)
{
}

lemma SucAdd(x: Unary, y: Unary)
  ensures Suc(Add(x, y)) == Add(Suc(x), y)
{
}

lemma AddZero(x: Unary)
  ensures Add(Zero, x) == x
{
}

function Sub(x: Unary, y: Unary): Unary
  requires !Less(x, y)
{
  match y
  case Zero => x
  case Suc(y') => Sub(x.pred, y')
}

lemma SubCorrect(x: Unary, y: Unary)
  requires !Less(x, y)
  ensures UnaryToNat(Sub(x, y)) == UnaryToNat(x) - UnaryToNat(y)
{
}

function Mul(x: Unary, y: Unary): Unary {
  match x
  case Zero => Zero
  case Suc(x') => Add(Mul(x', y), y)
}

lemma MulCorrect(x: Unary, y: Unary)
  ensures UnaryToNat(Mul(x, y)) == UnaryToNat(x) * UnaryToNat(y)
{
  match x
  case Zero =>
  case Suc(x') =>
    calc {
      UnaryToNat(Mul(x, y));
      == UnaryToNat(Add(Mul(x', y), y));
      == { AddCorrect(Mul(x', y), y); } UnaryToNat(Mul(x', y)) + UnaryToNat(y);
    }
}

// Let's finish up with another classic container structure,
// which we will spend more time with next time.

datatype BST<V> = Leaf
  | Node(key: int, value: V, left: BST<V>, right: BST<V>)
// Here we have binary search trees (BSTs), specialized to integers
// as keys, because it takes some more work to handle arbitrary key types.
// (We'll see one way next time.)

// In fact, this is a good excuse to introduce another classic container,
// much more commonly found in functional languages than elsewhere.
datatype Option<T> = None | Some(value: T)

function Lookup<V>(t: BST<V>, k: int): Option<V>
{
  match t
  case Leaf => None
  case Node(k', v, l, r) =>
    if k == k' then
      Some(v)
    else if k < k' then
      Lookup(l, k)
    else
      Lookup(r, k)
}

function Insert<V>(t: BST<V>, k: int, v: V): BST<V>
{
  match t
  case Leaf => Node(k, v, Leaf, Leaf)
  case Node(k', v', l, r) =>
    if k == k' then
      Node(k, v, l, r)
    else if k < k' then
      Node(k', v', Insert(l, k, v), r)
    else
      Node(k', v', l, Insert(r, k, v))
}

// This property gets a nice automated proof.
lemma Lookup_Insert<V>(t: BST<V>, k: int, v: V, k': int)
  ensures Lookup(Insert(t, k, v), k') == if k' == k then Some(v) else Lookup(t, k')
{
}

// Let's also reason about the ranges of values that appear in BSTs.
// A range has optional lower and upper bounds, pushed into pairs like so.
type range = (Option<int>, Option<int>)

predicate InRange(n: int, r: range) {
  (r.0.None? || r.0.value < n)
    && (r.1.None? || n < r.1.value)
}

// A BST is well-formed w.r.t. a range if all its keys fall in the range.
predicate WellFormedBST<V>(t: BST<V>, rng: range) {
  match t
  case Leaf => true
  case Node(k, _, l, r) =>
    InRange(k, rng)
    && WellFormedBST(l, (rng.0, Some(k)))
    && WellFormedBST(r, (Some(k), rng.1))
}

// The next two proofs are interesting for forcing us to invoke IHes explicitly.
lemma LookupRange<V>(t: BST<V>, k: int, rng: range)
  requires WellFormedBST(t, rng)
  requires Lookup(t, k).Some?
  ensures InRange(k, rng)
{
  match t
  case Leaf =>
  case Node(k', v, l, r) =>
    if k == k' {
    } else if k < k' {
      LookupRange(l, k, (rng.0, Some(k')));
    } else {
      LookupRange(r, k, (Some(k'), rng.1));
    }
}

lemma InsertWellFormed<V>(t: BST<V>, rng: range, k : int, v: V)
  requires WellFormedBST(t, rng)
  requires InRange(k, rng)
  ensures WellFormedBST(Insert(t, k, v), rng)
{
  match t
  case Leaf =>
  case Node(k', v', l, r) =>
    if k == k' {
    } else if k < k' {
      InsertWellFormed(l, (rng.0, Some(k')), k, v);
    } else {
      InsertWellFormed(r, (Some(k'), rng.1), k, v);
    }
}