From bc562242c928fd7203a70c288dd389be713bdfa4 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sun, 11 Feb 2018 21:52:31 -0600
Subject: [PATCH] add `Heap.remove`

---
 src/util/Heap.ml      | 34 ++++++++++++++++++----------------
 src/util/Heap_intf.ml |  7 ++++++-
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/src/util/Heap.ml b/src/util/Heap.ml
index 1e309ab1..254b92a3 100644
--- a/src/util/Heap.ml
+++ b/src/util/Heap.ml
@@ -29,7 +29,7 @@ module Make(Elt : RANKED) = struct
   *)
 
   (* [elt] is above or at its expected position. Move it up the heap
-     (towards high indices) to restore the heap property *)
+     (towards the root) to restore the heap property *)
   let percolate_up {heap} (elt:Elt.t) : unit =
     let pi = ref (parent (Elt.idx elt)) in
     let i = ref (Elt.idx elt) in
@@ -42,6 +42,8 @@ module Make(Elt : RANKED) = struct
     Vec.set heap !i elt;
     Elt.set_idx elt !i
 
+  (* move [elt] towards the leaves of the heap to restore the heap
+     property *)
   let percolate_down {heap} (elt:Elt.t): unit =
     let sz = Vec.size heap in
     let li = ref (left (Elt.idx elt)) in
@@ -69,7 +71,9 @@ module Make(Elt : RANKED) = struct
 
   let[@inline] in_heap x = Elt.idx x >= 0
 
-  let[@inline] decrease s x = assert (in_heap x); percolate_up s x
+  let[@inline] decrease s x =
+    assert (in_heap x);
+    percolate_up s x
 
   (*
   let increase cmp s n =
@@ -112,20 +116,6 @@ module Make(Elt : RANKED) = struct
   let[@inline] grow_to_at_least s sz =
     Vec.grow_to_at_least s.heap sz
 
-  (*
-  let update cmp s n =
-    assert (heap_property cmp s);
-    begin
-      if in_heap s n then
-        begin
-          percolate_up cmp s (Vec.get s.indices n);
-          percolate_down cmp s (Vec.get s.indices n)
-        end
-      else insert cmp s n
-    end;
-    assert (heap_property cmp s)
-  *)
-
   let remove_min ({heap} as s) =
     if Vec.size heap=0 then raise Not_found;
     let x = Vec.get heap 0 in
@@ -140,4 +130,16 @@ module Make(Elt : RANKED) = struct
     );
     x
 
+  let remove ({heap} as s) (elt:elt) : unit =
+    assert (in_heap elt);
+    let i = Elt.idx elt in
+    Vec.fast_remove heap i;
+    Elt.set_idx elt ~-1;
+    assert (not (in_heap elt));
+    (* element put in place of [x] might be too high *)
+    if Vec.size heap > i then (
+      percolate_down s (Vec.get heap i);
+    );
+    ()
+
 end
diff --git a/src/util/Heap_intf.ml b/src/util/Heap_intf.ml
index a5d11100..c3df6013 100644
--- a/src/util/Heap_intf.ml
+++ b/src/util/Heap_intf.ml
@@ -19,7 +19,8 @@ module type S = sig
   (** Create a heap *)
 
   val decrease : t -> elt -> unit
-  (** [decrease h x] decreases the value associated to [x] within [h] *)
+  (** [decrease h x] decreases the value associated to [x] within [h],
+      making it closer to the root (so, more prioritary) *)
 
   val in_heap : elt -> bool
 
@@ -46,6 +47,10 @@ module type S = sig
   (** Remove and return the integer that has the lowest value from the heap
       @raise Not_found if the heap is empty *)
 
+  val remove : t -> elt -> unit
+  (** Remove element from the heap.
+      precond: [in_heap elt] *)
+
   val filter : t -> (elt -> bool) -> unit
   (** Filter out values that don't satisfy the predicate *)
 end