(* finite sets *)

type 'a set
logic set_empty : 'a set
logic set_add : 'a, 'a set -> 'a set
logic set_rmv : 'a, 'a set -> 'a set

logic In : 'a, 'a set -> prop
predicate Is_empty(s : 'a set) = forall x:'a. not In(x, s)
predicate Incl(s1 : 'a set, s2 : 'a set) = forall x:'a. In(x, s1) -> In(x, s2)

axiom set_empty_def : Is_empty(set_empty)

axiom set_add_def : 
  forall x:'a. forall y:'a. forall s:'a set. 
  In(x, set_add(y,s)) <-> (x = y or In(x, s))
axiom set_rmv_def :
  forall x:'a. forall y:'a. forall s:'a set. 
  In(x, set_rmv(y,s)) <-> (x <> y and In(x, s))

logic set_card : 'a set -> int

axiom card_nonneg : 
  forall s:'a set. set_card(s) >= 0

axiom card_set_add : 
  forall x:'a. forall s:'a set. 
  not In(x,s) -> set_card(set_add(x,s)) = 1 + set_card(s)

axiom card_set_rmv : 
  forall x:'a. forall s:'a set. 
  In(x,s) -> set_card(s) = 1 + set_card(set_rmv(x, s))

axiom card_Incl :
  forall s1,s2 : 'a set. Incl(s1,s2) -> set_card(s1) <= set_card(s2)

(* iteration on a set *)

parameter set_has_next : 
  s:'a set ref -> 
    {} bool reads s { if result then Is_empty(s) else not Is_empty(s) }

parameter set_next :
  s:'a set ref -> 
  { not Is_empty(s) } 'a writes s { In(result,s@) and s = set_rmv(result,s@) }

(* maps *)

type ('a, 'b) map

logic map_get : ('a,'b) map, 'a -> 'b
logic map_set : ('a,'b) map, 'a, 'b -> ('a, 'b) map

axiom map_get_set_eq :
  forall m : ('a,'b) map . forall i : 'a. forall v : 'b.
  map_get(map_set(m,i,v), i) = v

axiom map_get_set_neq : 
  forall m : ('a, 'b) map. forall i : 'a. forall j : 'a. forall v : 'b.
  i <> j -> map_get(map_set(m,i,v), j) = map_get(m,j)

(* graph *)

type vertex

logic V : vertex set

logic g_succ : vertex -> vertex set

axiom g_succ_sound : 
  forall x:vertex. Incl(g_succ(x), V)

logic weight : vertex, vertex -> int (* edge weight, if there is an edge *)

axiom weight_nonneg : forall x,y:vertex. weight(x,y) >= 0

parameter eq_vertex : 
  x:vertex -> y:vertex -> {} bool { if result then x=y else x<>y }

(* visited vertices *)

parameter S : vertex set ref

parameter S_add : 
  x:vertex -> {} unit writes S { S = set_add(x, S@) }

(* current distances *)

parameter d : (vertex, int) map ref

(* priority queue *)

parameter Q : vertex set ref

parameter Q_is_empty : 
  unit -> 
    {} bool reads Q { if result then Is_empty(Q) else not Is_empty(Q) }

parameter init :
  src:vertex ->
    {} 
    unit writes S,Q,d 
    { Is_empty(S) and 
      Q = set_add(src, set_empty) and 
      d = map_set(d@, src, 0)  }

parameter relax :
  u:vertex -> v:vertex -> 
    {} 
    unit reads S writes d,Q 
    { (In(v,S) and Q = Q@ and d = d@) 
      or
      (In(v,Q) and map_get(d,u) + weight(u,v) >= map_get(d,v) and 
        Q = Q@ and d = d@) 
      or	
      (In(v,Q) and map_get(d,u) + weight(u,v) < map_get(d,v) and 
        Q = Q@ and d = map_set(d@, v, map_get(d,u) + weight(u,v)))
      or
      (not In(v,S) and not In(v,Q) and Q = set_add(v,Q@) and
        d = map_set(d@, v, map_get(d,u) + weight(u,v))) }	

predicate Min(m:vertex, Q:vertex set, d:(vertex, int) map) =
  In(m, Q) and forall x:vertex. In(x, Q) -> map_get(d,x) <= map_get(d,m)

parameter Q_extract_min :
  unit ->
    { not Is_empty(Q) } 
    vertex reads d writes Q 
    { Min(result, Q@, d) and Q = set_rmv(result, Q@) }

(* paths and shortest paths 
   path(x,y,d) = 
	there is a path from x to y of length d
   shortest_path(x,y,d) = 
	there is a path from x to y of length d, and no shorter path *)

logic path : vertex, vertex, int -> prop 

axiom path_nil : 
  forall x:vertex. path(x,x,0)

axiom path_cons : 
  forall x,y,z:vertex. forall d:int. 
    path(x,y,d) -> In(z,g_succ(y)) -> path(x,z,d+weight(y,z))

axiom length_nonneg : forall x,y:vertex. forall d:int. path(x,y,d) -> d >= 0

predicate shortest_path(x:vertex, y:vertex, d:int) =
  path(x,y,d) and forall d':int. path(x,y,d') -> d <= d'

