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]}
        if a[i] < a[n-1] then
          let aux = ref a[i] in
            a[i] <- a[!lt];
            a[!lt] <- !aux;
            lt := !lt + 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