Source file store.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
(**************************************************************************)
(*                                                                        *)
(*  SPDX-License-Identifier LGPL-2.1                                      *)
(*  Copyright (C)                                                         *)
(*  CEA (Commissariat à l'énergie atomique et aux énergies alternatives)  *)
(*                                                                        *)
(**************************************************************************)

(* -------------------------------------------------------------------------- *)
(* --- UnionFind Store with explicit integer keys                         --- *)
(* -------------------------------------------------------------------------- *)

module Imap = Map.Make(Int)

module S =
struct
  type 'a rref = int
  type 'a store = {
    mutable rid : int ;
    mutable map : 'a Imap.t ;
  }

  let new_store () = { rid = 0 ; map = Imap.empty  }
  let copy r = { rid = r.rid ; map = r.map }

  let make s v =
    let k = succ s.rid in
    s.rid <- k ; s.map <- Imap.add k v s.map ; k

  let get s k = Imap.find k s.map
  let set s k v = s.map <- Imap.add k v s.map

  let eq _s i j = (i == j)

end

module Ufind :
sig
  type 'a rref
  type 'a store
  val id : 'a rref -> int
  val new_store : unit -> 'a store
  val get : 'a store -> 'a rref -> 'a
  val set : 'a store -> 'a rref -> 'a -> unit
  val make : 'a store -> 'a -> 'a rref
  val find : 'a store -> 'a rref -> 'a rref
  val merge : 'a store -> ('a -> 'a -> 'a) -> 'a rref -> 'a rref -> 'a rref
end =
struct
  include UnionFind.Make(S)
  let id = Fun.id
end

module type NodeData =
sig
  type 'a t
  val get_id : 'a t -> int
  val set_id : 'a t -> int -> unit
end


module Make(D : NodeData) =
struct

  type store = {
    values : data Ufind.store ;
    keymap : (int,node) Hashtbl.t ;
  }
  and data = node D.t
  and node = { rref : data Ufind.rref ; store : store }

  let check ~fn (m : store) (m' : store) =
    if m == m' then m else
      invalid_arg
        (Printf.sprintf "Region.Store.%s (inconsistent maps)" fn)
      [@ coverage off]

  let checklock ~fn store =
    if Hashtbl.length store.keymap > 0 then
      invalid_arg
        (Printf.sprintf "Region.Store.%s (locked map)" fn)
      [@ coverage off]

  let create () = { values = Ufind.new_store () ; keymap = Hashtbl.create 0 }

  let ufind n = Ufind.find n.store.values n.rref
  let find n = { n with rref = ufind n }

  let by_rank a b = Int.compare (Ufind.id a.rref) (Ufind.id b.rref)
  let find_all ns = List.sort_uniq by_rank @@ List.map find ns

  let find_all2 xs ys =
    let rec bag xs ys =
      match xs , ys with
      | [] , w | w , [] -> List.map find w
      | x::xs , y::ys ->
        if x.rref == y.rref then
          find x :: bag xs ys
        else
          find x :: find y :: bag xs ys
    in List.sort_uniq by_rank (bag xs ys)

  let store a = a.store
  let get a = Ufind.get a.store.values a.rref
  let set a v = Ufind.set a.store.values a.rref v
  let any a b = if Ufind.id a.rref <= Ufind.id b.rref then a else b

  let fresh store v =
    checklock ~fn:"fresh" store ;
    { rref = Ufind.make store.values v ; store }

  let eq a b =
    let store = check ~fn:"eq" a.store b.store in
    S.eq store a.rref b.rref

  let merge f a b =
    let store = check ~fn:"merge" a.store b.store in
    checklock ~fn:"merge" store ;
    { rref = Ufind.merge store.values f a.rref b.rref ; store }

  let noid = (-1)
  let is_locked store = Hashtbl.length store.keymap > 0

  let lock a : bool =
    let d = get a in
    0 < D.get_id d ||
    begin
      let uid = succ @@ Hashtbl.length a.store.keymap in
      Hashtbl.add a.store.keymap uid a ;
      D.set_id d uid ; false
    end

  let id a = checklock ~fn:"id" a.store ; D.get_id @@ get a
  let of_id store k =
    try Hashtbl.find store.keymap k
    with Not_found -> invalid_arg "Region.Store.of_id" [@ coverage off]

  let pretty fmt a =
    if is_locked a.store then
      Format.fprintf fmt "R%04x" (D.get_id @@ get a)
    else
      Format.fprintf fmt "#%04X" (Ufind.id a.rref)

  type marks = Z.t ref
  let marks () = ref Z.zero

  let marked m n =
    let uid = D.get_id @@ get n in
    Z.testbit !m uid

  let test_and_mark m n =
    let uid = D.get_id @@ get n in
    Z.testbit !m uid ||
    ( m := Z.(!m lor (one lsl uid)) ; false )

  let once f =
    let m = ref Z.zero in
    fun n ->
      let uid = D.get_id @@ get n in
      Z.testbit !m uid ||
      ( m := Z.(!m lor (one lsl uid)) ; f n ; false )

end