// We've been using Dafny *modules* for bookkeeping on the problem sets, and
// let's take a closer look at them now.  Fundamentally, modules group related
// definitions and lemmas, providing both *namespace* control (which names can be
// referenced from which other places) as well as a further kind of
// *abstraction* that we get to shortly.

// Here is an example of packaging list definitions inside a module.
module ListLibrary {
  datatype List<T> = Nil | Cons(head: T, tail: List<T>)

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

  function Append<T>(xs: List<T>, ys: List<T>): List<T> {
    match xs
    case Nil => ys
    case Cons(h, t) => Cons(h, Append(t, ys))
  }

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

  function ReverseAux<T>(xs: List<T>, acc: List<T>): List<T> {
    match xs
    case Nil => acc
    case Cons(h, t) => ReverseAux(t, Cons(h, acc))
  }

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

  lemma SnocReverseAux<T>(xs: List<T>, acc: List<T>, x: T)
    ensures Snoc(ReverseAux(xs, acc), x) == ReverseAux(xs, Snoc(acc, x))
  {
  }

  lemma ReverseCons<T>(x: T, xs: List<T>)
    ensures Reverse(Cons(x, xs)) == Snoc(Reverse(xs), x)
  {
    calc {
      Reverse(Cons(x, xs));
      == ReverseAux(xs, Cons(x, Nil));
      == { SnocReverseAux(xs, Nil, x); } Snoc(ReverseAux(xs, Nil), x);
      == Snoc(Reverse(xs), x);
    }
  }

  lemma SnocAppend<T>(xs: List<T>, ys: List<T>, x: T)
    ensures Snoc(Append(xs, ys), x) == Append(xs, Snoc(ys, x))
  {
  }

  lemma ReverseAuxNonempty<T>(xs: List<T>, acc: List<T>)
    requires xs.Cons?
    ensures ReverseAux(xs, acc).Cons?
  {
  }

  lemma ReverseNonempty<T>(xs: List<T>)
    requires xs.Cons?
    ensures Reverse(xs).Cons?
  {
    ReverseAuxNonempty(xs, Nil);
  }

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

// Now a function outside the module can access the contents with the "dot"
// notation.
function FirstN(n: nat): ListLibrary.List<nat> {
  if n == 0 then
    ListLibrary.Nil
  else
    ListLibrary.Append(FirstN(n-1), ListLibrary.Cons(n, ListLibrary.Nil))
}

// Check out how one module may reference contents of another.

module ModuleA {
  function Plus(x: int, y: int): int {
    x + y
  }
}

module ModuleB {
  import ModuleA
  // Note that this prior line is needed to avoid a name-resolution error.
  
  function Double(x: int): int {
    ModuleA.Plus(x, x)
  }
}

// Alternatively, one can assign a local shorthand for an imported module.
module ModuleB' {
  import A = ModuleA
  
  function Double(x: int): int {
    A.Plus(x, x)
  }
}

// Finally, one can import a module *without* prefixing its names:
// discouraged in general, we have used it on the problem sets for simplicity.
module ModuleB'' {
  import opened ModuleA
  
  function Double(x: int): int {
    Plus(x, x)
  }
}

// Now on to *abstraction*, the even more important behavior of modules (and
// arguably one of the top ideas we hope this class gives you practice thinking
// about).

// Let's package up two not-so-related definitions.
module ModuleC {
  // What is this next command?  We explain shortly.
  export
    provides Color
    reveals Double

  datatype Color = Blue | Yellow | Green | Red

  function Double(x: int): nat
    requires 0 <= x
    ensures Double(x) % 2 == 0
  {
    if x == 0 then 0 else 2 + Double(x - 1)
  }
}

// Another module tries to use the contents.
module ModuleD {
  import ModuleC

  function Test(): nat
    ensures Test() == 6
  {
    //var c : ModuleC.Color := ModuleC.Yellow;
    // The previous commented line fails, saying ModuleC.Yellow is not
    // available.  Why?  Because ModuleC merely *provides*, not *reveals*,
    // Color, which means we know the type exists, but we don't know its
    // definition (so, in particular, we don't know which constructors exist).
    ModuleC.Double(3)
  }
}

