include "pset03_definitions.dfy"

// We refer to a few definitions from this other file, which you probably want
// to take a quick look at.

// One META-HINT about proofs in this pset: Dafny's automatic induction is very
// good at this kind of exercise!  Our solution contains just *four* lines of
// proof, across the whole set of exercises (before the optional one at the
// end).  However, we often have to do work choosing the right auxiliary
// definitions and lemmas.  To find the right ones to define/prove, it is
// helpful to start working through more manual proofs with 'calc', case
// analysis, and so on.  Then you can delete these proofs as you realize how.

module Pset3 {
  import opened Defs

  // PART 1: Binary numbers and arithmetic

  // Let's define binary numbers from first principles!
  // Specifically, a binary number is just a list of bits
  // (Booleans: 'true' for 1, 'false' for 0).
  // See the 'type' definition in the file we included above.
  // There is still some ambiguity in this choice of type.
  // Please represent binary numbers such that the least-significant bit appears
  // first.  This would be the bit that appears *last* in conventional notation,
  // so we are effectively storing binary numbers in *reverse order* to simplify
  // your job.  (You may see how it helps, as you proceed through the
  // assignment.)

  // Your first challenge: implement both of these conversions between our
  // binary-numbers type and the built-in natural-number type.
  function BinaryToNat(b: Binary): nat
  function NatToBinary(n: nat): Binary

  // Your second challenge: prove this lemma about round-tripping in one order.
  lemma NatToBinary_and_back(n: nat)
    ensures BinaryToNat(NatToBinary(n)) == n

  // Next, we would like to prove the round-tripping law in the other order.
  // However, it is not literally true!  (At least, it isn't with the
  // implementations we chose.)  Define a predicate to capture when a binary
  // number is *well-formed*.  HINT: we can use well-formedness to be sure there
  // are not multiple ways to write the same number.  Make sure this predicate
  // accepts just one *canonical* binary representation per number.
  predicate WellFormedBinary(b: Binary)

  // Then you should be able to prove the remaining round-tripping law.
  lemma BinaryToNat_and_back(b: Binary)
    requires WellFormedBinary(b)
    ensures NatToBinary(BinaryToNat(b)) == b

  // Also, show that your own translation respects this well-formedness notion.
  lemma NatToBinary_WellFormed(n: nat)
    ensures WellFormedBinary(NatToBinary(n))

  // The last subchallenge: define *addition* on binary numbers.
  // For full credit, DO NOT use your translations back and forth with 'nat',
  // and indeed don't use any translations of whole bitstrings to normal number
  // types.  Try to write code that operates as directly on the binary
  // representation as possible!
  // HINT: you will probably find it useful to define multiple helper functions,
  // some of which may get their own correctness lemmas.
  function Add(n: Binary, m: Binary): Binary
    requires Length(n) == Length(m)

  // Finally, prove your implementation correct.
  lemma Add_correct(n: Binary, m: Binary)
    requires Length(n) == Length(m)
    ensures BinaryToNat(Add(n, m)) == BinaryToNat(n) + BinaryToNat(m)

  // PART 2: Binary tries

  // *Tries* are kind of a B-list data structure, not nearly as famous as e.g.
  // binary search trees but still covered in many computer-science curricula.
  // Check out 'pset03_definitions.dfy' for our definition of *binary tries* and
  // what it means for a binary number to belong to one.  Here, we use tries to
  // represent *sets of binary numbers*.  And it's a rather efficient choice of
  // data structure, with every relevant operation taking time linear in the
  // length of the binary number.

  // How are tries laid out?
  // - A Leaf(false) is an empty set.
  // - A Leaf(true) is a one-element set of just the empty binary number.
  // - A Node(inc, left, right) stands for a set including:
  //   -- The set that left stands for, *with a 0 added to the front of every
  //      binary number*
  //   -- The set that right stands for, *with a 1 added to the front of every
  //      binary number*
  //   -- If inc is true, then we also add in the empty binary number.
  // This description is a bit (<-- pun?) imprecise, but see the definition of
  // Member for full detail.

