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 1 in
      for i = 1 to n - 1 do
        invariant{forall k: int. 0 <= k < i -> a[k] <= a[0]}
        if a[i] < a[0] then
          let aux = ref a[i] in
            a[i] <- a[!lt];
            a[!lt] <- !aux;
            lt := !lt + 1;
      done;
      let aux = ref a[0] in
        a[0] <- a[!lt];
        a[!lt] <- !aux;
    assert{partitioned a !lt n};
    (a, lt)
end