module ModuleE {
  datatype Parity = Even | Odd

  function F(x: int): Parity {
    if x % 2 == 0 then Even else Odd
  }

  // export provides F
  // This first attempted export set is rejected, because we can't expose F
  // without revealing that the Parity type exists.

  // export provides Parity reveals F
  // We also can't reveal F without also revealing Even and Odd that appear
  // inside.

  export provides Parity, F
  // This one works fine.  The type signature of F makes perfect sense without
  // knowing the definition of Parity, even as the *body* of F depends on those
  // details.

  // We could also go for:
  // export reveals Parity, F
  // Revealing all definitions is actually the default, if we don't include an
  // export command.
}

// OK, now we're ready for the big idea of this class, which you probably got
// a sense for in a prerequisite course, but we're about to make it very
// concrete.  Parts of a program should be *abstracted* behind good interfaces
// that achieve *information hiding*, so users of such a program part don't need
// to know how it is implemented, and indeed other code can keep working as the
// part we're focusing on evolves.

// Here is a good interface for an immutable queue, where we enqueue new
// elements at the back and dequeue elements from the front, so dequeuing always
// returns the "oldest" element not yet returned.
module ImmutableQueue1 {
  type Queue<A>
  function Empty<A>(): Queue<A>
  function Enqueue<A>(q: Queue<A>, x: A): Queue<A>
  function Dequeue<A>(q: Queue<A>): (A, Queue<A>)
    requires q != Empty()
}
// Note how there is no implementation for any identifier in the module, but
// this interface is enough to code a client module against.  Well, it's enough
// if we don't care about doing proofs!  Let's see how we can enable client code
// to be proved against an arbitrary queue implementation.

module ImmutableQueue2 {
  import LL = ListLibrary

  type Queue<A>
  function Empty<A>(): Queue<A>
  function Enqueue<A>(q: Queue<A>, x: A): Queue<A>
  function Dequeue<A>(q: Queue<A>): (A, Queue<A>)
    requires q != Empty()

  // Most important new ingredient: an *abstraction function* that explains how
  // to translate a "secret" internal type to a standard representation.  In
  // this case, we explain queues as lists.
  ghost function Elements<A>(q: Queue<A>): LL.List<A>

