// Let's bring back our friend the list from last time.
datatype List<T> = Nil | Cons(head: T, tail: List<T>)

// Here are a few definitions copied over.
function Length<T>(xs: List<T>): nat {
  match xs
  case Nil => 0
  case Cons(_, tail) => 1 + Length(tail)
}

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))
}

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)
}

// Now something new: higher-order functions.

function Map<A, B>(f: A -> B, ls: List<A>): List<B> {
  match ls
  case Nil => Nil
  case Cons(x, ls') => Cons(f(x), Map(f, ls'))
}

lemma LengthMap<A, B>(f: A -> B, ls: List<A>)
  ensures Length(Map(f, ls)) == Length(ls)
{
}

lemma MapAppend<A, B>(f: A -> B, ls1: List<A>, ls2: List<A>)
  ensures Map(f, Append(ls1, ls2)) == Append(Map(f, ls1), Map(f, ls2))
{
}

// We can work with functions in an even more first-class way.

function Identity<A>(): A -> A {
  x => x
}

lemma MapIdentity<A>(ls: List<A>)
  ensures Map(Identity(), ls) == ls
{
}

function Compose<A, B, C>(f: A -> B, g: B -> C): A -> C {
  x => g(f(x))
}

lemma MapCompose<A, B, C>(f: A -> B, g: B -> C, ls: List<A>)
  ensures Map(Compose(f, g), ls) == Map(g, Map(f, ls))
{
}

// This function is a classic for stepping through all elements of a list.
function Foldl<A, B>(f : (A, B) -> B, acc: B, ls: List<A>): B {
  match ls
  case Nil => acc
  case Cons(x, ls') => Foldl(f, f(x, acc), ls')
}

// For instance, we can use Foldl to add up the elements of a list.
lemma Foldl_example()
  ensures Foldl((n, m) => n + m, 0, Cons(1, Cons(2, Cons(3, Nil)))) == 3+(2+(1+0))
{
}
// The order in which we wrote that sum at the end is suggestive.
// Foldl somewhat naturally applies list elements in reverse order.
// Let's formalize that connection to list reversal.

lemma FoldlReverseAux<A>(ls: List<A>, acc: List<A>)
  ensures Foldl((h, t) => Cons(h, t), acc, ls) == ReverseAux(ls, acc)
{
}

lemma FoldlReverse<A>(ls: List<A>)
  ensures Foldl((h, t) => Cons(h, t), Nil, ls) == Reverse(ls)
{
  FoldlReverseAux(ls, Nil);
}
// Yup, it worked!

