A simple Priority Queue in OCaml

While working on some programming quiz in OCaml, I needed to use a priority queue (or heap) in order to solve some graph traversal problem. However, unfortunately, OCaml doesn’t provide any sort of priority queue or heap in its standard library. And after a brief search I haven’t found any that satisfied me in terms of their interface.

However, in I came across the idea to simply use a Set keyed by (int * 'a) where int is the priority and 'a the payload. That’s clever! The access complexity for the min (or max) element is now log N but that is still better than not having a proper heap. I also didn’t feel like writing my own heap based on an array or a simple tree structure; making use of the well-known Base.Set module seemed more attractive. I generally use Jane Street’s Base and Core all the time, as they are a bit more pleasant to use and more complete than the standard library.

It’s reasonably easy to implement such a queue for just one type, but coming up with a generic approach that works for more than one data type was a bit more difficult; it involved wrangling with first-class modules and comparators for a bit. Being a bit rusty with those, it took me the better part of a morning to arrive at a satisfactory solution. I’ve done some research afterwards, and the interface is actually fairly similar to the bheap package, which I could have used in the first place.

Also, the OCaml user manual (ironically, after searching and not finding a library implementation) uses a PrioQueue as example to demonstrate basic modules. A basic heap implementation is shown, demonstrating how easy it is implemented in elegant OCaml. However, it’s only implemented for int, whereas the implementation below works for any type with a Comparator.S implementation. (not that it would be very difficult to adapt the manual’s example to work with any type).

The approach I took is generic over any datatype that comes in a module providing compare and sexp_of_t and is slightly less elegant than Base.Set (with comparator witness etc.) but more easily understandable to the more fresh-faced OCaml programmers like me. It is essentially the structure used by the Stdlib container types, which at this point I had already forgotten and was thus forced to rediscover:

(* Simple approach: use a functor! Comparator.S is prevalent through existing data types. *)
module PrioQueue (Inner : Comparator.S) = struct
  (* This is the new element type, wrapping both a priority and the payload.
     Most importantly, it provides a compare function, and also `sexp_of_t` as required
     by Base.Set.
  *)
  module Elt = struct
    (* just a tuple is enough *)
    type t = int * Inner.t

    (* First compare priorities, then fall back to payload comparison. *)
    let compare (a, x) (b, y) =
      match Int.compare a b with 0 -> Inner.comparator.compare x y | x -> x

    let sexp_of_t (a, x) =
      Sexp.List [ Int.sexp_of_t a; Inner.comparator.sexp_of_t x ]
  end

  (* Comparator for Set: The Comparator complication wants to take the existing
     element module `Elt` and add a `comparator` definition based on its
     `compare` and `sexp_of_t` functions.
    *)
  module EltComp = struct
    include Comparator.Make (Elt)

    type t = Elt.t
  end

  (* Now we can write down both element type and comparator witness type. *)
  type t = Set.M(EltComp).t

  (* The priority queue (heap) functionality now basically comes for free. *)
  let empty : t = Set.empty (module EltComp)
  let is_empty = Set.is_empty
  let add q ~prio elt = Set.add q (prio, elt)
  let min_elt_exn q = snd (Set.min_elt_exn q)

  let remove_min_elt_exn q =
    let ((_, elt) as min) = Set.min_elt_exn q in
    (Set.remove q min, elt)

  let sexp_of_t q =
    let elts = Set.to_list q in
    Sexp.List (List.map elts ~f:Elt.sexp_of_t)

  (* add more of your liking ... *)
end

This can already be used with Base data types:

(* Declaration necessary for every data type to be contained. *)
module Int_pq = PrioQueue (Int)

let () =
  let q = Int_pq.empty in
  let q = Int_pq.add q ~prio:1 1 in
  let q = Int_pq.add q ~prio:2 2 in
  let q = Int_pq.add q ~prio:0 3 in
  assert (Int_pq.min_elt_exn q = 3)

However, we run into (slight) trouble when declaring our own module. We’d need to use the annoying indirection described in Real World OCaml where a module includes the inner module T for the entire functionality, and then uses Comparator.Make to provide the comparison functions.

This can be helped by a small wrapper functor:

module type Comparable = sig
  type t [@@deriving sexp, compare]
end

module PrimPrioQueue (Inner : Comparable) = struct
  module CompPos = struct
    include Inner
    include Comparable.Make (Inner)
  end

  include PrioQueue (CompPos)
end

A “primitive” priority queue PrimPrioQueue can now be created from any module containing as little as compare and sexp_of_t, where PrimPrioQueue is now responsible for building the comparator into the module. The inner type is still transparent.

Using this approach worked with no more nasty surprises, probably a bit slower than a true heap, but robust enough to not have to think about it anymore after the initial setup.