axiom path_inversion :
  forall src,v:vertex. forall d:int. path(src,v,d) ->
    (v = src and d = 0) or
    (exists v':vertex. path(src,v',d - weight(v',v)) and In(v,g_succ(v')))

axiom path_shortest_path :
  forall src,v:vertex. forall d:int. path(src,v,d) ->
    exists d':int. shortest_path(src,v,d') and d' <= d

(* lemmas (those requiring induction) *)

axiom main_lemma :
  forall src,v:vertex. forall d:int. 
  path(src,v,d) -> not shortest_path(src,v,d) ->
    exists v':vertex. exists d':int. 
      shortest_path(src,v',d') and In(v,g_succ(v')) and d' + weight(v',v) < d

axiom completeness_lemma :
  forall s:vertex set. forall src:vertex. forall dst:vertex. forall d:int. 
    (* if s is closed under g_succ *)
    (forall v:vertex. In(v,s) -> forall w:vertex. In(w,g_succ(v)) -> In(w,s))
    ->
    (* and if s contains src *)
    In(src, s) -> 
    (* then any vertex reachable from s is also in s *)
    path(src, dst, d) -> In(dst, s)

(* definitions used in loop invariants *)

predicate Inv_src(src:vertex, S:vertex set, Q:vertex set) =
  In(src,S) or In(src,Q)

predicate Inv(src:vertex, S:vertex set, Q:vertex set, d:(vertex,int) map) =
  Inv_src(src, S, Q)
  (* S,Q are contained in V *)
  and Incl(S,V) and Incl(Q,V)
  (* S and Q are disjoint *)
  and 
  (forall v:vertex. In(v,Q) -> In(v,S) -> false)
  (* we already found the shortest paths for vertices in S *)
  and
  (forall v:vertex. In(v,S) -> shortest_path(src,v,map_get(d,v)))
  (* there are paths for vertices in Q *)
  and
  (forall v:vertex. In(v,Q) -> path(src,v,map_get(d,v)))
  (* vertices at distance < min(Q) are already in S *)
  and
  (forall m:vertex. Min(m,Q,d) -> 
     forall x:vertex. forall dx:int. shortest_path(src,x,dx) -> 
       dx < map_get(d,m) -> In(x,S)) 

predicate InvSucc(src:vertex, S:vertex set, Q: vertex set) =
  (* successors of vertices in S are either in S or in Q *)
  (forall x:vertex. In(x, S) -> 
     forall y:vertex. In(y,g_succ(x)) -> In(y, S) or In(y, Q))

predicate InvSucc2(src:vertex, S:vertex set, Q: vertex set,
	           u:vertex, su:vertex set) =
  (* successors of vertices in S are either in S or in Q,
     unless they are successors of u still in su *)
  (forall x:vertex. In(x, S) -> 
     forall y:vertex. In(y,g_succ(x)) -> 
       (x<>u or (x=u and not In(y,su))) -> In(y, S) or In(y, Q))

(* code *)

let shortest_path (src:vertex) (dst:vertex) =
  { In(src,V) and In(dst,V) }
  init src;
  while not (Q_is_empty void) do
    { invariant 
	Inv(src, S, Q, d) and InvSucc(src, S, Q)
      variant 
	set_card(V) - set_card(S) }
    let u = Q_extract_min void in
    assert { shortest_path(src, u, map_get(d,u)) };
    S_add u;
    let su = ref (g_succ u) in
    while not (set_has_next su) do 
      { invariant 
	  Incl(su,g_succ(u)) and 
	  Inv(src, S, Q, d) and InvSucc2(src, S, Q, u, su)
        variant 
	  set_card(su) }
      let v = set_next su in
      relax u v
    done
  done
  { (forall v:vertex. In(v,S) -> shortest_path(src,v,map_get(d,v))) 
    and
    (forall v:vertex. not In(v,S) -> forall dv:int. not path(src,v,dv)) }
    
(*
Local Variables: 
compile-command: "gwhy dijkstra.why"
End: 
*)