wip: CCAllocCache, an allocation cache for short-lived arrays

_oasis

@@ -77,7 +77,7 @@ Library "containers_data"
                     CCPersistentHashtbl, CCDeque, CCFQueue, CCBV, CCMixtbl,
                     CCMixmap, CCRingBuffer, CCIntMap, CCPersistentArray,
                     CCMixset, CCHashconsedSet, CCGraph, CCHashSet, CCBitField,
-                    CCHashTrie, CCBloom, CCWBTree, CCRAL
+                    CCHashTrie, CCBloom, CCWBTree, CCRAL, CCAllocCache
   BuildDepends:     bytes
   # BuildDepends:     bytes, bisect_ppx
   FindlibParent:    containers

benchs/run_benchs.ml

@@ -1032,6 +1032,72 @@ module Thread = struct
   )
 end
 
+module Alloc = struct
+  module type ALLOC_ARR = sig
+    type 'a t
+    val name : string
+    val create : int -> 'a t
+    val make : 'a t -> int -> 'a -> 'a array
+    val free : 'a t -> 'a array -> unit
+  end
+
+  let dummy =
+    let module A = struct
+      type _ t = unit
+      let name = "dummy"
+      let create _ = ()
+      let make _ i x = Array.make i x
+      let free _ _ = ()
+    end in
+    (module A : ALLOC_ARR)
+
+  let alloc_cache ~buck_size =
+    let module A = struct
+      type 'a t = 'a CCAllocCache.Arr.t
+      let name = Printf.sprintf "alloc_cache(%d)" buck_size
+      let create n = CCAllocCache.Arr.create ~buck_size n
+      let make = CCAllocCache.Arr.make
+      let free = CCAllocCache.Arr.free
+    end in
+    (module A : ALLOC_ARR)
+
+  (* repeat [n] times:
+    - repeat [batch] times:
+      - allocate [batch] arrays of size from 1 to batch+1
+    - free those arrays
+  *)
+  let bench1 ~batch n =
+    let make (module C : ALLOC_ARR) () =
+      let c = C.create (batch*2) in
+      let tmp = Array.make (batch * batch) [||] in (* temporary storage *)
+      for _ = 1 to n do
+        for j = 0 to batch-1 do
+          for k = 0 to batch-1 do
+            tmp.(j*batch + k) <- C.make c (k+1) '_';
+          done;
+        done;
+        Array.iter (C.free c) tmp (* free the whole array *)
+      done
+    in
+    B.throughputN 3 ~repeat
+      [ "dummy", make dummy, ()
+      ; "cache(5)", make (alloc_cache ~buck_size:5), ()
+      ; "cache(20)", make (alloc_cache ~buck_size:20), ()
+      ; "cache(50)", make (alloc_cache ~buck_size:50), ()
+      ]
+
+  let () = B.Tree.register (
+    "alloc" @>>>
+      [ "bench1(batch=5)" @>>
+          app_ints (bench1 ~batch:5) [100; 1_000]
+      ; "bench1(batch=15)" @>>
+          app_ints (bench1 ~batch:15) [100; 1_000]
+      ; "bench1(batch=50)" @>>
+          app_ints (bench1 ~batch:50) [100; 1_000]
+      ]
+    )
+end
+
 let () =
   try B.Tree.run_global ()
   with Arg.Help msg -> print_endline msg

doc/intro.txt

@@ -65,6 +65,7 @@ such as:
 Various data structures.
 
 {!modules:
+CCAllocCache
 CCBitField
 CCBloom
 CCBV

src/data/CCAllocCache.ml

@@ -0,0 +1,72 @@
+
+(* This file is free software, part of Logtk. See file "license" for more details. *)
+
+(** {1 Simple Cache for Allocations} *)
+
+module Arr = struct
+  type 'a t = {
+    caches: 'a array array array;
+      (* array of buckets, where each bucket is an array of arrays *)
+    max_buck_size: int;
+    sizes: int array; (* number of cached arrays in each bucket *)
+  }
+
+  let create ?(buck_size=16) n =
+    if n<1 then invalid_arg "AllocCache.Arr.create";
+    { max_buck_size=buck_size;
+      sizes=Array.make n 0;
+      caches=Array.init n (fun _ -> Array.make buck_size [||]);
+    }
+
+  let make c i x =
+    if i=0 then [||]
+    else if i<Array.length c.sizes then (
+      let bs = c.sizes.(i) in
+      if bs = 0 then Array.make i x
+      else (
+        (* remove last array *)
+        let ret = c.caches.(i).(bs-1) in
+        c.sizes.(i) <- bs - 1;
+        ret
+      )
+    ) else Array.make i x
+
+  let free c a =
+    let n = Array.length a in
+    if n > 0 && n < Array.length c.sizes then (
+      let bs = c.sizes.(n) in
+      if bs < c.max_buck_size then (
+        (* store [a] *)
+        c.caches.(n).(bs) <- a;
+        c.sizes.(n) <- bs + 1
+      )
+    )
+
+  let with_ c i x ~f =
+    let a = make c i x in
+    try
+      let ret = f a in
+      free c a;
+      ret
+    with e ->
+      free c a;
+      raise e
+end
+
+(*$inject
+  let c = Arr.create ~buck_size:2 20
+
+*)
+
+(*$Q
+  Q.small_int (fun n -> Array.length (Arr.make c n '_') = n)
+*)
+
+(*$T
+  let a = Arr.make c 1 '_' in Array.length a = 1
+  let a = Arr.make c 2 '_' in Array.length a = 2
+  let a = Arr.make c 3 '_' in Array.length a = 3
+  let a = Arr.make c 4 '_' in Array.length a = 4
+*)
+
+

src/data/CCAllocCache.mli

@@ -0,0 +1,34 @@
+
+(* This file is free software, part of Logtk. See file "license" for more details. *)
+
+(** {1 Simple Cache for Allocations}
+
+  Be very careful not to use-after-free or double-free.
+
+  {b NOT THREAD SAFE}
+
+  @since NEXT_RELEASE
+
+*)
+
+module Arr : sig
+  type 'a t
+  (** Cache for 'a arrays *)
+
+  val create: ?buck_size:int -> int -> 'a t
+  (** [create n] makes a new cache of arrays up to length [n]
+      @param buck_size number of arrays cached for each array length
+      @param n maximum size of arrays put in cache *)
+
+  val make : 'a t -> int -> 'a -> 'a array
+  (** [make cache i x] is like [Array.make i x],
+      but might return a cached array instead of allocating one.
+      {b NOTE}: if the array is already allocated then it
+        will NOT be filled with [x] *)
+
+  val free : 'a t -> 'a array -> unit
+  (** Return array to the cache. The array's elements will not be GC'd *)
+
+  val with_ : 'a t -> int -> 'a -> f:('a array -> 'b) -> 'b
+  (** Combines {!make} and {!free} *)
+end