OCaml Lru_cache using Core v0.11

lru_cache.ml

(** Replacement for deprecated Core_extended.Cache *)
open Core

module type S = sig
  type key
  type 'a t

  val make : max_entries:int -> 'a t
  val lookup : 'a t -> key -> 'a option
  val set : 'a t -> key:key -> data:'a -> unit
  val size : 'a t -> int
  val clear : 'a t -> unit
  val data : 'a t -> (key * 'a) list
end

module Make (T : Hash_queue.Key) : S with type key = T.t = struct
  type key = T.t

  module H = Hash_queue.Make (T)

  type 'a t =
    { hash_queue : 'a H.t
    ; max_entries : int
    }

  let make ~max_entries =
    assert (max_entries >= 1);
    { hash_queue = H.create (); max_entries }
  ;;

  let lookup { hash_queue } key = H.lookup_and_move_to_back hash_queue key

  let set { hash_queue; max_entries } ~key ~data =
    (* Remove instead of replacing to ensure the new key is at the back of the queue *)
    if H.mem hash_queue key then H.remove_exn hash_queue key;
    (* Remove from max entries to stay within the given limits *)
    if H.length hash_queue >= max_entries
    then (H.dequeue hash_queue : 'a option) |> ignore;
    (* Finally add the new key *)
    H.enqueue_exn hash_queue key data
  ;;

  let size { hash_queue } = H.length hash_queue
  let clear { hash_queue } = H.clear hash_queue

  let data { hash_queue } =
    H.keys hash_queue
    |> List.map ~f:(fun key ->
           (* We intentionally use H.lookup and not lookup because we don't want to reset the queue order *)
           key, H.lookup_exn hash_queue key)
  ;;
end

module Int_cache = Make (Int)

let%test_unit "cache behavior" =
  let max_entries = 2 in
  let c = Int_cache.make ~max_entries in
  List.iter [ 1; 2; 3; 4; 5 ] ~f:(fun n -> Int_cache.set c ~key:n ~data:(Int.to_string n));
  [%test_result: int] (Int_cache.size c) ~expect:max_entries;
  Int_cache.data c |> [%test_result: (int * string) list] ~expect:[ 4, "4"; 5, "5" ];
  (* If we reset 4 and then add 3, the 5 should be dropped instead of 4 *)
  Int_cache.set c ~key:4 ~data:"new 4";
  Int_cache.set c ~key:6 ~data:"6";
  [%test_result: int] (Int_cache.size c) ~expect:max_entries;
  Int_cache.data c |> [%test_result: (int * string) list] ~expect:[ 4, "new 4"; 6, "6" ];
  (* clear the queue *)
  Int_cache.clear c;
  [%test_result: int] (Int_cache.size c) ~expect:0;
  Int_cache.data c |> [%test_result: (int * string) list] ~expect:[]
;;

(* Rename this so externally it's just Cache.Int *)
module Int = Int_cache

The .mli:

(** Replacement for deprecated Core_extended.Cache *)
open Core

module type S = sig
  type key
  type 'a t

  val make : max_entries:int -> 'a t

  (** [lookup t key] returns the data for that key if it exists and makes that data the most recently-used (will be
  removed last, unless another set or lookup is called) *)
  val lookup : 'a t -> key -> 'a option

  (** [set t ~key ~data] adds [data] to the cache as the most recently used data  (will be
  removed last, unless another set or lookup is called) *)
  val set : 'a t -> key:key -> data:'a -> unit

  val size : 'a t -> int
  val clear : 'a t -> unit

  (** [data t] returns an alist of the current data in the cache. It's length will always be equal to [size]
  and will never exceed [max_entries] *)
  val data : 'a t -> (key * 'a) list
end

module Make (T : Hash_queue.Key) : S
module Int : S with type key = int