(* 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

  type data = A.key

  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 f ?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 f ?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 cache.finder ?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