// Here is a classic algebraic law relating Map and Fold.
lemma FoldlMap<A, B, C>(f: A -> B, g : (B, C) -> C, acc: C, ls: List<A>)
  ensures Foldl(g, acc, Map(f, ls)) == Foldl((x, a) => g(f(x), a), acc, ls)
{
  match ls
  case Nil =>
  case Cons(x, ls') =>
    calc {
      Foldl(g, acc, Map(f, ls));
      == Foldl(g, acc, Map(f, Cons(x, ls')));
      == Foldl(g, acc, Cons(f(x), Map(f, ls')));
      == Foldl(g, g(f(x), acc), Map(f, ls'));
      == { FoldlMap(f, g, g(f(x), acc), ls'); } Foldl((x, a) => g(f(x), a), g(f(x), acc), ls');
      == Foldl((x, a) => g(f(x), a), acc, ls);
    }

    // Or just the following line suffices:
    // FoldlMap(f, g, g(f(x), acc), ls');
}

// Now back to binary search trees.
// This time, using higher-order functions, we venture into making trees generic
// in their key types!

datatype Option<T> = None | Some(value: T)

datatype BST<K, V> = Leaf
  | Node(key: K, value: V, left: BST<K, V>, right: BST<K, V>)

// Note how Lookup (and many more operations to come) takes in a "less-than"
// comparison function as an argument.
function Lookup<K(==), V>(lt: (K, K) -> bool, t: BST<K, V>, k: K): Option<V>
{
  match t
  case Leaf => None
  case Node(k', v, l, r) =>
    if k == k' then
      Some(v)
    else if lt(k, k') then
      Lookup(lt, l, k)
    else
      Lookup(lt, r, k)
}

function Insert<K(==), V>(lt: (K, K) -> bool, t: BST<K, V>, k: K, v: V): BST<K, 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 lt(k, k') then
      Node(k', v', Insert(lt, l, k, v), r)
    else
      Node(k', v', l, Insert(lt, r, k, v))
}

// Nifty: we can prove this old law, without needing to know anything about the
// implementation of 'lt'.
lemma Lookup_Insert<K, V>(lt: (K, K) -> bool, t: BST<K, V>, k: K, v: V, k': K)
  ensures Lookup(lt, Insert(lt, t, k, v), k') == if k' == k then Some(v) else Lookup(lt, t, k')
{
}

// OK, let's try a more stringent example.  This function flips the tree on its
// side.  What should the effect be on the tree, interpreted as a proper binary
// search tree?
function Mirror<K, V>(t: BST<K, V>): BST<K, V> {
  match t
  case Leaf => Leaf
  case Node(k, v, l, r) => Node(k, v, Mirror(r), Mirror(l))
}

// This function transformer will be helpful.
// It flips the senses of the two arguments, which is just what we need to
// explain what Mirror accomplishes.
function Flip<A, B, C>(f: (A, B) -> C): (B, A) -> C {
  (b, a) => f(a, b)
}

// So Mirror preserves Lookup, when we flip the less-than test.
// Note that, working through the proof, we realize that this property depends
// on 'lt' acting like a real total order.
lemma LookupMirror<K, V>(lt: (K, K) -> bool, t: BST<K, V>, k: K)
  requires forall a, b : K :: lt(a, b) == (!lt(b, a) && a != b)
  ensures Lookup(lt, t, k) == Lookup(Flip(lt), Mirror(t), k)
{
  // This proof actually goes through automatically,
  // but we show a 'calc' that helped us realize what precondition to add.
  match t
  case Leaf =>
  case Node(k', v', l, r) =>
    if k == k' {
    } else if lt(k, k') {
      calc {
        Lookup(lt, t, k);
        == Lookup(lt, Node(k', v', l, r), k);
        == Lookup(lt, l, k);
        == { LookupMirror(lt, l, k); } Lookup(Flip(lt), Mirror(l), k);
        == { assert Flip(lt)(k', k); } Lookup(Flip(lt), Node(k', v', Mirror(r), Mirror(l)), k);
        == Lookup(Flip(lt), Mirror(Node(k', v', l, r)), k);
        == Lookup(Flip(lt), Mirror(t), k);
      }
    } else {
    }
}

// This property may seem strange to need to state explicitly: Lookup acts the
// same, when passed either of two functions that act the same.  However, that
// property *does not* hold for any old Dafny function!  Or, at least, Dafny
// doesn't automatically conclude it for us; we need to inspect each function.
// The intuitive explanation is that, sure, different sorting algorithms may
// always produce equal outputs, but do we *really* want to consider those
// algorithms equal?
lemma LookupExtensional<K, V>(lt: (K, K) -> bool, lt': (K, K) -> bool, t: BST<K, V>, k: K)
  requires forall a, b : K :: lt(a, b) == lt'(a, b)
  ensures Lookup(lt, t, k) == Lookup(lt', t, k)
{
}

// The precondition two lemmas back was a little delicate, so let's check it
// works with a concrete key type.
lemma LookupMirrorInt<V>(t: BST<int, V>, k: int)
  ensures Lookup((a, b) => a < b, t, k) == Lookup((b, a) => a < b, Mirror(t), k)
{
  var lt := (a: int, b: int) => a < b;
  var lt' := (b: int, a: int) => a < b;

  LookupMirror(lt, t, k);
  LookupExtensional(Flip(lt), lt', Mirror(t), k);
}

// Let's relate trees and lists.

// How many nodes in a tree?
function Size<K, V>(t: BST<K, V>): nat {
  match t
  case Leaf => 0
  case Node(_, _, l, r) => 1 + Size(l) + Size(r)
}

// Sorted list of keys from a tree
// (Well, it's sorted if the tree is well-formed.)
function Keys<K, V>(t: BST<K, V>): List<K> {
  match t
  case Leaf => Nil
  case Node(k, _, l, r) => Append(Keys(l), Cons(k, Keys(r)))
}

lemma SizeKeys<K, V>(t: BST<K, V>)
  ensures Length(Keys(t)) == Size(t)
{
}

// Checking if an element appears in a list
predicate Member<A(==)>(x: A, xs: List<A>) {
  match xs
  case Nil => false
  case Cons(y, xs') => x == y || Member(x, xs')
}

lemma MemberAppend<A>(x: A, xs1: List<A>, xs2: List<A>)
  ensures Member(x, Append(xs1, xs2)) == (Member(x, xs1) || Member(x, xs2))
{
}

// OK, now we're ready to characterize well-formedness of BSTs.
// This type captures a range of keys, optionally omitting a bound on either end.
type range<K> = (Option<K>, Option<K>)

predicate InRange<K>(lt: (K, K) -> bool, n: K, r: range<K>) {
  && (r.0.None? || lt(r.0.value, n))
  && (r.1.None? || lt(n, r.1.value))
}

// Now we can redefine our well-formedness predicate from last time, but now
// it's type-generic.
predicate WellFormedBST<K, V>(lt: (K, K) -> bool, t: BST<K, V>, rng: range) {
  match t
  case Leaf => true
  case Node(k, _, l, r) =>
    && InRange(lt, k, rng)
    && WellFormedBST(lt, l, (rng.0, Some(k)))
    && WellFormedBST(lt, r, (Some(k), rng.1))
}

// Let's check that well-formedness captures the intuitive property.
// Note that this property depends on a different algebraic property of 'lt':
// *transitivity*.
lemma RangeKeys<K, V>(lt: (K, K) -> bool, t: BST<K, V>, rng: range<K>, k: K)
  requires WellFormedBST(lt, t, rng)
  requires Member(k, Keys(t))
  requires forall a, b, c :: lt(a, b) && lt(b, c) ==> lt(a, c)
  ensures InRange(lt, k, rng)
{
  match t
  case Leaf =>
  case Node(k', v', l, r) =>
    // This proof goes through with automatic induction and one call to MemberAppend:
    MemberAppend(k, Keys(l), Cons(k', Keys(r)));
    // Here's a more detailed proof that does not rely on automatic induction:
    /*
    assert Member(k, Keys(l)) || k == k' || Member(k, Keys(r)) by {
      assert Member(k, Append(Keys(l), Cons(k', Keys(r))));
      MemberAppend(k, Keys(l), Cons(k', Keys(r)));
    }
    if
    case Member(k, Keys(l)) =>
      RangeKeys(lt, l, (rng.0, Some(k')), k);
    case k == k' =>
    case Member(k, Keys(r)) =>
      RangeKeys(lt, r, (Some(k'), rng.1), k);
    */
}

// The following we probably won't have time for in class, especially given how
// long some of the proofs are, but they should be educational to demonstrate
// more-intricate lemma sequences.

// Working up to further BST properties, let's define what it means for two
// lists to share no elements.
predicate Disjoint<A(==)>(ls1: List<A>, ls2: List<A>) {
  match ls1
  case Nil => true
  case Cons(x, ls1') => !Member(x, ls2) && Disjoint(ls1', ls2)
}

lemma DisjointCons<A>(ls1: List<A>, x: A, ls2: List<A>)
  ensures Disjoint(ls1, Cons(x, ls2)) == (!Member(x, ls1) && Disjoint(ls1, ls2))
{
}

// Checking if a list is duplicate-free.
predicate NoDup<A(==)>(xs: List<A>) {
  match xs
  case Nil => true
  case Cons(x, xs') => !Member(x, xs') && NoDup(xs')
}

// That predicate is helpful to characterize the interaction of NoDup and
// Append.
lemma NoDupAppend<A>(ls1: List<A>, ls2: List<A>)
  ensures NoDup(Append(ls1, ls2)) == (Disjoint(ls1, ls2) && NoDup(ls1) && NoDup(ls2))
{
  match ls1
  case Nil =>
  case Cons(v, ls1') =>
    calc {
      NoDup(Append(ls1, ls2));
      == NoDup(Append(Cons(v, ls1'), ls2));
      == NoDup(Cons(v, Append(ls1', ls2)));
      == !Member(v, Append(ls1', ls2)) && NoDup(Append(ls1', ls2));
      == { NoDupAppend(ls1', ls2); } !Member(v, Append(ls1', ls2)) && Disjoint(ls1', ls2) && NoDup(ls1') && NoDup(ls2);
      == { MemberAppend(v, ls1', ls2); } !Member(v, ls1') && !Member(v, ls2) && Disjoint(ls1', ls2) && NoDup(ls1') && NoDup(ls2);
    }
}

// Here's an even-more-direct but useful property.
lemma DisjointAppend<A>(ls1: List<A>, ls2: List<A>, ls: List<A>)
  ensures Disjoint(Append(ls1, ls2), ls) == (Disjoint(ls1, ls) && Disjoint(ls2, ls))
{
}

// Now let's prove an involved kind of "strengthened induction hypothesis"
// allowing us to conclude that two trees have disjoint keys.
lemma BSTDisjoint<K, V>(lt: (K, K) -> bool, t1: BST<K, V>, t2: BST<K, V>, lower1: Option<K>, upper1: K, lower2: K, upper2: Option<K>)
  requires WellFormedBST(lt, t1, (lower1, Some(upper1)))
  requires WellFormedBST(lt, t2, (Some(lower2), upper2))
  requires lt(upper1, lower2) || upper1 == lower2
  // That last one is critical: the two trees' ranges are ordered in a
  // particular way, overlapping in *at most one* element, and given how we
  // interpret range ends strictly, that implies *no overlap of elements*.
  
  requires forall a, b, c :: lt(a, b) && lt(b, c) ==> lt(a, c)
  requires forall a :: !lt(a, a) // New property: 'lt' is *irreflexive*.

  ensures Disjoint(Keys(t1), Keys(t2))
{
  match t1
  case Leaf =>
  case Node(k, v, l, r) =>
    // Note a little *proof by contradiction* here,
    // to show how we know a given fact doesn't hold.
    if Member(k, Keys(t2)) {
      assert InRange(lt, k, (lower1, Some(upper1)));
      assert lt(k, upper1);
      assert InRange(lt, k, (Some(lower2), upper2)) by { RangeKeys(lt, t2, (Some(lower2), upper2), k); }
      assert lt(lower2, k);
      assert lt(k, k); // by transitivity
      assert false;    // by antireflexivity
    }
    assert !Member(k, Keys(t2));
    
    calc {
      Disjoint(Append(Keys(l), Cons(k, Keys(r))), Keys(t2));
      == { DisjointAppend(Keys(l), Cons(k, Keys(r)), Keys(t2)); } !Member(k, Keys(t2)) && Disjoint(Keys(l), Keys(t2)) && Disjoint(Keys(r), Keys(t2));
      == { BSTDisjoint(lt, l, t2, lower1, k, lower2, upper2); } !Member(k, Keys(t2)) && true && Disjoint(Keys(r), Keys(t2));
      == { BSTDisjoint(lt, r, t2, Some(k), upper1, lower2, upper2); } !Member(k, Keys(t2)) && true && true;
      == true && true && true;
      == true;
    }
}

// Finally, we pull it together into one pleasingly simple property: well-formed
// BSTs have no duplicate keys.
lemma NoDupKeys<K, V>(lt: (K, K) -> bool, t: BST<K, V>, rng: range<K>)
  requires WellFormedBST(lt, t, rng)
  requires forall a :: !lt(a, a)
  requires forall a, b, c :: lt(a, b) && lt(b, c) ==> lt(a, c)
  ensures NoDup(Keys(t))
{
  match t
  case Leaf =>
  case Node(k, v, l, r) =>
    var L, R := Keys(l), Keys(r);
    var R' := Cons(k, R);

    calc {
      NoDup(Keys(t));
      NoDup(Append(L, R'));
      { NoDupAppend(L, R'); }
      Disjoint(L, R') && NoDup(L) && NoDup(R');
      { DisjointCons(L, k, R); }
      !Member(k, L) && Disjoint(L, R) && NoDup(L) && NoDup(R');
      !Member(k, L) && Disjoint(L, R) && NoDup(L) && !Member(k, R) && NoDup(R);
      { NoDupKeys(lt, l, (rng.0, Some(k))); }
      !Member(k, L) && Disjoint(L, R) && true && !Member(k, R) && NoDup(R);
      { NoDupKeys(lt, r, (Some(k), rng.1)); }
      !Member(k, L) && Disjoint(L, R) && true && !Member(k, R) && true;
    }

    assert !Member(k, L) by {
      if Member(k, Keys(l)) {
        RangeKeys(lt, l, (rng.0, Some(k)), k);
      }
    }

    assert !Member(k, R) by {
      if Member(k, Keys(r)) {
        RangeKeys(lt, r, (Some(k), rng.1), k);
      }
    }

    assert Disjoint(L, R) by {
      BSTDisjoint(lt, l, r, rng.0, k, k, rng.1);
    }
}