(* Building a perfect maze using the union-find data structure
   (from the VACID-0 Benchmarks http://vacid.codeplex.com/)

   Jean-Christophe Filliâtre (CNRS)
   Andrei Paskevich (Univ Paris Sud) *)

theory UnionFind_pure

  use int.Int

  type uf_pure

  predicate repr uf_pure int int
  function size uf_pure : int
  function num uf_pure : int

  axiom Repr_function_1:
    forall u:uf_pure, x:int.
    0 <= x < size u -> exists y:int. 0 <= y < size u /\ repr u x y

  axiom Repr_function_2:
    forall u:uf_pure, x y z:int.
    0 <= x < size u -> repr u x y -> repr u x z -> y = z

  predicate same (u:uf_pure) (x y:int) =
    forall r:int. repr u x r <-> repr u y r

  predicate same_reprs (u1 u2 : uf_pure) =
     forall x r:int. repr u1 x r <-> repr u2 x r

  axiom OneClass :
    forall u:uf_pure. num u = 1 ->
    forall x y:int. 0 <= x < size u -> 0 <= y < size u -> same u x y

end

module UnionFind_sig

  use int.Int
  use export UnionFind_pure

  type uf = private { ghost mutable state : uf_pure }

  val create (n:int) : uf
    requires { 0 <= n }
    ensures { num result.state = n /\ size result.state = n /\
      forall x:int. 0 <= x < n -> repr result.state x x }

  val find (u:uf) (x:int) : int writes {u}
    requires { 0 <= x < size u.state }
    ensures { repr u.state x result /\
      size u.state = size (old u.state) /\
      num u.state = num (old u.state) /\
      same_reprs u.state (old u.state) }

  val union (u:uf) (a:int) (b:int) : unit writes {u}
    requires { 0 <= a < size u.state /\
               0 <= b < size u.state /\ not same u.state a b }
    ensures { same u.state a b /\
      size u.state = size (old u.state) /\
      num u.state = num (old u.state) - 1 /\
      forall x y:int. 0 <= x < size u.state -> 0 <= y < size u.state ->
        same u.state x y <->
        same (old u.state) x y \/
        same (old u.state) x a /\ same (old u.state) b y \/
        same (old u.state) x b /\ same (old u.state) a y }

  val get_num_classes (u:uf) : int ensures { result = num u.state }

end

theory Graph

  type vertex

  type graph

  inductive path graph vertex vertex =
    | Path_refl : forall g:graph, x:vertex. path g x x
    | Path_sym  : forall g:graph, x y:vertex. path g x y -> path g y x
    | Path_trans:
        forall g:graph, x y z:vertex. path g x y -> path g y z -> path g x z

end

module Graph_sig

  use int.Int
  use ref.Ref

  clone export Graph with type vertex = int

  val graph : ref graph

  val num_edges : ref int

  val add_edge (a:int) (b:int) : unit writes {num_edges,graph}
    requires { not path !graph a b }
    ensures { !num_edges = old !num_edges + 1 /\
      (forall x y:int.
        path !graph x y <->
        path (old !graph) x y \/
        path (old !graph) x a /\ path (old !graph) b y \/
        path (old !graph) x b /\ path (old !graph) a y) }

end

module BuildMaze

  use int.Int
  use ref.Ref

  use UnionFind_sig
  use Graph_sig

  val rand (s:int) : int requires { 0 < s } ensures { 0 <= result < s }

  lemma Ineq1 :
    forall n x y:int. 0 <= n -> 0 <= x < n -> 0 <= y < n -> 0 <= x*n+y < n*n

  let add_edge_and_union (u: uf) (a:int) (b:int)
    requires { 0 <= a < size u.state /\ 0 <= b < size u.state /\
      not same u.state a b /\ not path !graph a b /\
      forall x y:int.
        0 <= x < size u.state -> 0 <= y < size u.state ->
        same u.state x y <-> path !graph x y }
    ensures { !num_edges = old !num_edges + 1 /\
      same u.state a b /\
      size u.state = size (old u.state) /\
      num u.state = num (old u.state) - 1 /\
      (forall x y:int.
         0 <= x < size u.state -> 0 <= y < size u.state ->
         same u.state x y <-> path !graph x y) }
  = add_edge a b;
    union u a b

let build_maze (n:int)
  requires { 1 <= n /\ !num_edges = 0 /\
    forall x y:int. x=y <-> path !graph x y }
  diverges (* termination is only almost sure,
              assuming rand is well distributed.
              See http://toccata.lri.fr/gallery/RandomWalk.en.html *)
  ensures { !num_edges = n*n-1 /\
    forall x y:int. 0 <= x < n*n -> 0 <= y < n*n -> path !graph x y }
  = let u = create (n*n) in
    assert { forall x y:int. 0 <= x < n*n -> 0 <= y < n*n ->
      same u.state x y -> (repr u.state x y && repr u.state x x && x = y) };
    while get_num_classes u > 1 do
      invariant {
        1 <= num u.state /\
        num u.state + !num_edges = size u.state = n*n /\
        forall x y:int. 0 <= x < n*n -> 0 <= y < n*n ->
                        (same u.state x y <-> path !graph x y) }
      let x = rand n in
      let y = rand n in
      let d = rand 2 in
      let w = if d = 0 then x+1 else x in
      let z = if d = 0 then y else y+1 in
      if w < n && z < n then begin
        let a = y * n + x in
        assert { 0 <= a < n*n };
        let b = w * n + z in
        assert { 0 <= b < n*n };
        if find u a <> find u b then
          add_edge_and_union u a b
      end
    done

end