(* Inverse of a permutation, in place

     Algorithm I
     The Art of Computer Programming, volume 1, Sec. 1.3.3, page 176

   The idea is to inverse each orbit at a time, using negative values to
   denote elements already inversed. The main loop scans the array and
   proceeds as follows: a negative value -x-1 is simply changed into x; a
   nonnegative value is the topmost element of a new orbit, which is
   inversed by the inner loop.

   Authors: Martin Clochard (École Normale Supérieure)
            Jean-Christophe Filliâtre (CNRS)
            Andrei Paskevich (Université Paris Sud)
*)

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

let function (~_) (x: int) : int = -x-1

(* To prove the termination of the inner loop, we count the number of
   nonnegative values in the array, using the following shortcut. *)
function numof (m: int -> int) (l r: int) : int =
  NumOf.numof (fun n -> m n >= 0) l r

predicate is_permutation (a: array int) =
  forall i. 0 <= i < length a ->
    0 <= a[i] < length a /\
    forall j. 0 <= j < length a -> i <> j -> a[i] <> a[j]

(* In order to use the same predicate "loopinvariant" in both loops,
   we generalize its definition wrt the paper, using an extra argument m'.

   In the inner loop, this argument is m-1 and the definition coincides with
   that of the paper.

   In the outer loop, this argument is m, which means that properties 4 and 6
   hold up to m included (instead of up to m excluded). *)
predicate loopinvariant (olda a: array int) (m m' n: int) =
    (forall e. 0 <= e < n -> -n <= a[e] < n)
 /\ (forall e. m < e < n -> 0 <= a[e])
 /\ (forall e. m < e < n -> olda[a[e]] = e)
 /\ (forall e. 0 <= e <= m' -> a[e] >= 0 -> a[e] = olda[e])
 /\ (forall e. 0 <= e <= m -> a[e] <= m)
 /\ (forall e. 0 <= e <= m' -> a[e] < 0 ->
      olda[~a[e]] = e /\ (~a[e] <= m -> a[~a[e]] < 0))

let inverse_in_place (a: array int) : unit
  requires { is_permutation a }
  ensures  { is_permutation a }
  ensures  { forall i. 0 <= i < length a -> (old a)[a[i]] = i }
= let n = length a in
  for m = n - 1 downto 0 do
    invariant { loopinvariant (old a) a m m n }
    let ref i = a[m] in
    if i >= 0 then begin
      (* unrolled loop once *)
      a[m] <- -1;
      let ref j = (~m) in
      let ref k = i in
      i := a[i];
      while i >= 0 do
        invariant { a[k] = i <= m /\ 0 <= k <= m /\ ~m <= j < 0 /\
                    (old a)[~ j] = k /\ a[~ j] < 0 /\ a[m] < 0 }
        invariant { forall e. 0 <= e < m -> a[e] < 0 -> a[e] <> j }
        invariant { loopinvariant (old a) a m (m-1) n }
        variant   { numof a.elts 0 n }
        a[k] <- j;
        j := ~ k;
        k := i;
        i := a[k]
      done;
      assert { k = m };
      i := j
    end;
    assert { (old a)[~ i] = m };
    a[m] <- ~ i
  done