  // Challenge 1: implement and prove insertion: building a new set containing
  // one additional binary number.
  function Insert(b: Binary, t: BTrie): BTrie
  lemma MemberInsert(b: Binary, t: BTrie, b': Binary)
    ensures Member(b', Insert(b, t)) == (b' == b || Member(b', t))
  lemma InsertInsert(b: Binary, t: BTrie)
    ensures Insert(b, Insert(b, t)) == Insert(b, t)

  // Challenge 2: implement and prove deletion
  function Delete(b: Binary, t: BTrie): BTrie
  lemma MemberDelete(b: Binary, t: BTrie, b': Binary)
    ensures Member(b', Delete(b, t)) == (b' != b && Member(b', t))
  lemma DeleteDelete(b: Binary, t: BTrie)
    ensures Delete(b, Delete(b, t)) == Delete(b, t)

  // Challenge 3: implement and prove merging of two tries, corresponding to
  // set union
  function Merge(t1: BTrie, t2: BTrie): BTrie
  lemma MemberMerge(b: Binary, t1: BTrie, t2: BTrie)
    ensures Member(b, Merge(t1, t2)) == (Member(b, t1) || Member(b, t2))

  // Challenge 4: We can actually generalize everything that we did!
  // Let's work with a generic Trie that can store any type T.
  // Find its definition (and some related functions) in Defs.

  // Implement GLookup, which should act similarly to Member:
  // If the binary string exists in the trie, return the value at the corresponding node,
  // otherwise return None.
  function GLookup<T>(b: Binary, gt: GTrie<T>): Option<T>

  // And implement GMerge, which should act similarly to Merge!
  // If there are two values, GMerge should combine them using op.
  function GMerge<T>(op: (T, T) -> T, gt1: GTrie<T>, gt2: GTrie<T>): GTrie<T>

  // Some example lemmas that should automatically verify with the right definitions 
  // of GMerge and GLookup.
  lemma DepthBounds<T>(gt: GTrie<T>, b: Binary)
    requires GLookup(b, gt).Some?
    ensures Length(b) <= Depth(gt)
  {
  }

  lemma GMergeComm<T>(op: (T, T) -> T, gt1: GTrie<T>, gt2: GTrie<T>)
    requires forall a, b :: op(a, b) == op(b, a)
    ensures GMerge(op, gt1, gt2) == GMerge(op, gt2, gt1)
  {
  }

  lemma GMergeAssoc<T>(op: (T, T) -> T, gt1: GTrie<T>, gt2: GTrie<T>, gt3: GTrie<T>)
    requires forall a, b, c :: op(op(a, b), c) == op(a, op(b, c))
    ensures GMerge(op, GMerge(op, gt1, gt2), gt3) == GMerge(op, gt1, GMerge(op, gt2, gt3))
  {
  }

  // And implement GCombine
  function GCombine<T>(op: (T,T) -> T, x: Option<T>, y: Option<T>): Option<T>

  // The primary spec of GCombine: Combining the results of lookups should be the same 
  // as looking up the merge. 
  // This should prove automatically, and may be helpful for future parts
  lemma GCombineIsLookupOverMerge<T>(op: (T, T) -> T, b: Binary, gt1: GTrie<T>, gt2: GTrie<T>)
    ensures GLookup(b, GMerge(op, gt1, gt2)) == GCombine(op, GLookup(b, gt1), GLookup(b, gt2))

  // Now let's show (partially) that the GTrie used with bool
  // is actually equivalent to the original BTrie.

  // What does it mean for the tries to be the same?
  // They should give the same result for a binary string.
  predicate GenericBoolTrieMatchesBTrieOnBinary(b: Binary, gt: GTrie<bool>, bt: BTrie)
  {
    Member(b,bt) <==> (GLookup(b,gt) == Some(true))
  }

  // Merging *should* just be applying an Or to the GTrie.
  // Challenge 5a: Show that if we have matching pairs of GTrie<bool> and BTrie,
  //               we can merge them to obtain another pair that matches.
  // This proof does not require many steps, if we reuse earlier lemmas from the pset.
  lemma GenericBoolTrieMatchesBTrieAfterMerge(b: Binary, gt1: GTrie<bool>, gt2: GTrie<bool>, bt1: BTrie, bt2: BTrie)
    requires GenericBoolTrieMatchesBTrieOnBinary(b, gt1, bt1)
    requires GenericBoolTrieMatchesBTrieOnBinary(b, gt2, bt2)
    ensures GenericBoolTrieMatchesBTrieOnBinary(b, GMerge(Or, gt1, gt2), Merge(bt1, bt2))

  // What about tries made of functions?
  // We can apply a function found in a trie, or use the identity if we don't find any function
  function FuncTrieApply<T>(b: Binary, ft: GTrie<(T->T)>, x: T): T {
    match GLookup(b,ft)
    case None => x
    case Some(f) => f(x)
  }

  // Challenge 5b: Show that applying functions after merging with function composition
  //               is correct.
  // This may feel similar to 5a. Try to reuse generic specifications!
  // That is, we hope this proof and the previous one are nearly one-liners, calling the
  // same generic lemma you come up with. We admit, we won't check for niceness of code
  // reuse in autograding, but it's a worthwhile exercise.
  lemma FuncCompositionTrie<T>(b: Binary, ft1: GTrie<(T -> T)>, ft2: GTrie<(T -> T)>, x: T)
    ensures FuncTrieApply(b, GMerge(FuncComposition, ft1, ft2), x) == FuncTrieApply(b, ft2, FuncTrieApply(b, ft1, x))

  // This next exercise is *optional* because we found it requires quite some
  // wrangling of Dafny to avoid a *timeout* of automatic verification.  That
  // is, if we're not very careful in laying out the proof, Dafny runs out of
  // patience and gives up after 10 seconds.

  // Optional Challenge 6: Prove that the depth of the GTrie after merging is the
  //                       Max of its two component tries.
  // Note:  Verification of this lemma is likely to time out.
  //        For brave souls attemping the proof, we highly suggest using fewer
  //        virtual cores for verification and setting a timeout in VSCode to
  //        avoid creating too many Z3 instances.
  //        The {:isolate_assertions} attribute tells Dafny to verify each assertion provided 
  //        so you can figure out where Dafny gets stuck.
  /*
  lemma {:isolate_assertions} GMergeDepth<T>(op: (T, T) -> T, gt1: GTrie<T>, gt2: GTrie<T>)
    ensures Depth(GMerge(op, gt1, gt2)) == Max(Depth(gt1), Depth(gt2))
  */
  // Only uncomment the lemma if you actually solve it, please.
  // Otherwise, our autograder might complain about a lemma with no proof.
}