(* Ocsigen
 * Copyright (C) 2009
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, with linking exception;
 * either version 2.1 of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
*)
(**
   Cache.

   @author Vincent Balat
   @author Raphaël Proust (adding timers)
*)


let (>>=) = Lwt.bind

module Dlist :
sig

  type 'a t
  type 'a node
  val create : ?timer:float -> int -> 'a t
  val add : 'a -> 'a t -> 'a option
  val newest : 'a t -> 'a node option
  val oldest : 'a t -> 'a node option

  (** Removes an element from its list.
      If the element is not in a list, it does nothing.
      If it is in a list, it calls the finaliser, then removes the element.
      If the finaliser fails with an exception,
      the element is removed and the exception is raised again.
  *)
  val remove : 'a node -> unit

  (** Removes the element from its list without finalising,
      then adds it as newest. *)
  val up : 'a node -> unit

  val size : 'a t -> int
  val maxsize : 'a t -> int

  (** returns the timer of the Dlist *)
  val get_timer : 'a t -> float option

  val value : 'a node -> 'a
  val list_of : 'a node -> 'a t option

  (** remove the n oldest values ;
      returns the list of removed values *)
  val remove_n_oldest : 'a t -> int -> 'a list

  (** fold over the elements from the cache starting from the newest
      to the oldest *)
  val fold : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b

  (** fold over the elements from the cache starting from the oldest
      to the newest *)
  val fold_back : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b

  (** lwt version of fold *)
  val lwt_fold : ('b -> 'a -> 'b Lwt.t) -> 'b -> 'a t -> 'b Lwt.t

  (** lwt version of fold_back *)
  val lwt_fold_back : ('b -> 'a -> 'b Lwt.t) -> 'b -> 'a t -> 'b Lwt.t

  (** Move a node from one dlist to another one, without finalizing.
      If one value is removed from the destination list (because its
      maximum size is reached), it is returned (after finalisation). *)
  val move : 'a node -> 'a t -> 'a option

  (** change the maximum size ;
      returns the list of removed values, if any. *)
  val set_maxsize : 'a t -> int -> 'a list

  (** record a function to be called automatically on a piece of data
      just before it disappears from the list
      (either by explicit removal or because the maximum size is exceeded) *)
  val add_finaliser_before : ('a node -> unit) -> 'a t -> unit

  (** replace all finalizers by a new one. Be very careful while using this. *)
  val set_finaliser_before : ('a node -> unit) -> 'a t -> unit

  (** returns the finalizers. *)
  val get_finaliser_before : 'a t -> ('a node -> unit)

  (** record a function to be called automatically on a piece of data
      just after it disappears from the list
      (either by explicit removal or because the maximum size is exceeded) *)
  val add_finaliser_after : ('a node -> unit) -> 'a t -> unit

  (** replace all finalizers by a new one. Be very careful while using this. *)
  val set_finaliser_after : ('a node -> unit) -> 'a t -> unit

  (** returns the finalizers. *)
  val get_finaliser_after : 'a t -> ('a node -> unit)
end = struct

  type 'a node =
    { mutable value : 'a;
      mutable succ : 'a node option; (* the node added just after *)
      mutable prev : 'a node option; (* the node added just before *)
      mutable mylist : 'a t option; (* the list to which it belongs *)
      mutable collection : float option; (* the timestamp for removal *)
    }

  (* Doubly-linked list with maximum size.
     The field [oldest] is the first
     element that must be removed if the list becomes too long.
  *)
  and 'a t =
    {mutable newest : 'a node option (* None = empty *);
     mutable oldest : 'a node option;
     mutable size : int;
     mutable maxsize : int;
     mutable finaliser_before : 'a node -> unit;
     mutable finaliser_after : 'a node -> unit;
     (*   *) time_bound : time_bound option;
    }
  and time_bound =
    { (*   *) timer : float;
            mutable collector : unit Lwt.t option;
    }

  (* Checks (by BY):

     let compute_length c =
      let rec aux i = function
        | Some {prev=p} -> aux (i + 1) p
        | None -> i
      in aux 0 c.newest

     let correct_node n =
      (match n.succ with
         | None -> true
         | Some n' -> n'.prev == Some n) &&
       (match n.prev with
          | None -> true
          | Some n' -> n'.succ == Some n)

     (* Check that a list is correct. To be completed
       1. by adding a check on nodes,
       2. by verifying that newest can be reached from oldest and respectively *)
     let correct_list l =
      (l.size <= l.maxsize) &&
      (compute_length l = l.size) &&
      (match l.oldest with
         | None -> true
         | Some n -> n.prev = None) &&
      (match l.newest with
         | None -> true
         | Some n -> n.succ = None)
  *)

  let create ?timer size =
    {newest = None;
     oldest = None;
     size = 0;
     maxsize = size;
     finaliser_before = (fun _ -> ());
     finaliser_after = (fun _ -> ());
     time_bound =
       (match timer with
        | None -> None
        | Some t -> Some {timer = t; collector = None;}
       );
    }

  (* Remove an element from its list - don't finalise *)
  let remove' node l =
    (* assertion (node.mylist = Some l' with l' == l); *)
    let oldest =
      match l.oldest with
      | Some n when node == n -> node.succ
      | _ -> l.oldest
    in
    let newest =
      match l.newest with
      | Some n when node == n -> node.prev
      | _ -> l.newest
    in
    (match node.succ with
     | None -> ()
     | Some s -> s.prev <- node.prev);
    (match node.prev with
     | None -> ()
     | Some s -> s.succ <- node.succ);
    l.oldest <- oldest;
    l.newest <- newest;
    node.mylist <- None;
    l.size <- l.size - 1

  (* Remove an element from its list - and finalise *)
  let remove node =
    match node.mylist with
    | None -> ()
    | Some l as a ->
      (try
         l.finaliser_before node;
         assert (node.mylist == a);
         remove' node l
       with e ->
         remove' node l;
         raise e);
      l.finaliser_after node

  (* These next functions are for the collecting thread *)

  (* computing the timestamp for a node *)
  let collect_timer = function
    | {time_bound = Some {timer = t; _}; _  } -> Some (t +. (Unix.gettimeofday ()))
    | {time_bound = None; _} -> None

  (* do collect. We first check if the node is still in the list and then if
   * its collection hasn't been rescheduled ! *)
  let collect dl n = match n.mylist with
    | Some l when l == dl ->
      begin match n.collection with
        | None -> assert false
        | Some c -> if c < Unix.gettimeofday () then remove n else ()
      end
    | None | Some _ -> ()

  let sleep_until = function (*/!\ COOPERATES*)
    | None -> assert false (* collection is set to None and collector to Some *)
    | Some t ->
      let duration = t -. Unix.gettimeofday () in
      if duration <= 0.
      then Lwt.return ()
      else Lwt_unix.sleep duration

  (* a function to set the collector. *)
  let rec update_collector r = match r.time_bound with
    | None (* Not time bounded dlist *)
    | Some {collector = Some _; _} -> () (* Already collecting *)
    | Some ({collector = None; _} as t) -> match r.oldest with
      | None -> () (* Empty dlist *)
      | Some n ->
        t.collector <- Some (sleep_until n.collection >>= fun () ->
                             collect r n;
                             t.collector <- None;
                             update_collector r;
                             Lwt.return ()
                            )


  (* Add a node that do not belong to any list to a list.
     The fields [succ] and [prev] are overridden.
     If the list is too long, the function returns the oldest value.
     The node added becomes the element [list] of the list *)
  (* do not finalise *)
  (* not exported *)
  let add_node node r =
    assert (node.mylist = None);
    node.mylist <- Some r;
    let res =
      match r.newest with
      | None ->
        node.succ <- None;
        node.prev <- None;
        r.newest <- Some node;
        r.oldest <- r.newest;
        r.size <- 1;
        None
      | Some rl ->
        node.succ <- None;
        node.prev <- r.newest;
        rl.succ <- Some node;
        r.newest <- Some node;
        r.size <- r.size + 1;
        if r.size > r.maxsize
        then r.oldest
        else None
    in
    node.collection <- collect_timer r;
    update_collector r;
    res

  let add x l =
    let create_one a = {
      value = a;
      succ = None;
      prev = None;
      mylist = None;
      collection = None;
    } in
    (* create_one not exported *)
    match add_node (create_one x) l with
    | None -> None
    | Some v -> remove v; Some v.value

  let newest a = a.newest
  let oldest a = a.oldest
  let size c = c.size
  let maxsize c = c.maxsize
  let get_timer c = match c.time_bound with
    | None -> None
    | Some tb -> Some tb.timer

  let value n = n.value
  let list_of n = n.mylist

  let up node =
    match node.mylist with
    | None -> ()
    | Some l ->
      match l.newest with
      | Some n when node == n -> ()
      | _ ->
        remove' node l;
        ignore (add_node node l) (* assertion: = None *)
  (* we must not change the physical address => use add_node *)

  let rec remove_n_oldest l n = (* remove the n oldest values
                                   (or less if the list is not long enough) ;
                                   returns the list of removed values *)
    if n <= 0
    then []
    else
      match l.oldest with
      | None -> []
      | Some node ->
        let v = node.value in
        remove node; (* and finalise! *)
        v::remove_n_oldest l (n-1)


  (* Move a node from one dlist to another one, without finalizing *)
  let move node l =
    (match node.mylist with
     | None -> ()
     | Some l -> remove' node l);
    match add_node node l with
    | None -> None
    | Some v -> remove v; Some v.value

  (* fold over the elements from the newest to the oldest *)
  let lwt_fold f accu {newest; _} =
    match newest with
    | None -> Lwt.return accu
    | Some newest ->
      let rec fold accu node =
        f accu node.value >>= fun accu ->
        match node.prev with
        | None -> Lwt.return accu
        | Some new_node when new_node == newest -> Lwt.return accu
        | Some new_node ->
          fold accu new_node
      in
      fold accu newest

  (* fold over the elements from the oldest to the newest *)
  let lwt_fold_back f accu {oldest; _} =
    match oldest with
    | None -> Lwt.return accu
    | Some oldest ->
      let rec fold accu node =
        f accu node.value >>= fun accu ->
        match node.succ with
        | None -> Lwt.return accu
        | Some new_node when new_node == oldest -> Lwt.return accu
        | Some new_node ->
          fold accu new_node
      in
      fold accu oldest

  (* fold over the elements from the newest to the oldest *)
  let fold f accu {newest; _} =
    match newest with
    | None -> accu
    | Some newest ->
      let rec fold accu node =
        let accu = f accu node.value in
        match node.prev with
        | None -> accu
        | Some new_node when new_node == newest -> accu
        | Some new_node ->
          fold accu new_node
      in
      fold accu newest

  (* fold over the elements from the oldest to the newest *)
  let fold_back f accu {oldest; _} =
    match oldest with
    | None -> accu
    | Some oldest ->
      let rec fold accu node =
        let accu = f accu node.value in
        match node.succ with
        | None -> accu
        | Some new_node when new_node == oldest -> accu
        | Some new_node ->
          fold accu new_node
      in
      fold accu oldest

  let set_maxsize l m =
    let size = l.size in
    if m >= size
    then (l.maxsize <- m; [])
    else if m <= 0
    then failwith "Dlist.set_maxsize"
    else
      let ll = remove_n_oldest l (size - m) in
      l.maxsize <- m;
      ll

  let set_finaliser_before f l = l.finaliser_before <- f

  let get_finaliser_before l = l.finaliser_before

  let add_finaliser_before f l =
    let oldf = l.finaliser_before in
    l.finaliser_before <- (fun n -> oldf n; f n)

  let set_finaliser_after f l = l.finaliser_after <- f

  let get_finaliser_after l = l.finaliser_after

  let add_finaliser_after f l =
    let oldf = l.finaliser_after in
    l.finaliser_after <- (fun n -> oldf n; f n)

end


module Weak =  Weak.Make(struct type t = unit -> unit
    let hash = Hashtbl.hash
    let equal = (==)
  end)

let clear_all = Weak.create 17

module Make =
  functor (A: sig
             type key
             type value
           end) ->
  struct

    module H = Hashtbl.Make(
      struct
        type t = A.key
        let equal a a' = a = a'
        let hash = Hashtbl.hash
      end)

    type t =
      { mutable pointers : A.key Dlist.t;
        mutable table : (A.value * A.key Dlist.node) H.t;
        finder : A.key -> A.value Lwt.t;
        clear: unit -> unit (* This function clears the cache. It is put inside the
                               cache structure so that it is garbage-collected only when the cache
                               is no longer referenced, as the functions themselves are put inside
                               a weak hash table *)
      }

    let mk ?timer size =
      let (l, t) as a = (Dlist.create ?timer size, H.create size) in
      Dlist.set_finaliser_after
        (fun n -> H.remove t (Dlist.value n))
        l;
      a

    let rec create f ?timer size =
      let rec cache =
        let (l, t) = mk ?timer size in
        {pointers = l;
         table = t;
         finder = f;
         clear = f_clear;
        }
      and f_clear = (fun () -> clear cache)
      in
      Weak.add clear_all f_clear;
      cache

    and clear cache =
      let size = Dlist.maxsize cache.pointers in
      let timer = Dlist.get_timer cache.pointers in
      let (l, t) = mk ?timer size in
      cache.pointers <- l;
      cache.table <- t


    (* not exported *)
    let poke cache node =
      assert (match Dlist.list_of node with
          | None -> false
          | Some l -> cache.pointers == l);
      Dlist.up node

    let find_in_cache cache k =
      let (v, node) = H.find cache.table k in
      poke cache node;
      v

    let remove cache k =
      try
        let (_v, node) = H.find cache.table k in
        assert (match Dlist.list_of node with
            | None -> false | Some l -> cache.pointers == l);
        Dlist.remove node
      with Not_found -> ()

    (* Add in a cache, under the hypothesis that the value is
       not already in the cache *)
    let add_no_remove cache k v =
      ignore (Dlist.add k cache.pointers);
      match Dlist.newest cache.pointers with
      | None -> assert false
      | Some n -> H.add cache.table k (v, n)

    let add cache k v =
      remove cache k;
      add_no_remove cache k v

    let size c =
      Dlist.size c.pointers

    let find cache k =
      (try
         Lwt.return (find_in_cache cache k)
       with Not_found ->
         cache.finder k >>= fun r ->
         (try (* it may have been added during cache.finder *)
            ignore (find_in_cache cache k)
          with Not_found -> add_no_remove cache k r);
         Lwt.return r)

    class cache f ?timer size_c =
      let c = create f ?timer size_c in
      object
        method clear () = clear c
        method find = find c
        method add = add c
        method size = size c
        method find_in_cache = find_in_cache c
        method remove = remove c
      end

  end

let clear_all_caches () = Weak.iter (fun f -> f ()) clear_all