  // Now these three lemmas uniquely determine the behavior of the three
  // remaining functions
  lemma EmptyCorrect<A>()
    ensures Elements(Empty<A>()) == LL.Nil
  lemma EnqueueCorrect<A>(q: Queue<A>, x: A)
    ensures Elements(Enqueue(q, x)) == LL.Snoc(Elements(q), x)
  lemma DequeueCorrect<A>(q: Queue<A>)
    requires q != Empty()
    ensures var (a, q') := Dequeue(q);
      LL.Cons(a, Elements(q')) == Elements(q)

  // Finally, let's provide (but not reveal) all of the above, so client code
  // isn't allowed to depend on the particular data structure we choose.
  export
    provides Queue, Empty, Enqueue, Dequeue
    provides LL, Elements
    provides EmptyCorrect, EnqueueCorrect, DequeueCorrect
}

// Let's check that simple client code can now use and reason about queues.
module QueueClient {
  import IQ = ImmutableQueue2

  function Client(): int
    ensures Client() == 20
  {
    var q := IQ.Empty();
    var q' := IQ.Enqueue(q, 20);
    // Oops, at first the precondition for Dequeue can't be proved!
    // The "assume false' is a placeholder to keep Dafny from complaining, as we
    // know this version of the code fails.  Let's try again.
    assume false;
    var (a, q'') := IQ.Dequeue(q');
    a
  }

  // The fix?  Keep invoking lemmas from ImmutableQueue.
  function Client2(): int
    ensures Client2() == 20
  {
    IQ.EmptyCorrect<int>();
    var q := IQ.Empty();
    IQ.EnqueueCorrect(q, 20);
    var q' := IQ.Enqueue(q, 20);
    IQ.DequeueCorrect(q');
    var (a, q'') := IQ.Dequeue(q');
    a
  }

  // What if we get more ambitious?
  function Client3(): IQ.Queue<int>
    ensures Client3() == IQ.Empty<int>()
  {
    IQ.EmptyCorrect<int>();
    var q := IQ.Empty();
    IQ.EnqueueCorrect(q, 20);
    var q' := IQ.Enqueue(q, 20);
    IQ.DequeueCorrect(q');
    var (a, q'') := IQ.Dequeue(q');
    // Huh: can't prove the postcondition, even though the following assertion
    // succeeds.
    assert IQ.Elements(q'') == IQ.LL.Nil;
    assume false;
    q''
  }
}

// The problem is that ImmutableQueue doesn't reveal that any queue with no
// elements equals Empty().  There are a variety of possible fixes, but let's
// take the most targeted one.
module ImmutableQueue3 {
  import LL = ListLibrary

  type Queue<A>
  function Empty<A>(): Queue<A>
  function Enqueue<A>(q: Queue<A>, x: A): Queue<A>
  function Dequeue<A>(q: Queue<A>): (A, Queue<A>)
    requires q != Empty()

  ghost function Elements<A>(q: Queue<A>): LL.List<A>

  lemma EmptyCorrect<A>()
    ensures Elements(Empty<A>()) == LL.Nil
  lemma EnqueueCorrect<A>(q: Queue<A>, x: A)
    ensures Elements(Enqueue(q, x)) == LL.Snoc(Elements(q), x)
  lemma DequeueCorrect<A>(q: Queue<A>)
    requires q != Empty()
    ensures var (a, q') := Dequeue(q);
      LL.Cons(a, Elements(q')) == Elements(q)

  // New lemma
  lemma EmptyUnique<A>(q: Queue<A>)
    ensures Elements(q) == LL.Nil ==> q == Empty()

  export
    provides Queue, Empty, Enqueue, Dequeue
    provides LL, Elements
    provides EmptyCorrect, EnqueueCorrect, DequeueCorrect, EmptyUnique
}

module QueueClient2 {
  import IQ = ImmutableQueue3

  function Client(): IQ.Queue<int>
    ensures Client() == IQ.Empty<int>()
  {
    IQ.EmptyCorrect<int>();
    var q := IQ.Empty();
    IQ.EnqueueCorrect(q, 20);
    var q' := IQ.Enqueue(q, 20);
    IQ.DequeueCorrect(q');
    var (a, q'') := IQ.Dequeue(q');
    IQ.EmptyUnique(q'');
    q''
  }
}

// To check our work, here is the most naive, list-based Queue implementation.
module NaiveImmutableQueue {
  import LL = ListLibrary

  type Queue<A> = LL.List<A>
  function Empty<A>(): Queue<A> {
    LL.Nil
  }
  function Enqueue<A>(q: Queue<A>, x: A): Queue<A> {
    LL.Snoc(q, x)
  }
  function Dequeue<A>(q: Queue<A>): (A, Queue<A>)
    requires q != Empty()
  {
    (q.head, q.tail)
  }

  ghost function Elements<A>(q: Queue<A>): LL.List<A> {
    q
  }

  lemma EmptyCorrect<A>()
    ensures Elements(Empty<A>()) == LL.Nil
  {
  }
  lemma EnqueueCorrect<A>(q: Queue<A>, x: A)
    ensures Elements(Enqueue(q, x)) == LL.Snoc(Elements(q), x)
  {
  }
  lemma DequeueCorrect<A>(q: Queue<A>)
    requires q != Empty()
    ensures var (a, q') := Dequeue(q);
      LL.Cons(a, Elements(q')) == Elements(q)
  {
  }

  lemma EmptyUnique<A>(q: Queue<A>)
    ensures Elements(q) == LL.Nil ==> q == Empty()
  {
  }

  export
    provides Queue, Empty, Enqueue, Dequeue
    provides LL, Elements
    provides EmptyCorrect, EnqueueCorrect, DequeueCorrect, EmptyUnique
}
// Trivial because of how directly we line up with Elements!
// However, this version has a slow implementation of Enqueue.
// It takes time proportional to the queue length.

// Here's a classic, more-clever implementation.
// A queue is a pair of lists, which we treat as *stacks*, where we only
// directly add or remove elements to/from their fronts.
module CleverImmutableQueue {
  import LL = ListLibrary

  datatype Queue<A> = FRQ(front: LL.List<A>, rear: LL.List<A>)
  function Empty<A>(): Queue<A> {
    FRQ(LL.Nil, LL.Nil)
  }
  function Enqueue<A>(q: Queue<A>, x: A): Queue<A> {
    FRQ(q.front, LL.Cons(x, q.rear))
  }

  // Clearly the last two operations now take constant time.
  // How about Dequeue?
  
  function Dequeue<A>(q: Queue<A>): (A, Queue<A>)
    requires q != Empty()
  {
    match q.front
    case Cons(x, front') =>
      (x, FRQ(front', q.rear))
    case Nil =>
      var front := LL.Reverse(q.rear);
      LL.ReverseCons(q.rear.head, q.rear.tail);
      (front.head, FRQ(front.tail, LL.Nil))
  }

  // The only step whose time varies substantially with queue size is the
  // Reverse.  That'll take time proportional to the queue length.  The cool
  // thing about this data structure is that it nonetheless achieves
  // *amortized constant time* for all operations.  See an algorithms class like
  // 6.122 for the details, but the intuition is that we encounter the case with
  // Reverse sufficiently rarely that its cost averages out nicely across calls
  // to Dequeue.  Luckily, proof of this data structure can proceed
  // independently of running-time details.

  // This function explains the key idea for how we get away with representing
  // queues via two stacks.
  ghost function Elements<A>(q: Queue<A>): LL.List<A> {
    LL.Append(q.front, LL.Reverse(q.rear))
  }

  // The rest of the proof work is just bureaucracy with list lemmas.  Enjoy!
  
  lemma EmptyCorrect<A>()
    ensures Elements(Empty<A>()) == LL.Nil
  {
  }

  lemma EnqueueCorrect<A>(q: Queue<A>, x: A)
    ensures Elements(Enqueue(q, x)) == LL.Snoc(Elements(q), x)
  {
    calc {
      Elements(Enqueue(q, x));
      == LL.Append(q.front, LL.Reverse(LL.Cons(x, q.rear)));
      == { LL.ReverseCons(x, q.rear); } LL.Append(q.front, LL.Snoc(LL.Reverse(q.rear), x));
      == { LL.SnocAppend(q.front, LL.Reverse(q.rear), x); } LL.Snoc(LL.Append(q.front, LL.Reverse(q.rear)), x);
      == LL.Snoc(Elements(q), x);
    }
  }

  lemma DequeueCorrect<A>(q: Queue<A>)
    requires q != Empty()
    ensures var (a, q') := Dequeue(q);
      LL.Cons(a, Elements(q')) == Elements(q)
  {
    match q.front
    case Cons(x, front') =>
    case Nil =>
      assert q.rear.Cons?;
      var front := LL.Reverse(q.rear);
      var (a, q') := Dequeue(q);
      assert front.Cons? by { LL.ReverseNonempty(q.rear); }
      assert a == front.head;
      assert q' == FRQ(front.tail, LL.Nil);
      calc {
        LL.Cons(a, Elements(q'));
        == LL.Cons(front.head, Elements(FRQ(front.tail, LL.Nil)));
        == LL.Cons(front.head, LL.Append(front.tail, LL.Reverse(LL.Nil)));
        == { LL.AppendNil(front.tail); } LL.Cons(front.head, front.tail);
        == Elements(q);
      }
  }

  lemma EmptyUnique<A>(q: Queue<A>)
    ensures Elements(q) == LL.Nil ==> q == Empty()
  {
    match q.front
    case Nil => {
      match q.rear
      case Nil =>
      case Cons(h, t) =>
        assert Elements(q) == LL.Reverse(LL.Cons(h, t));
        LL.ReverseNonempty(LL.Cons(h, t));
    }
    case Cons(_, _) =>
  }

  export
    provides Queue, Empty, Enqueue, Dequeue
    provides LL, Elements
    provides EmptyCorrect, EnqueueCorrect, DequeueCorrect, EmptyUnique
}

// In case we have time for another example, here's another fun one that we've
// sometimes used in 6.101, so it's nice to show that now we can put it on a formal
// footing.  It's an implementation of finite sets optimized for the common case
// of intervals of integers.
module IntSet {
  datatype Set =
    | Interval(lower: int, upper: int)
    // Really an interval: the case we optimize for
    | EmptySet
    // Self-explanatory!
    | Fallback(s: set<int>)
    // For when we aren't sure a set is really an interval.
    // Performance may suffer, but we should maintain correctness.

  function Empty(): Set {
    EmptySet
  }

  predicate Member(x: int, s: Set) {
    match s
    case EmptySet => false
    case Fallback(s) => x in s
    case Interval(lower, upper) => lower <= x <= upper
  }

  function IntervalSet(lower: int, upper: int): set<int> {
    set i | lower <= i <= upper
  }

  function Add(x: int, s: Set): Set {
    match s
    case EmptySet => Interval(x, x)
    case Fallback(s) => Fallback(s + {x})
    case Interval(lower, upper) =>
      if lower <= x <= upper then
        s
      else if x == lower-1 && lower <= upper then
        Interval(x, upper)
      else if x == upper+1 && lower <= upper then
        Interval(lower, x)
      else
        Fallback({x} + IntervalSet(lower, upper))
  }
  // Note a tricky part above: we need to check whether intervals are nonempty,
  // or the cases for extending intervals are unsound.

  ghost function Elements(s: Set): set<int> {
    match s
    case EmptySet => {}
    case Fallback(s) => s
    case Interval(lower, upper) => IntervalSet(lower, upper)
  }

  lemma EmptyCorrect()
    ensures Elements(Empty()) == {}
  {
  }

  lemma MemberCorrect(x: int, s: Set)
    ensures Member(x, s) == (x in Elements(s))
  {
  }

  lemma AddCorrect(x: int, s: Set)
    ensures Elements(Add(x, s)) == {x} + Elements(s)
  {
  }

  // Nice: all three of these correctness proofs go through automatically.

  export
    provides Set, Empty, Member, Add
    provides Elements, EmptyCorrect, MemberCorrect, AddCorrect
}

// Now we see how easy it is to implement functionality on top, like this check
// for whether a list contains duplicates.

// First, our Platonic definition of duplicate-freedom.
ghost predicate Unique(ls: seq<int>) {
  forall i, j :: 0 <= i < j < |ls| ==> ls[i] != ls[j]
}

// Now a version that should work efficiently for sequences that happen to
// contain intervals.
predicate NoDup(ls: seq<int>)
{
  NoDup'(ls, IntSet.Empty())
}

predicate NoDup'(ls: seq<int>, soFar: IntSet.Set)
{
  |ls| == 0
  || (!IntSet.Member(ls[0], soFar)
    && NoDup'(ls[1..], IntSet.Add(ls[0], soFar)))
}

lemma NoDupCorrect(ls: seq<int>)
  ensures NoDup(ls) == Unique(ls)
{
  IntSet.EmptyCorrect();
  NoDup'Correct(ls, IntSet.Empty());
}

lemma NoDup'Correct(ls: seq<int>, soFar: IntSet.Set)
  ensures NoDup'(ls, soFar) == (Unique(ls) && forall i :: i in ls ==> i !in IntSet.Elements(soFar))
{
  if |ls| == 0 {
  } else {
    IntSet.MemberCorrect(ls[0], soFar);
    IntSet.AddCorrect(ls[0], soFar);
  }
}