module RefinedSpec

use export int.Int
use import map.Map as M
use import set.Fset as F
use seq.Seq, seq.Mem, seq.Distinct

predicate same_ext (m1 m2: 'a -> 'b) = forall x: 'a. m1 x = m2 x

(* auxiliary lemmas on sequences *)
lemma append_empty: forall s: seq 'a.
  s ++ empty == s

lemma empty_length: forall s: seq 'a.
  length s = 0 <-> s == empty

predicate disjoint_seq (s1 s2: seq 'a) =
  forall i j. 0 <= i < length s1 ->
		0 <= j < length s2 -> s1[i] <> s2[j]

(* Arbitrary type for a tree node *)
type elt

(* Verifies if two nodes are equal *)
val equal (e1 e2 : elt) : bool
  ensures { result <-> e1 = e2 }

(* Indicates if two nodes are connected by an edge in the tree *)
predicate edge (x y : elt) (f : elt -> elt) =
  x <> y /\ f x = y

(* recursive predicate for expressing a path between two nodes *)
(* in the main text a few cosmetic changes were done, namely *)
(* rename of -f- to -parent- and expand the edge definition *)
predicate path (f: elt -> elt) (x y: elt) (p: seq elt) =
  let n = length p in
  n = 0 /\ x = y
  \/
  n > 0 /\
  p[0] = x /\
  edge p[n - 1] y f /\
  distinct p /\
  (forall i. 0 <= i < n - 1 -> edge p[i] p[i + 1] f) /\
  (forall i. 0 <= i < n     -> p[i] <> y)

predicate reachability (f: elt -> elt) (x y: elt) =
  exists p. path f x y p

(* If there is an edge between nodes x and y
 the path is defined as the singleton seq with node x  *)
axiom path_to_parent: forall x y : elt, f : elt -> elt.
  edge x y f -> path f x y (cons x empty)

(* If there is a path from [from] to [middle] and a path from [middle] to 
  [until] then there is a path from [from] to [until] *)
axiom path_transitivity: forall from middle until f pth1 pth2.
  path f from middle pth1 -> path f middle until pth2 ->
  disjoint_seq pth1 pth2 -> from <> until ->
  (forall j. 0 <= j < length pth1 -> pth1[j] <> until) ->
	path f from until (pth1 ++ pth2)

(* Recursive path composition *)
axiom path_composition: forall n x y: elt, f : elt -> elt, pth : seq elt.
  n <> y -> not (mem y pth) ->
  distinct (snoc pth x) -> path f n x pth -> edge x y f ->
  path f n y (snoc pth x)

 (* If there is a path between two nodes, that path is unique *)
axiom path_uniqueness: forall x y: elt, f: elt -> elt,
			      pth1 pth2: seq elt.
  path f x y pth1 -> path f x y pth2 -> pth1 == pth2

 (* If node np is not reachable to node c, then np will
    not belong to any path that contains node c *)
axiom path_exclusion: forall f x c np p.
  not (reachability f np c) -> path f x c p -> not (mem np p)

 (* Given a path between two nodes, there is no overlap between any two consecutive subpaths *)
axiom path_separation: forall final initial middle : elt, f : elt -> elt,
		       	      p1 p2 : seq elt.
  path f middle final p2 -> path f initial middle p1 ->
  final <> initial -> middle <> initial -> middle <> final ->
  disjoint_seq p1 p2

constant n: elt (* constant used for defining a state witness *)

type state = {
  (* parent relation: up-pointers to direct ancestor *)
  mutable parent     : elt -> elt;
  (* parent root *)
  mutable root 	 : elt;
  (* nodes in the parent *)
  mutable nodes      : fset elt;
  (* rank for each node *)
  mutable rank       : elt -> int;
  (* tombstone nodes  *)
  mutable tombstones : fset elt;
  (* ancestors relation: all ancestors of a node *)
  mutable ancestors  : elt -> fset elt;
} invariant { F.mem root nodes }
  invariant { parent root = root }
  invariant { forall x. F.mem x nodes -> F.mem (parent x) nodes }
  invariant { forall x. F.mem x nodes -> reachability parent x root }
  invariant { forall x. F.mem x nodes ->
		reachability parent root x -> x = root }
  invariant { forall x y. F.mem x nodes /\ F.mem y nodes /\ x <> root /\
		x <> y /\ F.mem y (ancestors x) -> rank x > rank y }
  invariant { forall x. F.mem x nodes /\ x <> root ->
		ancestors x = F.add (parent x) (ancestors (parent x))}
by { parent = (fun _ -> n); root = n; nodes = F.singleton n ;
     rank = (fun _ -> 1); tombstones = F.empty;
     ancestors = (fun _ -> F.empty) }

(* Paths already present in the tree remain in the tree after executing 
   the add operation *)
axiom remaining_nodes_add: forall n w p: elt, s: state, l: seq elt.
  path s.parent w s.root l -> not (mem n l) ->
  F.mem w s.nodes -> F.mem p s.nodes -> not (F.mem n s.nodes) ->
  w <> n -> n <> p -> path (M.set s.parent n p) w s.root l

(* Descendants of the node being moved continue to be its descendants *)
axiom descendants_move: forall x c np: elt, f: elt -> elt, p: seq elt.
  x <> np -> c <> np -> x <> c -> not (reachability f np c) ->
  path f x c p -> distinct (cons c p) -> not (mem np p) ->
  (path (M.set f c np) x c p)

(*  Paths nodes unreachable to the node being moved are not affected *)
axiom remaining_nodes_move: forall x c np: elt, s: state, p: seq elt.
  c <> np -> x <> c -> not (reachability s.parent np c) ->
  path s.parent x s.root p -> (not reachability s.parent x c) ->
  distinct p -> (path (M.set s.parent c np) x s.root p)

val ghost add (n p : elt) (s : state) : unit
  requires { [@expl:add1] not F.mem n s.nodes }
  requires { [@expl:add2] F.mem p s.nodes }
  writes   { s.parent, s.nodes, s.rank, s.ancestors }
  ensures  { s.parent = M.set (old s.parent) n p }
  ensures  { edge n p s.parent }
  ensures  { s.nodes = F.add n (old s).nodes }
  ensures  { s.rank = M.set (old s).rank n ((s.rank p) + 1) }
  ensures  { s.ancestors = M.set ((old s).ancestors) n
		(F.add p ((old s).ancestors p)) }

val ghost remove (n : elt) (s : state) : unit
  requires { n <> s.root }
  writes   { s.tombstones }
  ensures  { s.tombstones = F.add n (old s).tombstones }

val ghost move (c np : elt) (s : state) : unit
  requires { F.mem np s.nodes }
  requires { F.mem c s.nodes }
  requires { not (reachability s.parent np c) }
  requires { c <> s.root }
  requires { c <> np }
  writes   { s.parent, s.ancestors, s.rank }
  ensures  { [@expl:post1] s.parent = M.set (old s.parent) c np }

let ghost move_refined (c1 np1 c2 np2 : elt) (pr1 pr2 : int)
		       (s : state) : unit
  requires { pr1 <> pr2 }
  requires { F.mem c1 s.nodes /\ F.mem c2 s.nodes }
  requires { F.mem np1 s.nodes /\ F.mem np2 s.nodes }
  requires { not (reachability s.parent np1 c1) /\
  	     not (reachability s.parent np2 c2) }
  requires { c1 <> s.root /\ c2 <> s.root }
  requires { c1 <> np1 /\ c2 <> np2 }
  ensures  { (s.parent = M.set (old s.parent) c2 np2) \/
             (s.parent = M.set (old s.parent) c1 np1) \/
             (same_ext s.parent (old s).parent) }
= if (equal c1 c2) then
  if (pr1 < pr2) then move c2 np2 s else move c1 np1 s
  else if (s.rank np1 < s.rank c1) then move c1 np1 s
  else if (F.mem c2 (diff (F.add np1 (s.ancestors np1))
                           (s.ancestors c1))) then
  if (s.rank np2 < s.rank c2) then ()
  else if (pr1 < pr2) then move c2 np2 s else move c1 np1 s
  else move c1 np1 s

end