Skip to content

Commit 62782ba

Browse files
lyrmlyrm
committed
Priority queue.
1 parent 88cfb69 commit 62782ba

File tree

8 files changed

+758
-33
lines changed

8 files changed

+758
-33
lines changed

‎src/priority_queue.ml

Lines changed: 350 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,350 @@
1+
(* Copyright (c) 2024 Carine Morel
2+
3+
Permission to use, copy, modify, and/or distribute this software for any
4+
purpose with or without fee is hereby granted, provided that the above
5+
copyright notice and this permission notice appear in all copies.
6+
7+
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
8+
REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
9+
AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
10+
INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
11+
LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
12+
OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
13+
PERFORMANCE OF THIS SOFTWARE. *)
14+
15+
(* Based on the skiplist implementation written by Vesa Karvonen. *)
16+
17+
module Atomic = Multicore_magic.Transparent_atomic
18+
19+
type ('k, 'v, _) node =
20+
| Null : ('k, 'v, [> `Null ]) node
21+
| Node : {
22+
key : 'k;
23+
value : 'v;
24+
next : ('k, 'v) links;
25+
mutable incr : Size.once;
26+
}
27+
-> ('k, 'v, [> `Node ]) node
28+
| Mark : {
29+
node : ('k, 'v, [< `Null | `Node ]) node;
30+
decr : Size.once;
31+
}
32+
-> ('k, 'v, [> `Mark ]) node
33+
34+
and ('k, 'v) link =
35+
| Link : ('k, 'v, [< `Null | `Node | `Mark ]) node -> ('k, 'v) link
36+
[@@unboxed]
37+
38+
and ('k, 'v) links = ('k, 'v) link Atomic.t array
39+
40+
type 'k compare = 'k -> 'k -> int
41+
(* Encoding the [compare] function using an algebraic type would allow the
42+
overhead of calling a closure to be avoided for selected primitive types like
43+
[int]. *)
44+
45+
type ('k, 'v) t = { compare : 'k compare; root : ('k, 'v) links; size : Size.t }
46+
47+
(* *)
48+
49+
(** [get_random_height max_height] gives a random value [n] in the range from
50+
[1] to [max_height] with the desired distribution such that [n] is twice as
51+
likely as [n + 1]. *)
52+
let rec get_random_height max_height =
53+
let m = (1 lsl max_height) - 1 in
54+
let x = Random.bits () land m in
55+
if x = 1 then
56+
(* We reject [1] to get the desired distribution. *)
57+
get_random_height max_height
58+
else
59+
(* We do a binary search for the highest 1 bit. Techniques in
60+
61+
Using de Bruijn Sequences to Index a 1 in a Computer Word
62+
by Leiserson, Prokop, and Randall
63+
64+
could perhaps speed this up a bit, but this is likely not performance
65+
critical. *)
66+
let n = 0 in
67+
let n, x = if 0xFFFF < x then (n + 0x10, x lsr 0x10) else (n, x) in
68+
let n, x = if 0x00FF < x then (n + 0x08, x lsr 0x08) else (n, x) in
69+
let n, x = if 0x000F < x then (n + 0x04, x lsr 0x04) else (n, x) in
70+
let n, x = if 0x0003 < x then (n + 0x02, x lsr 0x02) else (n, x) in
71+
let n, _ = if 0x0001 < x then (n + 0x01, x lsr 0x01) else (n, x) in
72+
max_height - n
73+
74+
(* *)
75+
76+
let[@inline] is_marked = function
77+
| Link (Mark _) -> true
78+
| Link (Null | Node _) -> false
79+
80+
(* *)
81+
82+
(** [find_path t key preds succs lowest] search fo the position after all the
83+
nodes with key [key], updating [preds] and [succs] and removing nodes with
84+
marked references along the way, and always descending down to [lowest]
85+
level. The boolean return value is only meaningful when [lowest] is given as
86+
[0]. *)
87+
let rec find_path t key preds succs lowest =
88+
let prev = t.root in
89+
let level = Array.length prev - 1 in
90+
let prev_at_level = Array.unsafe_get prev level in
91+
find_path_rec t key prev prev_at_level preds succs level lowest
92+
(Atomic.get prev_at_level)
93+
94+
and find_path_rec t key prev prev_at_level preds succs level lowest = function
95+
| Link Null ->
96+
if level < Array.length preds then begin
97+
Array.unsafe_set preds level prev_at_level;
98+
Array.unsafe_set succs level Null
99+
end;
100+
if lowest < level then
101+
let level = level - 1 in
102+
let prev_at_level = Array.unsafe_get prev level in
103+
find_path_rec t key prev prev_at_level preds succs level lowest
104+
(Atomic.get prev_at_level)
105+
| Link (Node r as curr) -> begin
106+
let next_at_level = Array.unsafe_get r.next level in
107+
match Atomic.get next_at_level with
108+
| Link (Null | Node _) as next ->
109+
let c = t.compare key r.key in
110+
111+
if 0 <= c then
112+
(* key >= r.key *)
113+
find_path_rec t key r.next next_at_level preds succs level lowest
114+
next
115+
else begin
116+
if level < Array.length preds then begin
117+
Array.unsafe_set preds level (Array.unsafe_get prev level);
118+
Array.unsafe_set succs level curr
119+
end;
120+
121+
if lowest < level then
122+
let level = level - 1 in
123+
let prev_at_level = Array.unsafe_get prev level in
124+
find_path_rec t key prev prev_at_level preds succs level lowest
125+
(Atomic.get prev_at_level)
126+
else begin
127+
if level = 0 then begin
128+
if r.incr != Size.used_once then begin
129+
Size.update_once t.size r.incr;
130+
r.incr <- Size.used_once
131+
end;
132+
() (* Return *)
133+
end
134+
end
135+
end
136+
| Link (Mark r) ->
137+
(* The [curr_node] is being removed from the skiplist and we help with
138+
that. *)
139+
if level = 0 then Size.update_once t.size r.decr;
140+
find_path_rec t key prev prev_at_level preds succs level lowest
141+
(let after = Link r.node in
142+
if Atomic.compare_and_set prev_at_level (Link curr) after then
143+
after
144+
else Atomic.get prev_at_level)
145+
end
146+
| Link (Mark _) ->
147+
(* The node corresponding to [prev] is being removed from the skiplist.
148+
This means we might no longer have an up-to-date view of the skiplist
149+
and so we must restart the search. *)
150+
find_path t key preds succs lowest
151+
152+
(* *)
153+
154+
(** [find_node t key] tries to find the first node with the specified [key],
155+
removing nodes with marked references along the way, and stopping as soon as
156+
the node is found. *)
157+
let rec find_node t ?timestamp key =
158+
let prev = t.root in
159+
let level = Array.length prev - 1 in
160+
let prev_at_level = Array.unsafe_get prev level in
161+
find_node_rec t ~timestamp key prev prev_at_level level
162+
(Atomic.get prev_at_level)
163+
164+
and find_node_rec t ~timestamp key prev prev_at_level level :
165+
_ -> (_, _, [< `Null | `Node ]) node = function
166+
| Link Null ->
167+
if 0 < level then
168+
let level = level - 1 in
169+
let prev_at_level = Array.unsafe_get prev level in
170+
find_node_rec t ~timestamp key prev prev_at_level level
171+
(Atomic.get prev_at_level)
172+
else Null
173+
| Link (Node r as curr) -> begin
174+
let next_at_level = Array.unsafe_get r.next level in
175+
match Atomic.get next_at_level with
176+
| Link (Null | Node _) as next ->
177+
let c = t.compare key r.key in
178+
if 0 < c then
179+
find_node_rec t ~timestamp key r.next next_at_level level next
180+
else if 0 < level then
181+
let level = level - 1 in
182+
let prev_at_level = Array.unsafe_get prev level in
183+
find_node_rec t ~timestamp key prev prev_at_level level
184+
(Atomic.get prev_at_level)
185+
else if c = 0 then begin
186+
if r.incr != Size.used_once then begin
187+
Size.update_once t.size r.incr;
188+
r.incr <- Size.used_once
189+
end;
190+
curr
191+
end
192+
else Null
193+
| Link (Mark r) ->
194+
if level = 0 then Size.update_once t.size r.decr;
195+
find_node_rec t ~timestamp key prev prev_at_level level
196+
(let after = Link r.node in
197+
if Atomic.compare_and_set prev_at_level (Link curr) after then
198+
after
199+
else Atomic.get prev_at_level)
200+
end
201+
| Link (Mark _) -> find_node t key
202+
203+
(* *)
204+
205+
let create ?(max_height = 10) ~compare () =
206+
(* The upper limit of [30] comes from [Random.bits ()] as well as from
207+
limitations with 32-bit implementations. It should not be a problem in
208+
practice. *)
209+
if max_height < 1 || 30 < max_height then
210+
invalid_arg "Skiplist: max_height must be in the range [1, 30]";
211+
let root = Array.init max_height @@ fun _ -> Atomic.make (Link Null) in
212+
let size = Size.create () in
213+
{ compare; root; size }
214+
215+
let max_height_of t = Array.length t.root
216+
217+
(* *)
218+
219+
let rec add t key value preds succs =
220+
find_path t key preds succs 0;
221+
222+
let (Node r as node : (_, _, [ `Node ]) node) =
223+
let next = Array.map (fun succ -> Atomic.make (Link succ)) succs in
224+
let incr = Size.new_once t.size Size.incr in
225+
Node { key; value; incr; next }
226+
in
227+
if
228+
let succ = Link (Array.unsafe_get succs 0) in
229+
Atomic.compare_and_set (Array.unsafe_get preds 0) succ (Link node)
230+
then begin
231+
if r.incr != Size.used_once then begin
232+
Size.update_once t.size r.incr;
233+
r.incr <- Size.used_once
234+
end;
235+
236+
(* The node is now considered as added to the skiplist. *)
237+
let rec update_levels level : unit =
238+
if Array.length r.next = level then begin
239+
if is_marked (Atomic.get (Array.unsafe_get r.next (level - 1))) then begin
240+
(* The node we finished adding has been removed concurrently. To
241+
ensure that no references we added to the node remain, we call
242+
[find_node] which will remove nodes with marked references along
243+
the way. *)
244+
find_node t key |> ignore
245+
end
246+
end
247+
else if
248+
let succ = Link (Array.unsafe_get succs level) in
249+
Atomic.compare_and_set (Array.unsafe_get preds level) succ (Link node)
250+
then update_levels (level + 1)
251+
else
252+
let _found = find_path t key preds succs level in
253+
let rec update_nexts level' =
254+
if level' < level then update_levels level
255+
else
256+
let next = Array.unsafe_get r.next level' in
257+
match Atomic.get next with
258+
| Link (Null | Node _) as before ->
259+
let succ = Link (Array.unsafe_get succs level') in
260+
if before != succ then
261+
(* It is possible for a concurrent remove operation to have
262+
marked the link. *)
263+
if Atomic.compare_and_set next before succ then
264+
update_nexts (level' - 1)
265+
else update_levels level
266+
else update_nexts (level' - 1)
267+
| Link (Mark _) ->
268+
(* The node we were trying to add has been removed concurrently.
269+
To ensure that no references we added to the node remain, we
270+
call [find_node] which will remove nodes with marked
271+
references along the way. *)
272+
find_path t key preds succs level
273+
in
274+
update_nexts (Array.length r.next - 1)
275+
in
276+
update_levels 1
277+
end
278+
else add t key value preds succs
279+
280+
let add t key value : unit =
281+
let height = get_random_height (Array.length t.root) in
282+
let preds =
283+
(* Init with [Obj.magic ()] is safe as the array is fully overwritten by
284+
[find_path] called at the start of the recursive [try_add]. *)
285+
Array.make height (Obj.magic ())
286+
in
287+
let succs = Array.make height Null in
288+
add t key value preds succs
289+
290+
(* *)
291+
292+
let length t = Size.get t.size
293+
294+
(* *)
295+
296+
let rec find_min t : (_, _, [< `Node | `Null ]) node =
297+
let root = t.root in
298+
let root_at_level0 = Array.unsafe_get root 0 in
299+
match Atomic.get root_at_level0 with
300+
| Link (Mark _) -> assert false
301+
| Link Null -> Null
302+
| Link (Node r) as curr_link -> (
303+
let next_at_level_0 = Array.unsafe_get r.next 0 in
304+
match Atomic.get next_at_level_0 with
305+
| Link (Null | Node _) -> Node r
306+
| Link (Mark next_marked) ->
307+
Size.update_once t.size next_marked.decr;
308+
if
309+
Atomic.compare_and_set root_at_level0 curr_link
310+
(Link next_marked.node)
311+
then find_min t
312+
else find_min t)
313+
314+
let rec try_remove t key next level link = function
315+
| Link (Mark r) ->
316+
if level = 0 then begin
317+
Size.update_once t.size r.decr;
318+
false
319+
end
320+
else
321+
let level = level - 1 in
322+
let link = Array.unsafe_get next level in
323+
try_remove t key next level link (Atomic.get link)
324+
| Link ((Null | Node _) as succ) ->
325+
let decr =
326+
if level = 0 then Size.new_once t.size Size.decr else Size.used_once
327+
in
328+
let marked_succ = Mark { node = succ; decr } in
329+
if Atomic.compare_and_set link (Link succ) (Link marked_succ) then
330+
if level = 0 then
331+
let _node = find_node t key in
332+
true
333+
else
334+
let level = level - 1 in
335+
let link = Array.unsafe_get next level in
336+
try_remove t key next level link (Atomic.get link)
337+
else try_remove t key next level link (Atomic.get link)
338+
339+
let remove_min_opt t =
340+
let rec loop t =
341+
match find_min t with
342+
| Null -> None
343+
| Node { next; key; value; _ } ->
344+
let level = Array.length next - 1 in
345+
let link = Array.unsafe_get next level in
346+
if try_remove t key next level link (Atomic.get link) then
347+
Some (key, value)
348+
else loop t
349+
in
350+
loop t

0 commit comments

Comments
 (0)