feat(util): Vec_float based on bigarrays

src/util/Sidekick_util.ml

@@ -2,14 +2,18 @@
 module Fmt = CCFormat
 
 module Util = Util
+
 module Vec = Vec
 module VecI32 = VecI32
+module Vec_float = Vec_float
 module Bitvec = Bitvec
-module Log = Log
+
+module IArray = IArray
 module Backtrack_stack = Backtrack_stack
 module Backtrackable_tbl = Backtrackable_tbl
+
+module Log = Log
 module Error = Error
-module IArray = IArray
 module Bag = Bag
 module Stat = Stat
 module Hash = Hash

src/util/Vec_float.ml

@@ -0,0 +1,77 @@
+
+module A = Bigarray.Array1
+
+type float_arr = (float, Bigarray.float64_elt, Bigarray.c_layout) A.t
+
+type t = {
+  mutable data: float_arr;
+  mutable sz: int;
+}
+
+let mk_arr_ sz : float_arr = A.create Bigarray.float64 Bigarray.c_layout sz
+
+let create () : t =
+  { sz=0; data=mk_arr_ 16 }
+
+let[@inline] clear self = self.sz <- 0
+let[@inline] shrink self n = if n < self.sz then self.sz <- n
+let[@inline] size self = self.sz
+let[@inline] is_empty self = self.sz = 0
+
+(* ensure capacity is [new_cap] *)
+let resize_cap_ self new_cap =
+  assert (A.dim self.data < new_cap);
+  let new_data = mk_arr_ new_cap in
+  A.blit self.data (A.sub new_data 0 (A.dim self.data));
+  self.data <- new_data
+
+let ensure_cap self (n:int) =
+  if n > A.dim self.data then (
+    let new_cap = max n (A.dim self.data * 2 + 10) in
+    resize_cap_ self new_cap;
+  )
+
+let ensure_size self n =
+  if n > self.sz then (
+    ensure_cap self n;
+    self.sz <- n
+  )
+
+let[@inline] push (self:t) x : unit =
+  ensure_cap self (self.sz+1);
+  self.data.{self.sz} <- x;
+  self.sz <- 1 + self.sz
+
+let[@inline] pop self =
+  if self.sz > 0 then (
+    let x = self.data.{self.sz-1} in
+    self.sz <- self.sz - 1;
+    x
+  ) else failwith "vec_float.pop: empty"
+
+let[@inline] get self i : float =
+  assert (i >= 0 && i < self.sz);
+  A.unsafe_get self.data i
+
+let[@inline] set self i x : unit =
+  assert (i >= 0 && i < self.sz);
+  A.unsafe_set self.data i x
+
+let[@inline] iter ~f self =
+  for i=0 to self.sz - 1 do
+    f self.data.{i}
+  done
+
+let[@inline] iteri ~f self =
+  for i=0 to self.sz - 1 do
+    f i self.data.{i}
+  done
+
+let[@inline] to_iter self k = iter ~f:k self
+
+let pp out self =
+  Format.fprintf out "[@[";
+  iteri self
+    ~f:(fun i x -> if i>0 then Format.fprintf out ",@ "; Format.fprintf out "%.3f" x);
+  Format.fprintf out "@]]"
+

src/util/Vec_float.mli

@@ -0,0 +1,33 @@
+
+(** Vectors of floats
+
+    These vectors are more optimized than {!Vec}. *)
+
+type t
+
+val create : unit -> t
+
+val ensure_size : t -> int -> unit
+
+val size : t -> int
+
+val clear : t -> unit
+
+val is_empty : t -> bool
+
+val push : t -> float -> unit
+
+val pop : t -> float
+
+val get : t -> int -> float
+
+val set : t -> int -> float -> unit
+
+val shrink : t -> int -> unit
+
+val iter : f:(float -> unit) -> t -> unit
+val iteri : f:(int -> float -> unit) -> t -> unit
+
+val to_iter : t -> float Iter.t
+
+val pp : t CCFormat.printer