perf(solver): use VecI32 in the heap

this reduces the size of the heap (4 bits per element, not 8), and
reduces GC work by not scanning the bigarray

src/sat/Heap.ml

@@ -9,14 +9,14 @@ module Make(Elt : RANKED) = struct
 
   type t = {
     store : elt_store;
-    heap : elt Vec.t;
+    heap : VecI32.t; (* vec of elements *)
   }
 
   let _absent_index = -1
 
   let create store : t =
     { store;
-      heap = Vec.create(); }
+      heap = VecI32.create(); }
 
   let[@inline] left i = (i lsl 1) + 1 (* i*2 + 1 *)
   let[@inline] right i = (i + 1) lsl 1 (* (i+1)*2 *)
@@ -24,29 +24,32 @@ module Make(Elt : RANKED) = struct
 
   (*
   let rec heap_property cmp ({heap=heap} as s) i =
-    i >= (Vec.size heap)  ||
+    i >= (VecI32.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
   *)
 
+  let[@inline] get_elt_ self i = Elt.of_int_unsafe (VecI32.get self.heap i)
+  let[@inline] set_elt_ self i elt = VecI32.set self.heap i (elt:Elt.t:>int)
+
   (* [elt] is above or at its expected position. Move it up the heap
      (towards high indices) to restore the heap property *)
   let percolate_up (self:t) (elt:Elt.t) : unit =
     let pi = ref (parent (Elt.heap_idx self.store elt)) in
     let i = ref (Elt.heap_idx self.store elt) in
-    while !i <> 0 && Elt.cmp self.store elt (Vec.get self.heap !pi) do
-      Vec.set self.heap !i (Vec.get self.heap !pi);
-      Elt.set_heap_idx self.store (Vec.get self.heap !i) !i;
+    while !i <> 0 && Elt.cmp self.store elt (get_elt_ self !pi) do
+      set_elt_ self !i (get_elt_ self !pi);
+      Elt.set_heap_idx self.store (get_elt_ self !i) !i;
       i  := !pi;
       pi := parent !i
     done;
-    Vec.set self.heap !i elt;
+    set_elt_ self !i elt;
     Elt.set_heap_idx self.store elt !i
 
   let percolate_down (self:t) (elt:Elt.t): unit =
-    let sz = Vec.size self.heap in
+    let sz = VecI32.size self.heap in
     let li = ref (left (Elt.heap_idx self.store elt)) in
     let ri = ref (right (Elt.heap_idx self.store elt)) in
     let i = ref (Elt.heap_idx self.store elt) in
@@ -54,20 +57,21 @@ module Make(Elt : RANKED) = struct
       try
         while !li < sz do
           let child =
-            if !ri < sz && Elt.cmp self.store (Vec.get self.heap !ri) (Vec.get self.heap !li)
+            if !ri < sz &&
+               Elt.cmp self.store (get_elt_ self !ri) (get_elt_ self !li)
             then !ri
             else !li
           in
-          if not (Elt.cmp self.store (Vec.get self.heap child) elt) then raise_notrace Exit;
-          Vec.set self.heap !i (Vec.get self.heap child);
-          Elt.set_heap_idx self.store (Vec.get self.heap !i) !i;
+          if not (Elt.cmp self.store (get_elt_ self child) elt) then raise_notrace Exit;
+          set_elt_ self !i (get_elt_ self child);
+          Elt.set_heap_idx self.store (get_elt_ self !i) !i;
           i  := child;
           li := left !i;
           ri := right !i
         done;
       with Exit -> ()
     end;
-    Vec.set self.heap !i elt;
+    set_elt_ self !i elt;
     Elt.set_heap_idx self.store elt !i
 
   let[@inline] in_heap self x = Elt.heap_idx self.store x >= 0
@@ -81,34 +85,34 @@ module Make(Elt : RANKED) = struct
 
   let filter (self:t) filt : unit =
     let j = ref 0 in
-    let lim = Vec.size self.heap in
+    let lim = VecI32.size self.heap in
     for i = 0 to lim - 1 do
-      if filt (Vec.get self.heap i) then (
-        Vec.set self.heap !j (Vec.get self.heap i);
-        Elt.set_heap_idx self.store (Vec.get self.heap i) !j;
+      if filt (get_elt_ self i) then (
+        set_elt_ self !j (get_elt_ self i);
+        Elt.set_heap_idx self.store (get_elt_ self i) !j;
         incr j;
       ) else (
-        Elt.set_heap_idx self.store (Vec.get self.heap i) _absent_index;
+        Elt.set_heap_idx self.store (get_elt_ self i) _absent_index;
       );
     done;
-    Vec.shrink self.heap (lim - !j);
+    VecI32.shrink self.heap (lim - !j);
     for i = (lim / 2) - 1 downto 0 do
-      percolate_down self (Vec.get self.heap i)
+      percolate_down self (get_elt_ self i)
     done
 
-  let size s = Vec.size s.heap
-
-  let is_empty s = Vec.is_empty s.heap
+  let[@inline] size s = VecI32.size s.heap
+  let[@inline] is_empty s = VecI32.is_empty s.heap
 
   let clear self : unit =
-    Vec.iter (fun e -> Elt.set_heap_idx self.store e _absent_index) self.heap;
-    Vec.clear self.heap;
+    VecI32.iter self.heap
+      ~f:(fun e -> Elt.set_heap_idx self.store (Elt.of_int_unsafe e) _absent_index);
+    VecI32.clear self.heap;
     ()
 
   let insert self elt =
     if not (in_heap self elt) then (
-      Elt.set_heap_idx self.store elt (Vec.size self.heap);
-      Vec.push self.heap elt;
+      Elt.set_heap_idx self.store elt (VecI32.size self.heap);
+      VecI32.push self.heap (elt:Elt.t:>int);
       percolate_up self elt;
     )
 
@@ -127,22 +131,21 @@ module Make(Elt : RANKED) = struct
   *)
 
   let remove_min self =
-    match Vec.size self.heap with
+    match VecI32.size self.heap with
     | 0 -> raise Not_found
     | 1 ->
-      let x = Vec.pop self.heap in
+      let x = Elt.of_int_unsafe (VecI32.pop self.heap) in
       Elt.set_heap_idx self.store x _absent_index;
       x
     | _ ->
-      let x = Vec.get self.heap 0 in
-      let new_hd = Vec.pop self.heap in (* heap.last() *)
-      Vec.set self.heap 0 new_hd;
+      let x = get_elt_ self 0 in
+      let new_hd = Elt.of_int_unsafe (VecI32.pop self.heap) in (* heap.last() *)
+      set_elt_ self 0 new_hd;
       Elt.set_heap_idx self.store x _absent_index;
       Elt.set_heap_idx self.store new_hd 0;
       (* enforce heap property again *)
-      if Vec.size self.heap > 1 then (
+      if VecI32.size self.heap > 1 then (
         percolate_down self new_hd;
       );
       x
-
 end [@@inline]

src/sat/Heap_intf.ml

@@ -1,7 +1,7 @@
 
 module type RANKED = sig
   type store
-  type t
+  type t = private int
 
   val heap_idx : store -> t -> int
   (** Index in heap. return -1 if never set *)
@@ -10,6 +10,9 @@ module type RANKED = sig
   (** Update index in heap *)
 
   val cmp : store -> t -> t -> bool
+
+  val of_int_unsafe : int -> t
+  (** turn an integer back into an element *)
 end
 
 module type S = sig

src/sat/Solver.ml

@@ -377,6 +377,7 @@ module Make(Plugin : PLUGIN)
         (kind_of_clause c) c.cid (Store.Atom.debug_a self) arr debug_premise cp
   end
 
+  (* TODO: mostly, move into a functor outside that works on integers *)
   module Proof =  struct
     exception Resolution_error of string
 
@@ -675,11 +676,11 @@ module Make(Plugin : PLUGIN)
   module H = (Heap.Make [@specialise]) (struct
     type store = Store.t
     type t = var
-    open Store
     let[@inline] cmp store i j =
       Var.weight store j < Var.weight store i (* comparison by weight *)
     let heap_idx = Var.heap_idx
     let set_heap_idx = Var.set_heap_idx
+    let of_int_unsafe = Var.of_int_unsafe
   end)
 
   (* cause of "unsat", possibly conditional to local assumptions *)