quicksort.dfy

predicate sorted<T>(le: (T, T) -> bool, s: seq<T>)
{
  forall i, j :: 0 <= i <= j < |s| ==> le(s[i], s[j])
}

predicate hasSameElements<T>(a: seq<T>, b: seq<T>)
{
  multiset(a) == multiset(b)
}

predicate allBelow<T>(le: (T, T) -> bool, a: seq<T>, b: seq<T>)
{
  forall x, y :: x in a && y in b ==> le(x, y)
}

lemma ConcatSortedLemma<T>(le: (T, T) -> bool, a: seq<T>, b: seq<T>)
  requires sorted(le, a)
  requires sorted(le, b)
  requires allBelow(le, a, b)
  ensures sorted(le, a + b)
{
}

lemma Concat3SortedLemma<T>(le: (T, T) -> bool, a: seq<T>, b: T, c: seq<T>)
  requires forall a :: le(a, a)
  requires forall a, b, c :: le(a, b) && le(b, c) ==> le(a, c)  
  requires sorted(le, a)
  requires sorted(le, c)
  requires allBelow(le, a, [b])
  requires allBelow(le, [b], c)
  ensures sorted(le, a + [b] + c)
{
  ConcatSortedLemma(le, a, [b]);
  forall x, y | x in (a + [b]) && y in c
    ensures le(x, y)
  {
    assert le(x, b);
    assert le(b, y);
  }
  assert allBelow(le, a + [b], c);
  ConcatSortedLemma(le, a + [b], c);
}

method Partition<T>(le: (T, T) -> bool, a: array<T>, start: nat, stop: nat) returns (pivot: nat)
  // TODO: le probably needs to be a total order
  requires forall x, y :: le(x, y) || le(y, x)
  requires 0 <= start < stop <= a.Length
  modifies a
  ensures start <= pivot < stop
  ensures hasSameElements(a[..], old(a[..]))
  ensures a[..start] == old(a[..start])
  ensures hasSameElements(a[start..stop], old(a[start..stop]))
  ensures allBelow(le, a[start..pivot], [a[pivot]])
  ensures allBelow(le, [a[pivot]], a[pivot + 1..stop])
  ensures a[stop..] == old(a[stop..])
{
  // This algorithm is the core of quicksort.
  // The elements of a[start..stop] are split into two growing regions, a[start..lo] and a[hi..stop - 1].
  // The region between, a[lo..hi], contains unsorted elements. The pivot value is stashed in a[stop - 1] until the end.
  var lo := start;
  var hi := stop - 1;
  var pval := a[stop - 1];
  while lo < hi
    invariant start <= lo <= hi <= stop - 1
    invariant allBelow(le, a[start..lo], [pval])
    invariant allBelow(le, [pval], a[hi..stop - 1])
    invariant hasSameElements(a[..], old(a[..]))
    invariant hasSameElements(a[start..stop], old(a[start..stop]))
  {
    if le(pval, a[hi - 1]) {
      hi := hi - 1;
    } else {
      assert le(a[hi - 1], pval);
      a[lo], a[hi - 1] := a[hi - 1], a[lo];
      lo := lo + 1;
    }
  }
  a[stop - 1], a[hi] := a[hi], a[stop - 1];
  pivot := hi;
}

/*
lemma SortRangeLemma<T>(le: (T, T) -> bool, orig: seq<T>, part: seq<T>, a: array<T>, start: nat, pivot: nat, stop: nat)
  decreases stop - start
  // TODO: f probably needs to be a total order
  requires forall a :: le(a, a)
  requires forall a, b, c :: le(a, b) && le(b, c) ==> le(a, c)
  requires 0 <= start <= pivot < stop <= a.Length
  requires |orig| == |part| == a.Length
  requires hasSameElements(orig, part)
  requires orig[..start] == part[start..]
  requires orig[stop..] == part[stop..]
  requires allBelow(le, part[start..pivot], [part[pivot]])
  requires allBelow(le, [part[pivot]], part[pivot + 1..stop])
  requires hasSameElements(part, a[..])
  requires part[..start] == a[..start]
  requires part[stop..] == a[stop..]
  requires sorted(le, a[start..pivot])
  requires sorted(le, a[pivot..stop])
  ensures hasSameElements(orig, a[..])
  ensures orig[..start] == a[..start]
  ensures hasSameElements(orig[start..stop], a[start..stop])
  ensures sorted(le, a[start..stop])
  ensures orig[stop..] == a[stop..]
{
  Concat3SortedLemma(le, a[start..pivot], a[pivot], a[pivot + 1..stop]);
  assert a[start..pivot] + [a[pivot]] + a[pivot + 1..stop] == a[start..stop];
}

method SortRange<T>(le: (T, T) -> bool, a: array<T>, start: nat, stop: nat)
  decreases stop - start
  // TODO: f probably needs to be a total order
  requires forall a :: le(a, a)
  requires forall a, b :: le(a, b) || le(b, a)
  requires forall a, b, c :: le(a, b) && le(b, c) ==> le(a, c)
  requires 0 <= start <= stop <= a.Length
  modifies a
  ensures hasSameElements(a[..], old(a[..]))
  ensures a[..start] == old(a[..start])
  ensures hasSameElements(a[start..stop], old(a[start..stop]))
  ensures sorted(le, a[start..stop])
  ensures a[stop..] == old(a[stop..])
{
  if stop - start > 1 {
    var pivot := Partition(le, a, start, stop);
    ghost var partitioned := a[..];
    SortRange(le, a, start, pivot);
    SortRange(le, a, pivot + 1, stop);
    SortRangeLemma(le, old(a[..]), partitioned, a, start, pivot, stop);
  }
}

method QuickSort<T>(le: (T, T) -> bool, a: array<T>)
  // TODO: f probably needs to be a total order
  requires forall a :: le(a, a)
  requires forall a, b :: le(a, b) || le(b, a)
  requires forall a, b, c :: le(a, b) && le(b, c) ==> le(a, c)
  modifies a
  ensures hasSameElements(a[..], old(a[..]))
  ensures sorted(le, a[..])
{
  SortRange(le, a, 0, a.Length);
}

*/