leaner Heap

heap.ml

@@ -25,20 +25,22 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 (** {1 Imperative priority queue} *)
 
-module Tree = struct
-
-  type 'a t = 'a tree * ('a -> 'a -> int)
-    (** A splay tree with the given comparison function *)
+type 'a t = {
+  mutable tree : 'a tree;
+  cmp : 'a -> 'a -> int;
+} (** A splay tree heap with the given comparison function *)
 and 'a tree =
   | Empty
   | Node of ('a tree * 'a * 'a tree)
   (** A splay tree containing values of type 'a *)
 
-  let empty ~cmp =
-    (Empty, cmp)
+let empty ~cmp = {
+  tree = Empty;
+  cmp;
+}
 
-  let is_empty (tree, _) =
-    match tree with
+let is_empty h =
+  match h.tree with
   | Empty -> true
   | Node _ -> false
 
@@ -73,28 +75,22 @@ module Tree = struct
       end
 
 (** Insert the element in the tree *)
-  let insert (tree, cmp) x =
-    let small, big = partition ~cmp x tree in
+let insert h x =
+  let small, big = partition ~cmp:h.cmp x h.tree in
   let tree' = Node (small, x, big) in
-    tree', cmp
-
-  (** Returns the top value, or raise Not_found is empty *)
-  let top (tree, _) =
-    match tree with
-    | Empty -> raise Not_found
-    | Node (_, x, _) -> x
+  h.tree <- tree'
 
 (** Access minimum value *)
-  let min (tree, _) =
+let min h =
   let rec min tree =
     match tree with
     | Empty -> raise Not_found
     | Node (Empty, x, _) -> x
     | Node (l, _, _) -> min l
-    in min tree
+  in min h.tree
 
 (** Get minimum value and remove it from the tree *)
-  let delete_min (tree, cmp) =
+let pop h =
   let rec delete_min tree = match tree with
   | Empty -> raise Not_found
   | Node (Empty, x, b) -> x, b
@@ -104,61 +100,29 @@ module Tree = struct
     let m, a' = delete_min a in
     m, Node (a', x, Node (b, y, c))
   in
-    let m, tree' = delete_min tree in
-    m, (tree', cmp)
+  let m, tree' = delete_min h.tree in
+  h.tree <- tree';
+  m
+
+let junk h =
+  ignore (pop h)
 
 (** Iterate on elements *)
-  let iter (tree, _) f =
+let iter h f =
   let rec iter tree =
     match tree with
     | Empty -> ()
     | Node (a, x, b) ->
       iter a; f x; iter b
-    in iter tree
-end
+  in iter h.tree
 
-type 'a t = 'a Tree.t ref
-  (** The heap is a reference to a splay tree *)
-
-(** Create an empty heap *)
-let empty ~cmp =
-  ref (Tree.empty ~cmp)
-
-(** Insert a value in the heap *)
-let insert heap x =
-  heap := Tree.insert !heap x
-
-(** Check whether the heap is empty *)
-let is_empty heap =
-  Tree.is_empty !heap
-
-(** Access the minimal value of the heap, or raises Empty *)
-let min (heap : 'a t) : 'a =
-  let elt = Tree.min !heap in
-  elt
-
-(** Discard the minimal element *)
-let junk heap =
-  let _, tree' = Tree.delete_min !heap in
-  heap := tree'
-
-(** Remove and return the mininal value (or raise Invalid_argument) *)
-let pop heap =
-  let elt, tree' = Tree.delete_min !heap in
-  heap := tree';
-  elt
-
-(** Iterate on the elements, in an unspecified order *)
-let iter heap k =
-  Tree.iter !heap (fun elt -> k elt)
-
-let size heap =
+let size h =
   let r = ref 0 in
-  iter heap (fun _ -> incr r);
+  iter h (fun _ -> incr r);
   !r
 
-let to_seq heap =
-  fun k -> iter heap k
+let to_seq h =
+  fun k -> iter h k
 
-let of_seq heap seq =
-  seq (fun elt -> insert heap elt)
+let of_seq h seq =
+  seq (fun elt -> insert h elt)