module Partition

  use import int.Int
  use import ref.Ref
  use import array.Array

  predicate partitioned (a: array int) (pivot: int) (n: int) =
      forall i:int. (0 <= i <= pivot ->  a[i] <= a[pivot]) /\
                    (pivot < i < n -> a[i] > a[pivot])

  let partition (a: array int) (n: int)
    requires { 2 <= n = length a }
  = let lt = ref 0 in
      let aux = ref a[0] in
        a[0] <- a[n-1];
        a[n-1] <- !aux;
      for i = !lt to n - 2 do
        invariant{0 <= !lt <= i /\
            forall k: int. (0 <= k < !lt -> a[k] <= a[n-1]) /\
                           (!lt <= k < i -> a[k] > a[n-1])}
        if a[i] <= a[n-1] then
          let aux = ref a[i] in
            a[i] <- a[!lt];
            a[!lt] <- !aux;
            lt := !lt + 1;
        else
            assert{ !lt <= i /\ i < n-1 /\ a[i] > a[n-1]}
      done;
      let aux = ref a[n-1] in
        a[n-1] <- a[!lt];
        a[!lt] <- !aux;
    assert{ partitioned a !lt n };
    (a, lt)
end