// Building on material from the previous class, let's now formalize the
// interface of a *priority queue*, where dequeue doesn't return the oldest item
// not already dequeued.  Rather, it returns the *lowest* integer not already
// dequeued.  Here is a specification in the style from last time.
module PriorityQueue0 {
  type PriorityQueue

  function Empty(): PriorityQueue
  predicate IsEmpty(pq: PriorityQueue)
  function Insert(pq: PriorityQueue, y: int): PriorityQueue
  function RemoveMin(pq: PriorityQueue): (int, PriorityQueue)
    requires !IsEmpty(pq)

  ghost function Elements(pq: PriorityQueue): multiset<int>

  lemma EmptyCorrect()
    ensures Elements(Empty()) == multiset{}
  lemma IsEmptyCorrect(pq: PriorityQueue)
    ensures IsEmpty(pq) <==> Elements(pq) == multiset{}
  lemma InsertCorrect(pq: PriorityQueue, y: int)
    ensures Elements(Insert(pq, y)) == Elements(pq) + multiset{y}
  lemma RemoveMinCorrect(pq: PriorityQueue)
    requires !IsEmpty(pq)
    ensures var (y, pq') := RemoveMin(pq);
      IsMin(y, Elements(pq)) &&
      Elements(pq') + multiset{y} == Elements(pq)

  ghost predicate IsMin(y: int, s: multiset<int>) {
    y in s && forall x :: x in s ==> y <= x
  }
}

// Let's use a rather clever implementation: Braun trees.
// It will turn out that there is a nontrivial *data-structure invariant* that
// underlies correctness, and the above spec style doesn't let us enforce that
// invariant.
module PriorityQueue {
  // Braun trees look innocuous at this level: just your garden-variety binary
  // tree.
  datatype PriorityQueue =
    | Leaf
    | Node(x: int, left: PriorityQueue, right: PriorityQueue)

  // This *validity* predicate can also be called the *representation invariant*
  // of the data structure.
  ghost predicate Valid(pq: PriorityQueue) {
    IsBinaryHeap(pq) && IsBalanced(pq)
  }

  // Property 1: the root of the tree is the lowest number (though there may be
  //             "ties" elsewhere in the tree).
  ghost predicate IsBinaryHeap(pq: PriorityQueue) {
    match pq
    case Leaf => true
    case Node(x, left, right) =>
      IsBinaryHeap(left) && IsBinaryHeap(right) &&
      (left == Leaf || x <= left.x) &&
      (right == Leaf || x <= right.x)
  }

  // Property 2: every left child is the same size as its right sibling or is
  //             one element larger.
  ghost predicate IsBalanced(pq: PriorityQueue) {
    match pq
    case Leaf => true
    case Node(_, left, right) =>
      IsBalanced(left) && IsBalanced(right) &&
      var L, R := |Elements(left)|, |Elements(right)|;
      L == R || L == R + 1
  }

  function Empty(): PriorityQueue {
    Leaf
  }

  predicate IsEmpty(pq: PriorityQueue) {
    pq == Leaf
  }

  // Cool maneuver for maintaining balance when inserting: always swap the left
  // and right children.
  function Insert(pq: PriorityQueue, y: int): PriorityQueue {
    match pq
    case Leaf => Node(y, Leaf, Leaf)
    case Node(x, left, right) =>
      if y < x then
        Node(y, Insert(right, x), left)
      else
        Node(x, Insert(right, y), left)
  }

  // To remove the minimum element, computing that number is the easy part,
  // given the tree's ordering property.  A helper updates the tree to remove
  // that minimum.
  function RemoveMin(pq: PriorityQueue): (int, PriorityQueue)
    requires !IsEmpty(pq)
  {
    (pq.x, DeleteMin(pq))
  }

  // DeleteMin maintains the theme of swapping left and right children.
  function DeleteMin(pq: PriorityQueue): PriorityQueue
    requires !IsEmpty(pq)
  {
    if pq.left == Leaf || pq.right == Leaf then
      pq.left
    else if pq.left.x <= pq.right.x then
      Node(pq.left.x, pq.right, DeleteMin(pq.left))
    else
      Node(pq.right.x, ReplaceRoot(pq.right, pq.left.x), DeleteMin(pq.left))
  }

  // Another intricate helper
  function ReplaceRoot(pq: PriorityQueue, y: int): PriorityQueue
    requires !IsEmpty(pq)
  {
    if pq.left == Leaf || (y <= pq.left.x && (pq.right == Leaf || y <= pq.right.x)) then
      Node(y, pq.left, pq.right)
    else if pq.right == Leaf then
      Node(pq.left.x, Node(y, Leaf, Leaf), Leaf)
    else if pq.left.x < pq.right.x then
      Node(pq.left.x, ReplaceRoot(pq.left, y), pq.right)
    else
      Node(pq.right.x, pq.left, ReplaceRoot(pq.right, y))
  }

  // At least computing element sets is straightforward.
  ghost function Elements(pq: PriorityQueue): multiset<int> {
    match pq
    case Leaf => multiset{}
    case Node(x, left, right) =>
      multiset{x} + Elements(left) + Elements(right)
  }

  // OK, now we can include Valid in preconditions and postconditions as
  // appropriate.
  lemma EmptyCorrect()
    ensures var pq := Empty();
      Valid(pq) &&
      Elements(pq) == multiset{}
  {
  }

  lemma IsEmptyCorrect(pq: PriorityQueue)
    requires Valid(pq)
    ensures IsEmpty(pq) <==> Elements(pq) == multiset{}
  {
  }

  lemma InsertCorrect(pq: PriorityQueue, y: int)
    requires Valid(pq)
    ensures var pq' := Insert(pq, y);
      Valid(pq') &&
      Elements(pq') == Elements(pq) + multiset{y}
  {
  }

  // The first three go through automatically.  Nice!  There is more work ahead
  // of us for RemoveMin.

  lemma RemoveMinCorrect(pq: PriorityQueue)
    requires Valid(pq) && !IsEmpty(pq)
    ensures var (y, pq') := RemoveMin(pq);
      Valid(pq') &&
      IsMin(y, Elements(pq)) &&
      Elements(pq') + multiset{y} == Elements(pq)
  {
    DeleteMinCorrect(pq);
  }

  // As we move to the DeleteMin helper, we add a postcondition about the size
  // of Elements(..): this condition is needed by IsBalanced, and while it
  // follows from the other postconditions, Dafny will not be able to prove
  // that fact automatically.

  // In addition, a single lemma for DeleteMinCorrect generates quite a bit of
  // work for Dafny, and it is likely to time out. In fact, the staff ran into
  // this problem with the proof of this example in *Program Proofs*!
  // An effective strategy was to break the lemma into pieces; we must then
  // also provide a `decreases` clause that allows the main lemma to call the
  // two sub-lemmas with the whole tree as input, but calls back to the main
  // lemma require a structurally included subtree.

  lemma DeleteMinCorrect(pq: PriorityQueue)
    requires Valid(pq) && !IsEmpty(pq)
    ensures var pq' := DeleteMin(pq);
      Valid(pq') &&
      Elements(pq') + multiset{pq.x} == Elements(pq) &&
      |Elements(pq')| == |Elements(pq)| - 1
    decreases pq, 1
  {
    if pq.left == Leaf || pq.right == Leaf {
    } else if pq.left.x <= pq.right.x {
      DeleteMinCorrectForLeftMin(pq);
    } else {
      DeleteMinCorrectForRightMin(pq);
    }
  }

  // We have two cases with sub-lemmas. The first is largely automatic, but
  // breaking it into its own lemma avoids the proof timing out.

  lemma DeleteMinCorrectForLeftMin(pq: PriorityQueue)
    requires Valid(pq) && !IsEmpty(pq)
    // And at this point in DeleteMin, we know:
    requires pq.left.Node? && pq.right.Node? && pq.left.x <= pq.right.x
    ensures var pq' := DeleteMin(pq);
      Valid(pq') &&
      Elements(pq') + multiset{pq.x} == Elements(pq) &&
      |Elements(pq')| == |Elements(pq)| - 1
    decreases pq, 0
  {
    DeleteMinCorrect(pq.left);
  }

  // In the second sub-lemma, we help Dafny reason about multisets...

  lemma DeleteMinCorrectForRightMin(pq: PriorityQueue)
    requires Valid(pq) && !IsEmpty(pq)
    // And at this point in DeleteMin, we know:
    requires pq.left.Node? && pq.right.Node? && pq.right.x < pq.left.x
    ensures var pq' := DeleteMin(pq);
      Valid(pq') &&
      Elements(pq') + multiset{pq.x} == Elements(pq) &&
      |Elements(pq')| == |Elements(pq)| - 1
    decreases pq, 0
  {
    var left, right := ReplaceRoot(pq.right, pq.left.x), DeleteMin(pq.left);
    var pq' := Node(pq.right.x, left, right);

    // We invoke a correctness lemma for ReplaceRoot and the induction
    // hypothesis for DeleteMin. Then we need to step Dafny through applying
    // those lemmas to show the desired equality of multisets; it is able to
    // prove the other postconditions without additional explanation.

    ReplaceRootCorrect(pq.right, pq.left.x); // so we know:
    assert Elements(pq.right) + multiset{pq.left.x} == Elements(left) + multiset{pq.right.x};
    DeleteMinCorrect(pq.left); // so we know:
    assert Elements(right) + multiset{pq.left.x} == Elements(pq.left);
    calc {
      Elements(pq') + multiset{pq.x};
      == // def'n Elements
      multiset{pq.right.x} + Elements(left) + Elements(right) + multiset{pq.x};
      == // ReplaceRootCorrect
      Elements(pq.right) + multiset{pq.left.x} + Elements(right) + multiset{pq.x};
      == // DeleteMinCorrect
      Elements(pq.right) + Elements(pq.left) + multiset{pq.x};
      == // def'n Elements
      Elements(pq);
    }
  }

  // Finally, the ReplaceRoot helper. Again we break the lemma into several
  // pieces to avoid timing out. And as we saw in the DeleteMin proofs above,
  // Dafny is able to check that the multisets of elements have all the desired
  // properties, but we must lead it through which properties to check and, in
  // some cases, which intermediate steps to take. These proofs proceed in a
  // similar vein.

  lemma ReplaceRootCorrect(pq: PriorityQueue, y: int)
    requires Valid(pq) && !IsEmpty(pq)
    ensures var pq' := ReplaceRoot(pq, y);
      Valid(pq') &&
      Elements(pq) + multiset{y} == Elements(pq') + multiset{pq.x} &&
      |Elements(pq)| == |Elements(pq')|
    decreases pq, 1
  {
    if pq.left == Leaf || (y <= pq.left.x && (pq.right == Leaf || y <= pq.right.x)) {
    } else if pq.right == Leaf {
    } else if pq.left.x < pq.right.x {
      ReplaceRootCorrectForLeftMin(pq, y);
    } else {
      ReplaceRootCorrectForRightMin(pq, y);
    }
  }

  // This time we have two symmetric cases with sub-lemmas...

  lemma ReplaceRootCorrectForLeftMin(pq: PriorityQueue, y: int)
    requires Valid(pq) && !IsEmpty(pq)
    // And at this point in ReplaceRoot, we know:
    requires pq.left.Node? && pq.right.Node? && pq.left.x < pq.right.x && pq.left.x < y
    ensures var pq' := ReplaceRoot(pq, y);
      Valid(pq') &&
      Elements(pq) + multiset{y} == Elements(pq') + multiset{pq.x} &&
      |Elements(pq)| == |Elements(pq')|
    decreases pq, 0
  {
    var left := ReplaceRoot(pq.left, y);
    var pq' := Node(pq.left.x, left, pq.right);
    ReplaceRootCorrect(pq.left, y); // so we know:
    assert Elements(pq.left) + multiset{y} == Elements(left) + multiset{pq.left.x};
    calc {
      Elements(pq) + multiset{y};
      == // def'n Elements
      multiset{pq.x} + Elements(pq.left) + Elements(pq.right) + multiset{y};
      ==
      multiset{pq.x} + Elements(pq.right) + (Elements(pq.left) + multiset{y});
      == // ReplaceRootCorrect
      multiset{pq.x} + Elements(pq.right) + Elements(left) + multiset{pq.left.x};
      == // def'n Elements
      Elements(pq') + multiset{pq.x};
    }
  }

  lemma ReplaceRootCorrectForRightMin(pq: PriorityQueue, y: int)
    requires Valid(pq) && !IsEmpty(pq)
    // And at this point in ReplaceRoot, we know:
    requires pq.left.Node? && pq.right.Node? && pq.right.x <= pq.left.x && pq.right.x < y
    ensures var pq' := ReplaceRoot(pq, y);
      Valid(pq') &&
      Elements(pq) + multiset{y} == Elements(pq') + multiset{pq.x} &&
      |Elements(pq)| == |Elements(pq')|
    decreases pq, 0
  {
    var right := ReplaceRoot(pq.right, y);
    var pq' := Node(pq.right.x, pq.left, right);
    ReplaceRootCorrect(pq.right, y); // so we know:
    assert Elements(pq.right) + multiset{y} == Elements(right) + multiset{pq.right.x};
    calc {
      Elements(pq) + multiset{y};
      == // def'n Elements
      multiset{pq.x} + Elements(pq.left) + Elements(pq.right) + multiset{y};
      ==
      multiset{pq.x} + Elements(pq.left) + (Elements(pq.right) + multiset{y});
      == // ReplaceRootCorrect
      multiset{pq.x} + Elements(pq.left) + Elements(right) + multiset{pq.right.x};
      == // def'n Elements
      Elements(pq') + multiset{pq.x};
    }
  }

  // Phew.

  ghost predicate IsMin(y: int, s: multiset<int>) {
    y in s && forall x :: x in s ==> y <= x
  }

  // This example is largely taken from the optional course textbook, and you
  // can find a detailed presentation of it there. Our rearrangement of the
  // proof is the result of debugging timeouts to understand how it could be
  // broken up effectively. A Dafny expert might use {:focus} or {:split_here}
  // annotations to break up the work without having to define the four
  // sub-lemmas with all their duplicative specifications.

  export
    provides PriorityQueue, Valid, Empty, IsEmpty, Insert, RemoveMin
    provides Elements, EmptyCorrect, IsEmptyCorrect, InsertCorrect, RemoveMinCorrect
    reveals IsMin
}

// It's interesting to consider alternative specification strategies for
// modules.  Here is an *intrinsic* restatement of PriorityQueue, meaning the
// functions themselves contain their correctness properties as postconditions,
// rather than using separate lemmas.  It's routine to wrap the implementation
// above in this style.
module PriorityQueueIntrinsic {
  export
    provides PriorityQueue, Valid, Empty, IsEmpty, Insert, RemoveMin
    provides Elements
    reveals IsMin
    provides PQ

  import PQ = PriorityQueue

  type PriorityQueue = PQ.PriorityQueue

  ghost predicate Valid(pq: PriorityQueue) {
    PQ.Valid(pq)
  }

  ghost function Elements(pq: PriorityQueue): multiset<int> {
    PQ.Elements(pq)
  }

  ghost predicate IsMin(m: int, s: multiset<int>) {
    PQ.IsMin(m, s)
  }

  function Empty(): PriorityQueue
    ensures var pq := Empty();
      Valid(pq) &&
      Elements(pq) == multiset{}
  {
    PQ.EmptyCorrect();
    PQ.Empty()
  }

  predicate IsEmpty(pq: PriorityQueue)
    requires Valid(pq)
    ensures IsEmpty(pq) <==> Elements(pq) == multiset{}
  {
    PQ.IsEmptyCorrect(pq);
    PQ.IsEmpty(pq)
  }

  function Insert(pq: PriorityQueue, y: int): PriorityQueue
    requires Valid(pq)
    ensures var pq' := Insert(pq, y);
      Valid(pq') &&
      Elements(pq') == multiset{y} + Elements(pq)
  {
    PQ.InsertCorrect(pq, y);
    PQ.Insert(pq, y)
  }

  function RemoveMin(pq: PriorityQueue): (int, PriorityQueue)
    requires Valid(pq) && !IsEmpty(pq)
    ensures var (y, pq') := RemoveMin(pq);
      Valid(pq') &&
      IsMin(y, Elements(pq)) &&
      Elements(pq') + multiset{y} == Elements(pq)
  {
    PQ.RemoveMinCorrect(pq);
    PQ.RemoveMin(pq)
  }
}

// Here's another fun example with a representation invariant: finite sets of
// integers as sorted sequences, allowing use of binary search.  We'll write
// this one in the intrinsic style to begin with.

module SortedListSet {
  export
    provides Set, Valid, Elements, Empty, Member, Add

  type Set = seq<int>

  // The invariant is exactly sortedness.
  ghost predicate Valid(s: Set) {
    forall i, j :: 0 <= i < j < |s| ==> s[i] < s[j]
  }

  ghost function Elements(s: Set): set<int> {
    set n | n in s
  }

  function Empty(): Set
    ensures var s := Empty();
      Valid(s) &&
      Elements(s) == {}
  {
    []
  }

  // Member goes through in a satisfyingly simple way.
  predicate Member(s: Set, n: int)
    requires Valid(s)
    ensures Member(s, n) <==> n in Elements(s)
  {
    if |s| == 0 then
      false
    else
      var mid := |s| / 2;
      if n == s[mid] then
        true
      else if n < s[mid] then
        Member(s[..mid], n)
      else
        Member(s[mid+1..], n)
  }

  // Add requires calling two symmetrical lemmas and making symmetrical
  // assertions.
  function Add(s: Set, n: int): Set
    requires Valid(s)
    ensures var s' := Add(s, n);
      Valid(s') &&
      Elements(s') == Elements(s) + {n}
  {
    if |s| == 0 then
      [n]
    else
      var mid := |s| / 2;
      if n == s[mid] then
        s
      else if n < s[mid] then
        var r := Add(s[..mid], n);
        ValidAddLeft(s[..mid], s[mid..], n, r);
        var r' := r + s[mid..];
        assert Elements(r') == Elements(r) + Elements(s[mid..]);
        r'
      else
        var r := Add(s[mid+1..], n);
        ValidAddRight(s[..mid+1], s[mid+1..], n, r);
        var r' := s[..mid+1] + r;
        assert Elements(r') == Elements(s[..mid+1]) + Elements(r);
        r'
  }

  // To get into the details, we name the situation when all elements of one
  // list are less than all elements of another.
  ghost predicate Below(s1: Set, s2: Set) {
    forall n, m :: n in s1 && m in s2 ==> n < m
  }

  lemma UseBelow(s1: Set, s2: Set, n: int, m: int)
    requires Below(s1, s2)
    requires n in s1
    requires m in s2
    ensures n < m
  {
  }

  // The next two lemmas show how Below is preserved by combining sets.

  lemma BelowUnionLeft(s1: Set, s2: Set, s: Set, s': Set)
    requires Below(s1, s)
    requires Below(s2, s)
    requires Elements(s1) + Elements(s2) == Elements(s')
    ensures Below(s', s)
  {
    forall n, m ensures n in s' && m in s ==> n < m {
      if n in s' && m in s {
        assert n in Elements(s');
      }
    }
  }

  lemma BelowUnionRight(s1: Set, s2: Set, s: Set, s': Set)
    requires Below(s, s1)
    requires Below(s, s2)
    requires Elements(s1) + Elements(s2) == Elements(s')
    ensures Below(s, s')
  {
    forall n, m ensures n in s && m in s' ==> n < m {
      if n in s && m in s' {
        assert m in Elements(s');
      }
    }
  }

  // And we come to these last two lemmas which abstract out details of the sets
  // involved in the two tricky cases of Insert.

  lemma ValidAddLeft(s1: Set, s2: Set, n: int, s1': Set)
    requires Valid(s1)
    requires Valid(s2)
    requires Below(s1, s2)
    requires Below([n], s2)
    requires Valid(s1')
    requires Elements(s1') == Elements(s1) + {n}
    ensures Valid(s1' + s2)
  {
    var s := s1' + s2;
    assert Below(s1', s2) by { BelowUnionLeft(s1, [n], s2, s1'); }
    forall i, j ensures 0 <= i < j < |s| ==> s[i] < s[j] {
      if 0 <= i < j < |s| {
        if i == |s1'| - 1 {
          UseBelow(s1', s2, s[i], s[j]);
        } else if i < |s1'| {
          if j < |s1'| {
          } else {
            assert s[i] in s1';
            assert s[j] in s2;
          }
        } else {
        }
      }
    }
  }

  lemma ValidAddRight(s1: Set, s2: Set, n: int, s2': Set)
    requires Valid(s1)
    requires Valid(s2)
    requires Below(s1, s2)
    requires Below(s1, [n])
    requires Valid(s2')
    requires Elements(s2') == Elements(s2) + {n}
    ensures Valid(s1 + s2')
  {
    var s := s1 + s2';
    assert Below(s1, s2') by { BelowUnionRight(s2, [n], s1, s2'); }
    forall i, j ensures 0 <= i < j < |s| ==> s[i] < s[j] {
      if 0 <= i < j < |s| {
        if i == |s1| - 1 {
          UseBelow(s1, s2', s[i], s[j]);
        } else if i < |s1| {
          if j < |s1| {
          } else {
            assert s[i] in s1;
            assert s[j] in s2';
          }
        } else {
        }
      }
    }
  }
}