Skip to content

Commit b2c260a

Browse files
committed
Add remove_min_exn function.
1 parent 62782ba commit b2c260a

File tree

2 files changed

+26
-13
lines changed

2 files changed

+26
-13
lines changed

src/priority_queue.ml

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ type 'k compare = 'k -> 'k -> int
4444

4545
type ('k, 'v) t = { compare : 'k compare; root : ('k, 'v) links; size : Size.t }
4646

47+
exception Empty
4748
(* *)
4849

4950
(** [get_random_height max_height] gives a random value [n] in the range from
@@ -293,6 +294,8 @@ let length t = Size.get t.size
293294

294295
(* *)
295296

297+
type ('a, _) poly = Option : ('a, 'a option) poly | Value : ('a, 'a) poly
298+
296299
let rec find_min t : (_, _, [< `Node | `Null ]) node =
297300
let root = t.root in
298301
let root_at_level0 = Array.unsafe_get root 0 in
@@ -336,15 +339,16 @@ let rec try_remove t key next level link = function
336339
try_remove t key next level link (Atomic.get link)
337340
else try_remove t key next level link (Atomic.get link)
338341

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
342+
let rec remove_min_as : type p v r. (p, v) t -> (p * v, r) poly -> r =
343+
fun t poly ->
344+
match find_min t with
345+
| Null -> ( match poly with Value -> raise Empty | Option -> None)
346+
| Node { next; key; value; _ } ->
347+
let level = Array.length next - 1 in
348+
let link = Array.unsafe_get next level in
349+
if try_remove t key next level link (Atomic.get link) then
350+
match poly with Value -> (key, value) | Option -> Some (key, value)
351+
else remove_min_as t poly
352+
353+
let remove_min_opt t = remove_min_as t Option
354+
let remove_min_exn t = remove_min_as t Value

src/priority_queue.mli

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,18 @@ val add : ('p, 'v) t -> 'p -> 'v -> unit
3838
val remove_min_opt : ('p, 'v) t -> ('p * 'v) option
3939
(** [remove_min_opt s] removes and returns [Some] of the element with the
4040
smallest priority from the priority queue [pq]. If the priority queue [pq]
41-
was empty [None] is returned. The elements with the same priority are
41+
was empty, [None] is returned. The elements with the same priority are
4242
removed in fifo order. *)
4343

44+
exception Empty
45+
46+
val remove_min_exn : ('p, 'v) t -> 'p * 'v
47+
(** [remove_min_exn s] removes and returns the element with the smallest
48+
priority from the priority queue [pq]. The elements with the same priority
49+
are removed in fifo order.
50+
51+
@raise Empty if the priority queue [pq] is empty.*)
52+
4453
val length : ('k, 'v) t -> int
4554
(** [length pq] computes and returns the number of elements in the priority
4655
queue [pq].

0 commit comments

Comments
 (0)