/
ptmap.ml
350 lines (299 loc) · 10.5 KB
/
ptmap.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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
(**************************************************************************)
(* *)
(* Copyright (C) Jean-Christophe Filliatre *)
(* *)
(* This software is free software; you can redistribute it and/or *)
(* modify it under the terms of the GNU Lesser General Public *)
(* License version 2.1, with the special exception on linking *)
(* described in file LICENSE. *)
(* *)
(* This software is distributed in the hope that it will be useful, *)
(* but WITHOUT ANY WARRANTY; without even the implied warranty of *)
(* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. *)
(* *)
(**************************************************************************)
(*s Maps of integers implemented as Patricia trees, following Chris
Okasaki and Andrew Gill's paper {\em Fast Mergeable Integer Maps}
({\tt\small http://www.cs.columbia.edu/\~{}cdo/papers.html\#ml98maps}).
See the documentation of module [Ptset] which is also based on the
same data-structure. *)
type key = int
type 'a t =
| Empty
| Leaf of int * 'a
| Branch of int * int * 'a t * 'a t
let empty = Empty
let is_empty t = t = Empty
let zero_bit k m = (k land m) == 0
let rec mem k = function
| Empty -> false
| Leaf (j,_) -> k == j
| Branch (_, m, l, r) -> mem k (if zero_bit k m then l else r)
let rec find k = function
| Empty -> raise Not_found
| Leaf (j,x) -> if k == j then x else raise Not_found
| Branch (_, m, l, r) -> find k (if zero_bit k m then l else r)
let find_opt k m = try Some (find k m) with Not_found -> None
(* Note: find_first/last have to look in both subtrees
as these are little-endian Patricia trees *)
let rec find_first_opt f = function
| Empty -> None
| Leaf (j,x) -> if f j then Some (j,x) else None
| Branch (_, _, l, r) ->
match find_first_opt f l, find_first_opt f r with
| Some (lk,lv) , Some (rk,rv) ->
if lk < rk then Some (lk,lv) else Some (rk,rv)
| Some v, None | None, Some v -> Some v
| None, None -> None
let find_first f = function
| Empty -> raise Not_found
| Leaf (j,x) -> if f j then (j,x) else raise Not_found
| Branch (_, _, l, r) ->
match find_first_opt f l, find_first_opt f r with
| Some (lk,lv) , Some (rk,rv) -> if lk < rk then (lk,lv) else (rk,rv)
| Some v, None | None, Some v -> v
| None, None -> raise Not_found
let rec find_last_opt f = function
| Empty -> None
| Leaf (j,x) -> if f j then Some (j,x) else None
| Branch (_, _, l, r) ->
match find_last_opt f l, find_last_opt f r with
| Some (lk,lv) , Some (rk,rv) ->
if lk > rk then Some (lk,lv) else Some (rk,rv)
| Some v, None | None, Some v -> Some v
| None, None -> None
let find_last f = function
| Empty -> raise Not_found
| Leaf (j,x) -> if f j then (j,x) else raise Not_found
| Branch (_, _, l, r) ->
match find_last_opt f l, find_last_opt f r with
| Some (lk,lv) , Some (rk,rv) -> if lk > rk then (lk,lv) else (rk,rv)
| Some v, None | None, Some v -> v
| None, None -> raise Not_found
let lowest_bit x = x land (-x)
let branching_bit p0 p1 = lowest_bit (p0 lxor p1)
let mask p m = p land (m-1)
let join (p0,t0,p1,t1) =
let m = branching_bit p0 p1 in
if zero_bit p0 m then
Branch (mask p0 m, m, t0, t1)
else
Branch (mask p0 m, m, t1, t0)
let match_prefix k p m = (mask k m) == p
let add k x t =
let rec ins = function
| Empty -> Leaf (k,x)
| Leaf (j,_) as t ->
if j == k then Leaf (k,x) else join (k, Leaf (k,x), j, t)
| Branch (p,m,t0,t1) as t ->
if match_prefix k p m then
if zero_bit k m then
Branch (p, m, ins t0, t1)
else
Branch (p, m, t0, ins t1)
else
join (k, Leaf (k,x), p, t)
in
ins t
let singleton k v =
add k v empty
let branch = function
| (_,_,Empty,t) -> t
| (_,_,t,Empty) -> t
| (p,m,t0,t1) -> Branch (p,m,t0,t1)
let remove k t =
let rec rmv = function
| Empty -> Empty
| Leaf (j,_) as t -> if k == j then Empty else t
| Branch (p,m,t0,t1) as t ->
if match_prefix k p m then
if zero_bit k m then
branch (p, m, rmv t0, t1)
else
branch (p, m, t0, rmv t1)
else
t
in
rmv t
let rec cardinal = function
| Empty -> 0
| Leaf _ -> 1
| Branch (_,_,t0,t1) -> cardinal t0 + cardinal t1
let rec iter f = function
| Empty -> ()
| Leaf (k,x) -> f k x
| Branch (_,_,t0,t1) -> iter f t0; iter f t1
let rec map f = function
| Empty -> Empty
| Leaf (k,x) -> Leaf (k, f x)
| Branch (p,m,t0,t1) -> Branch (p, m, map f t0, map f t1)
let rec mapi f = function
| Empty -> Empty
| Leaf (k,x) -> Leaf (k, f k x)
| Branch (p,m,t0,t1) -> Branch (p, m, mapi f t0, mapi f t1)
let rec fold f s accu = match s with
| Empty -> accu
| Leaf (k,x) -> f k x accu
| Branch (_,_,t0,t1) -> fold f t0 (fold f t1 accu)
let rec for_all p = function
| Empty -> true
| Leaf (k, v) -> p k v
| Branch (_,_,t0,t1) -> for_all p t0 && for_all p t1
let rec exists p = function
| Empty -> false
| Leaf (k, v) -> p k v
| Branch (_,_,t0,t1) -> exists p t0 || exists p t1
let rec filter pr = function
| Empty -> Empty
| Leaf (k, v) as t -> if pr k v then t else Empty
| Branch (p,m,t0,t1) -> branch (p, m, filter pr t0, filter pr t1)
let rec filter_map pr = function
| Empty -> Empty
| Leaf (k, v) -> (match pr k v with Some v' -> Leaf (k, v') | None -> Empty)
| Branch (p,m,t0,t1) -> branch (p, m, filter_map pr t0, filter_map pr t1)
let partition p s =
let rec part (t,f as acc) = function
| Empty -> acc
| Leaf (k, v) -> if p k v then (add k v t, f) else (t, add k v f)
| Branch (_,_,t0,t1) -> part (part acc t0) t1
in
part (Empty, Empty) s
let rec choose = function
| Empty -> raise Not_found
| Leaf (k, v) -> (k, v)
| Branch (_, _, t0, _) -> choose t0 (* we know that [t0] is non-empty *)
let rec choose_opt = function
| Empty -> None
| Leaf (k, v) -> Some (k, v)
| Branch (_, _, t0, _) -> choose_opt t0 (* we know that [t0] is non-empty *)
let split x m =
let coll k v (l, b, r) =
if k < x then add k v l, b, r
else if k > x then l, b, add k v r
else l, Some v, r
in
fold coll m (empty, None, empty)
let rec min_binding = function
| Empty -> raise Not_found
| Leaf (k, v) -> (k, v)
| Branch (_,_,s,t) ->
let (ks, _) as bs = min_binding s in
let (kt, _) as bt = min_binding t in
if ks < kt then bs else bt
let rec min_binding_opt = function
| Empty -> None
| Leaf (k, v) -> Some (k, v)
| Branch (_,_,s,t) ->
match (min_binding_opt s, min_binding_opt t) with
| None, None -> None
| None, bt -> bt
| bs, None -> bs
| (Some (ks, _) as bs), (Some (kt, _) as bt) ->
if ks < kt then bs else bt
let rec max_binding = function
| Empty -> raise Not_found
| Leaf (k, v) -> (k, v)
| Branch (_,_,s,t) ->
let (ks, _) as bs = max_binding s in
let (kt, _) as bt = max_binding t in
if ks > kt then bs else bt
let rec max_binding_opt = function
| Empty -> None
| Leaf (k, v) -> Some (k, v)
| Branch (_,_,s,t) ->
match max_binding_opt s, max_binding_opt t with
| None, None -> None
| None, bt -> bt
| bs, None -> bs
| (Some (ks, _) as bs), (Some (kt, _) as bt) ->
if ks > kt then bs else bt
let bindings m =
fold (fun k v acc -> (k, v) :: acc) m []
(* we order constructors as Empty < Leaf < Branch *)
let compare cmp t1 t2 =
let rec compare_aux t1 t2 = match t1,t2 with
| Empty, Empty -> 0
| Empty, _ -> -1
| _, Empty -> 1
| Leaf (k1,x1), Leaf (k2,x2) ->
let c = compare k1 k2 in
if c <> 0 then c else cmp x1 x2
| Leaf _, Branch _ -> -1
| Branch _, Leaf _ -> 1
| Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
let c = compare p1 p2 in
if c <> 0 then c else
let c = compare m1 m2 in
if c <> 0 then c else
let c = compare_aux l1 l2 in
if c <> 0 then c else
compare_aux r1 r2
in
compare_aux t1 t2
let equal eq t1 t2 =
let rec equal_aux t1 t2 = match t1, t2 with
| Empty, Empty -> true
| Leaf (k1,x1), Leaf (k2,x2) -> k1 = k2 && eq x1 x2
| Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
p1 = p2 && m1 = m2 && equal_aux l1 l2 && equal_aux r1 r2
| _ -> false
in
equal_aux t1 t2
let merge f m1 m2 =
let add m k = function None -> m | Some v -> add k v m in
(* first consider all bindings in m1 *)
let m = fold
(fun k1 v1 m -> add m k1 (f k1 (Some v1) (find_opt k1 m2))) m1 empty in
(* then bindings in m2 that are not in m1 *)
fold (fun k2 v2 m -> if mem k2 m1 then m else add m k2 (f k2 None (Some v2)))
m2 m
let update x f m =
match f (find_opt x m) with
| None -> remove x m
| Some z -> add x z m
let unsigned_lt n m = n >= 0 && (m < 0 || n < m)
let rec union f = function
| Empty, t -> t
| t, Empty -> t
| Leaf (k,v1), t ->
update k (function None -> Some v1 | Some v2 -> f k v1 v2) t
| t, Leaf (k,v2) ->
update k (function None -> Some v2 | Some v1 -> f k v1 v2) t
| (Branch (p,m,s0,s1) as s), (Branch (q,n,t0,t1) as t) ->
if m == n && match_prefix q p m then
(* The trees have the same prefix. Merge the subtrees. *)
branch (p, m, union f (s0,t0), union f (s1,t1))
else if unsigned_lt m n && match_prefix q p m then
(* [q] contains [p]. Merge [t] with a subtree of [s]. *)
if zero_bit q m then
branch (p, m, union f (s0,t), s1)
else
branch (p, m, s0, union f (s1,t))
else if unsigned_lt n m && match_prefix p q n then
(* [p] contains [q]. Merge [s] with a subtree of [t]. *)
if zero_bit p n then
branch (q, n, union f (s,t0), t1)
else
branch (q, n, t0, union f (s,t1))
else
(* The prefixes disagree. *)
join (p, s, q, t)
let union f s t = union f (s,t)
let to_seq m =
let rec prepend_seq m s = match m with
| Empty -> s
| Leaf (k, v) -> fun () -> Seq.Cons((k,v), s)
| Branch (_, _, l, r) -> prepend_seq l (prepend_seq r s)
in
prepend_seq m Seq.empty
let to_seq_from k m =
let rec prepend_seq m s = match m with
| Empty -> s
| Leaf (key, v) -> if key >= k then fun () -> Seq.Cons((key,v), s) else s
| Branch (_, _, l, r) -> prepend_seq l (prepend_seq r s)
in
prepend_seq m Seq.empty
let add_seq s m =
Seq.fold_left (fun m (k, v) -> add k v m) m s
let of_seq s =
Seq.fold_left (fun m (k, v) -> add k v m) empty s