From 3484efc6916bfac9fff821d80883c362671f152e Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Wed, 7 Oct 2020 11:44:21 -0400 Subject: [PATCH] feat: add mutable heap with increase/decrease to containers-data this code is adapted from msat --- src/data/CCMutHeap.ml | 146 +++++++++++++++++++++++++++++++++++++ src/data/CCMutHeap.mli | 5 ++ src/data/CCMutHeap_intf.ml | 51 +++++++++++++ 3 files changed, 202 insertions(+) create mode 100644 src/data/CCMutHeap.ml create mode 100644 src/data/CCMutHeap.mli create mode 100644 src/data/CCMutHeap_intf.ml diff --git a/src/data/CCMutHeap.ml b/src/data/CCMutHeap.ml new file mode 100644 index 00000000..a8a91506 --- /dev/null +++ b/src/data/CCMutHeap.ml @@ -0,0 +1,146 @@ +(* This code is extracted from Msat ( https://github.com/Gbury/mSAT ). + As such it is under the Apache 2 License. +*) + +module type RANKED = CCMutHeap_intf.RANKED + +module type S = CCMutHeap_intf.S + +module Make(Elt : RANKED) = struct + module Vec = CCVector + type elt = Elt.t + + type t = { + heap : elt Vec.vector; + } [@@unboxed] + + let _absent_index = -1 + + let create () = + { heap = Vec.create(); } + + let[@inline] left i = (i lsl 1) + 1 (* i*2 + 1 *) + let[@inline] right i = (i + 1) lsl 1 (* (i+1)*2 *) + let[@inline] parent i = (i - 1) asr 1 (* (i-1) / 2 *) + + (* + let rec heap_property cmp ({heap=heap} as s) i = + i >= (Vec.size heap) || + ((i = 0 || not(cmp (Vec. get heap i) (Vec.get heap (parent i)))) + && heap_property cmp s (left i) && heap_property cmp s (right i)) + + let heap_property cmp s = heap_property cmp s 1 + *) + + (* [elt] is above or at its expected position. Move it up the heap + (towards high indices) 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 + while !i <> 0 && Elt.cmp elt (Vec.get heap !pi) do + Vec.set heap !i (Vec.get heap !pi); + Elt.set_idx (Vec.get heap !i) !i; + i := !pi; + pi := parent !i + done; + Vec.set heap !i elt; + Elt.set_idx elt !i + + let percolate_down {heap} (elt:Elt.t): unit = + let sz = Vec.size heap in + let li = ref (left (Elt.idx elt)) in + let ri = ref (right (Elt.idx elt)) in + let i = ref (Elt.idx elt) in + begin + try + while !li < sz do + let child = + if !ri < sz && Elt.cmp (Vec.get heap !ri) (Vec.get heap !li) + then !ri + else !li + in + if not (Elt.cmp (Vec.get heap child) elt) then raise Exit; + Vec.set heap !i (Vec.get heap child); + Elt.set_idx (Vec.get heap !i) !i; + i := child; + li := left !i; + ri := right !i + done; + with Exit -> () + end; + Vec.set heap !i elt; + Elt.set_idx elt !i + + let[@inline] in_heap x = Elt.idx x >= 0 + + let[@inline] decrease s x = assert (in_heap x); percolate_up s x + + let[@inline] increase s x = assert (in_heap x); percolate_down s x + + let filter s filt = + let j = ref 0 in + let lim = Vec.size s.heap in + for i = 0 to lim - 1 do + if filt (Vec.get s.heap i) then ( + Vec.set s.heap !j (Vec.get s.heap i); + Elt.set_idx (Vec.get s.heap i) !j; + incr j; + ) else ( + Elt.set_idx (Vec.get s.heap i) _absent_index; + ); + done; + Vec.truncate s.heap (lim - !j); + for i = (lim / 2) - 1 downto 0 do + percolate_down s (Vec.get s.heap i) + done + + let size s = Vec.size s.heap + + let is_empty s = Vec.is_empty s.heap + + let clear {heap} = + Vec.iter (fun e -> Elt.set_idx e _absent_index) heap; + Vec.clear heap; + () + + let insert s elt = + if not (in_heap elt) then ( + Elt.set_idx elt (Vec.size s.heap); + Vec.push s.heap elt; + percolate_up s elt; + ) + + (* + 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) = + match Vec.size heap with + | 0 -> raise Not_found + | 1 -> + let x = Vec.pop_exn heap in + Elt.set_idx x _absent_index; + x + | _ -> + let x = Vec.get heap 0 in + let new_hd = Vec.pop_exn heap in (* heap.last() *) + Vec.set heap 0 new_hd; + Elt.set_idx x _absent_index; + Elt.set_idx new_hd 0; + (* enforce heap property again *) + if Vec.size heap > 1 then ( + percolate_down s new_hd; + ); + x + +end [@@inline] diff --git a/src/data/CCMutHeap.mli b/src/data/CCMutHeap.mli new file mode 100644 index 00000000..eba23161 --- /dev/null +++ b/src/data/CCMutHeap.mli @@ -0,0 +1,5 @@ +module type RANKED = CCMutHeap_intf.RANKED + +module type S = CCMutHeap_intf.S + +module Make(X : RANKED) : S with type elt = X.t diff --git a/src/data/CCMutHeap_intf.ml b/src/data/CCMutHeap_intf.ml new file mode 100644 index 00000000..152382a2 --- /dev/null +++ b/src/data/CCMutHeap_intf.ml @@ -0,0 +1,51 @@ + +(* This code is extracted from Msat ( https://github.com/Gbury/mSAT ). *) + +(** {1 Imperative Heaps} *) + +module type RANKED = sig + type t + val idx: t -> int (** Index in heap. return -1 if never set *) + val set_idx : t -> int -> unit (** Update index in heap *) + val cmp : t -> t -> bool +end + +module type S = sig + type elt + (** Type of elements *) + + type t + (** Heap of {!elt}, whose priority is increased or decreased + incrementally (see {!decrease} for instance) *) + + val create : unit -> t + (** Create a heap *) + + val decrease : t -> elt -> unit + (** [decrease h x] decreases the value associated to [x] within [h] *) + + val increase : t -> elt -> unit + (** [increase h x] increases the value associated to [x] within [h] *) + + val in_heap : elt -> bool + + val size : t -> int + (** Number of integers within the heap *) + + val is_empty : t -> bool + + val clear : t -> unit + (** Clear the content of the heap *) + + val insert : t -> elt -> unit + (** Insert a new element into the heap *) + + (*val update : (int -> int -> bool) -> t -> int -> unit*) + + val remove_min : t -> elt + (** Remove and return the integer that has the lowest value from the heap + @raise Not_found if the heap is empty *) + + val filter : t -> (elt -> bool) -> unit + (** Filter out values that don't satisfy the predicate *) +end