Source

ocaml-lib / cache.ml

type 't db_mod = {
    find : 't -> string -> string;
    replace : 't -> string -> string -> unit;
    remove : 't -> string -> unit
  }

type ('t,'a,'b) t = {
  db_mod : 't db_mod;
  db : 't;
  prefix : string;
  ht : ('a,'b) Hashtbl.t}

let string_of_key cache key = cache.prefix ^ Marshal.to_string key []
let string_of_data data = Marshal.to_string data []
let data_of_string s = Marshal.from_string s 0

let sync cache =
  Hashtbl.iter
    (fun key data -> cache.db_mod.replace cache.db (string_of_key cache key) (string_of_data data))
    cache.ht

let create db_mod db keyprefix size =
  let cache = {
    db_mod = db_mod;
    db = db;
    prefix = keyprefix;
    ht = Hashtbl.create size} in
  at_exit (fun () -> sync cache);
  cache

let mem cache key =
  Hashtbl.mem cache.ht key
  or
    try ignore (cache.db_mod.find cache.db (string_of_key cache key)); true
    with Not_found -> false

let get (cache : ('t,'a,'b) t) key =
  try Hashtbl.find cache.ht key
  with Not_found ->
    let data = (data_of_string (cache.db_mod.find cache.db (string_of_key cache key)) : 'b) in
    Hashtbl.add cache.ht key data;
    data

let unget cache key =
  try
    cache.db_mod.replace cache.db (string_of_key cache key) (string_of_data (Hashtbl.find cache.ht key));
    Hashtbl.remove cache.ht key
  with Not_found -> ()

let fold f cache e = Hashtbl.fold f cache.ht e

let add cache key data =
  Hashtbl.add cache.ht key data

let put cache key =
  try cache.db_mod.replace cache.db (string_of_key cache key) (string_of_data (Hashtbl.find cache.ht key))
  with Not_found -> ()

let remove cache key =
  if Hashtbl.mem cache.ht key
  then Hashtbl.remove cache.ht key;
  try cache.db_mod.remove cache.db (string_of_key cache key) with _ -> ()