(* Union-Find

   Authors: Martin Clochard (Univ Paris Sud)
            Jean-Christophe Filliâtre (CNRS)
	    Mário Pereira (Univ Paris Sud)
	    Simão Melo de Sousa (University Beira Interior)
*)

use int.Int
use int.MinMax
use map.Map
use array.Array
use ref.Ref

type uf = {
          size: int;
          link: array int;
          rank: array int;
  mutable ghost rep : int -> int;
  mutable ghost dst : array int; (* fields dst and maxd are used to prove *)
  mutable ghost maxd: int;       (* termination; see find's variant below *)
}
  invariant { 0 <= size = length link = length rank = length dst }
  invariant { forall x. 0 <= x < size -> 0 <= link[x] < size }
  invariant { forall x. 0 <= x < size -> 0 <= rep x < size }
  invariant { forall x. 0 <= x < size -> rep (rep x) = rep x }
  invariant { forall x. 0 <= x < size -> x = link[x] <-> rep x = x }
  invariant { forall x. 0 <= x < size -> let y = link[x] in
              x <> y -> rep x = rep y /\ dst[x] < dst[y] }
  invariant { 0 <= maxd }
  invariant { forall x. 0 <= x < size -> dst[x] <= maxd }
  by { size = 0; link = make 0 0; rank = make 0 0;
       rep  = (fun x -> x); dst = make 0 0; maxd = 0; }

let create (n: int) : uf
  requires { 0 <= n }
= let l = Array.make n 0 in
  for i = 0 to n - 1 do invariant { forall j. 0 <= j < i -> l[j] = j }
    l[i] <- i
  done;
  { size = n;
    link = l;
    rank = Array.make n 0;
    rep  = (fun x -> x);
    dst  = Array.make n 0;
    maxd = 0; }

(* with path compression *)
let rec find (uf: uf) (x: int) : int
  requires { 0 <= x < uf.size }
  variant  { uf.maxd - uf.dst[x] }
  ensures  { result = uf.rep x }
  ensures  { uf.dst[result] >= uf.dst[x] }
= let y = uf.link[x] in
  if y = x then x
  else begin let rx = find uf y in uf.link[x] <- rx; rx end

let same (uf: uf) (x: int) (y: int) : bool
  requires { 0 <= x < uf.size }
  requires { 0 <= y < uf.size }
  ensures  { result <-> uf.rep x = uf.rep y }
= let a = find uf x in
  let b = find uf y in
  a = b

let ghost inc_dst (dst: int -> int) (x y: int) : int -> int
  ensures { result y = 1 + max (dst x) (dst y) }
  ensures { forall z. z <> y -> result z = dst z }
= Map.set dst y (1 + max (dst x) (dst y))

let union (uf: uf) (x y: int) : ghost int
  requires { 0 <= x < uf.size }
  requires { 0 <= y < uf.size }
  ensures  { result = old (uf.rep x) || result = old (uf.rep y) }
  ensures  { forall z. 0 <= z < uf.size ->
             uf.rep z = if old (uf.rep z = uf.rep x || uf.rep z = uf.rep y)
                        then result else old (uf.rep z) }
= let x = find uf x in
  let y = find uf y in
  if x = y then x else
  let rkx = uf.rank[x] in
  let rky = uf.rank[y] in
  if rkx < rky then begin
    uf.link[x] <- y;
    uf.rep <- pure { fun z -> if uf.rep z = uf.rep x then y else uf.rep z };
    uf.dst[y] <- 1 + max uf.dst[x] uf.dst[y];
    uf.maxd <- uf.maxd + 1;
    y
  end else begin
    uf.link[y] <- x;
    uf.rep <- pure { fun z -> if uf.rep z = uf.rep y then x else uf.rep z };
    uf.dst[x] <- 1 + max uf.dst[x] uf.dst[y];
    uf.maxd <- uf.maxd + 1;
    if rkx = rky then uf.rank[x] <- rkx + 1;
    x
  end

(* Though this is not needed, one can recover from above a global
   invariant that states the existence of a path from each element to
   its representative. *)

inductive path (link: array int) (x y: int) =
  | Path0: forall link x  . path link x x
  | Path1: forall link x y. path link link[x] y -> path link x y

let rec lemma path_rep (uf: uf) (x: int)
  requires { 0 <= x < uf.size }
  ensures  { path uf.link x (uf.rep x) }
  variant  { uf.maxd - uf.dst[x] }
= if uf.link[x] <> x then path_rep uf uf.link[x]