vacid_0_union_find.mlw 3.54 KB
Newer Older
1
module M
2

3 4 5
  use int.Int
  use ref.Ref
  use array.Array
6

7
  type uf = { link : array int;
8
        ghost dist : array int; (* distance to representative *)
9
              num  : int;       (* number of classes *) }
10

Andrei Paskevich's avatar
Andrei Paskevich committed
11
  function size (u: uf) : int = length u.link
12

Andrei Paskevich's avatar
Andrei Paskevich committed
13
  predicate inv (u : uf) =
14
    let s = length u.link in
Andrei Paskevich's avatar
Andrei Paskevich committed
15 16
    length u.dist = s /\
    (forall i:int. 0 <= i < s -> 0 <= u.link[i] < s) /\
17
    (forall i:int. 0 <= i < s ->
Andrei Paskevich's avatar
Andrei Paskevich committed
18 19
       (  (u.dist[i] = 0 /\ u.link[i] = i)
       \/ (u.dist[i] > 0 /\ u.dist[u.link[i]] < u.dist[i])))
20 21 22 23 24 25 26 27

  inductive repr (u:uf) (x:int) (r:int) =
  | Repr_root: forall u:uf. inv u ->
               forall x:int. 0 <= x < size u ->
               (link u)[x] = x -> repr u x x
  | Repr_link: forall u:uf. inv u ->
               forall x:int. 0 <= x < size u ->
               forall r:int. repr u (link u)[x] r -> repr u x r
28

29 30 31
  lemma Repr_bounds_2:
    forall u:uf, x y:int. repr u x y -> 0 <= y < size u
  lemma Repr_dist_1:
32
    forall u:uf, x y:int. repr u x y -> u.dist[y] = 0
33

Andrei Paskevich's avatar
Andrei Paskevich committed
34
  predicate same (u:uf) (x y:int) =
35 36
    forall r:int. repr u x r <-> repr u y r

Andrei Paskevich's avatar
Andrei Paskevich committed
37
  predicate same_reprs (u1 u2 : uf) =
38
    forall x r:int. repr u1 x r <-> repr u2 x r
39

40
  axiom OneClass :
41
    forall u:uf. num u = 1 ->
42
    forall x y:int. 0 <= x < size u -> 0 <= y < size u -> same u x y
43

44
  let get_num_classes (u: uf) ensures { result = num u } = num u
45

46 47 48 49 50 51
  let create (n:int)
    requires { 0 <= n }
    ensures { inv result /\
      num result = n /\ size result = n /\
      forall x:int. 0 <= x < n -> repr result x x }
  = let l = make n 0 in
52
    for i = 0 to n-1 do
53
      invariant { forall j:int. 0 <= j < i -> l[j] = j }
54
      l[i] <- i
55
    done;
56
    { link = l; dist = make n 0; num = n }
57

58 59 60
  let path_compression (u: uf) x r
    requires { inv u /\ 0 <= x < size u /\ u.dist[x] > 0 /\ repr u x r }
    ensures { inv u /\ size u = size (old u) /\
Andrei Paskevich's avatar
Andrei Paskevich committed
61
      num u = num (old u) /\ same_reprs (old u) u }
62 63
  = u.link[x] <- r;
    u.dist[x] <- 1
64

65 66 67 68 69 70 71
  let rec find (u: uf) (x: int) variant { u.dist[x] }
    requires { inv u /\ 0 <= x < size u }
    ensures { inv u /\
      repr u x result /\
      size u = size (old u) /\ num u = num (old u) /\
      same_reprs u (old u) }
  = let y = (link u)[x] in
72 73 74 75 76 77 78 79 80
    if y <> x then begin
      let r = find u y in
      path_compression u x r;
      r
    end else
      x

  (***

Andrei Paskevich's avatar
Andrei Paskevich committed
81
  val ghost_find (u:ref uf) (x:int) : int
82 83
    requires { inv !u /\ 0 <= x < size !u }
    ensures { repr !u x result }
84

85 86 87 88 89
  let increment u r
    requires { inv !u /\ 0 <= r < size !u }
    ensures { forall i:int. 0 <= i < size !u ->
        result#i = (dist !u)#i + if repr !u i r then 1 else 0 }
  = let i = ref 0 in
90 91
    let d = ref (dist !u) in
    while !i < size !u do
Andrei Paskevich's avatar
Andrei Paskevich committed
92
      invariant { 0 <= !i <= size !u /\
93 94
                  forall j:int. 0 <= j < size !u ->
                  !d#j = dist(!u)#j +
Andrei Paskevich's avatar
Andrei Paskevich committed
95
                         if repr !u j r /\ j < !i then 1 else 0 }
96 97 98 99 100 101
      variant { size !u - !i }
      if ghost_find u !i = r then
        d := A.set !d !i (A.get !d !i + 1)
    done;
    !d

102 103
  let union u a b
    requires { inv !u /\
Andrei Paskevich's avatar
Andrei Paskevich committed
104
      0 <= a < size !u /\ 0 <= b < size !u /\ not same !u a b }
105
    ensures { inv !u /\
Andrei Paskevich's avatar
Andrei Paskevich committed
106 107
      same !u a b /\
      size !u = size (old !u) /\ num !u = num (old !u) - 1 /\
108 109
      forall x y:int. 0 <= x < size !u -> 0 <= y < size !u ->
        same !u x y <->
Andrei Paskevich's avatar
Andrei Paskevich committed
110 111 112
        same (old !u) x y \/
        same (old !u) x a /\ same (old !u) b y \/
        same (old !u) x b /\ same (old !u) a y }
113 114 115 116 117
  = let ra = find u a in
    let rb = find u b in
    let l = link !u in
    let d = increment u ra in
    u := UF (A.set l ra rb) d (size !u) (num !u - 1)
118 119

  ***)
120

121
end