hashweak.ml 6.36 KB
Newer Older
1 2 3
(**************************************************************************)
(*                                                                        *)
(*  Copyright (C) 2010-                                                   *)
MARCHE Claude's avatar
MARCHE Claude committed
4 5 6
(*    François Bobot                                                     *)
(*    Jean-Christophe Filliâtre                                          *)
(*    Claude Marché                                                      *)
7 8 9 10 11 12 13 14 15 16 17 18 19
(*    Andrei Paskevich                                                    *)
(*                                                                        *)
(*  This software is free software; you can redistribute it and/or        *)
(*  modify it under the terms of the GNU Library General Public           *)
(*  License version 2.1, with the special exception on linking            *)
(*  described in file LICENSE.                                            *)
(*                                                                        *)
(*  This software 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.                  *)
(*                                                                        *)
(**************************************************************************)

20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
module ProdConsume :
sig
  type 'a t
  val create : unit -> 'a t
  val add : 'a -> 'a t -> unit
  val iter_remove : ('a -> unit) -> 'a t -> unit
end
= struct
  (* One thing can produce, one thing can consume concurrently *)

  type 'a cell =
    | Empty
    | Cons of 'a * 'a list
  and 'a list = 'a cell ref

  let rec iter f = function
    | Empty -> ()
    | Cons (x,l) -> f x; iter f !l

  (* a reference on a mutable singly linked list *)
  type 'a t = 'a list ref

  let create () = ref (ref (Empty))
  let add x t = t := ref (Cons(x,!t))
  let iter_remove f t =
    if !(!t) = Empty then () else
    let r = !t in (* atomic one cell of the list *)
    let l = !r in (* the content of the cell *)
    r := Empty; (* Work even if there are some production,
                   just not anymore the head *)
    iter f l
end

Andrei Paskevich's avatar
Andrei Paskevich committed
53
module type S = sig
54 55 56 57 58

  type key

  type 'a t

Andrei Paskevich's avatar
Andrei Paskevich committed
59
  val create : int -> 'a t
60 61
    (* create a hashtbl with weak keys *)

Andrei Paskevich's avatar
Andrei Paskevich committed
62 63 64 65
  val clear : 'a t -> unit

  val copy : 'a t -> 'a t

66 67 68 69 70 71 72 73 74 75
  val find : 'a t -> key -> 'a
    (* find the value bound to a key.
       Raises Not_found if there is no binding *)

  val mem : 'a t -> key -> bool
    (* test if a key is bound *)

  val set : 'a t -> key -> 'a -> unit
    (* bind the key _once_ with the given value *)

Andrei Paskevich's avatar
Andrei Paskevich committed
76 77
  val remove : 'a t -> key -> unit
    (* remove the value *)
78

Andrei Paskevich's avatar
Andrei Paskevich committed
79
  val iter : (key -> 'a -> unit) -> 'a t -> unit
80

Andrei Paskevich's avatar
Andrei Paskevich committed
81
  val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
82

Andrei Paskevich's avatar
Andrei Paskevich committed
83
  val iterk : (key -> unit) -> 'a t -> unit
84

Andrei Paskevich's avatar
Andrei Paskevich committed
85
  val foldk : (key -> 'b -> 'b) -> 'a t -> 'b -> 'b
86

Andrei Paskevich's avatar
Andrei Paskevich committed
87
  val length : 'a t -> int
88

Andrei Paskevich's avatar
Andrei Paskevich committed
89 90
  val memoize : int -> (key -> 'a) -> (key -> 'a)
    (* create a memoizing function *)
91

Andrei Paskevich's avatar
Andrei Paskevich committed
92 93
  val memoize_option : int -> (key option -> 'a) -> (key option -> 'a)
    (* memoizing functions on option types *)
94

Andrei Paskevich's avatar
Andrei Paskevich committed
95
end
96

97 98 99
let new_tbl_tag = let c = ref (-1) in
  fun () -> (incr c; !c)

Andrei Paskevich's avatar
Andrei Paskevich committed
100 101 102 103
type tag = {
  tag_map : ((int,Obj.t) Hashtbl.t) Lazy.t;
  tag_tag : int;
}
104

Andrei Paskevich's avatar
Andrei Paskevich committed
105 106 107 108
let create_tag tag = {
  tag_map = lazy (Hashtbl.create 3);
  tag_tag = tag;
}
109

Andrei Paskevich's avatar
Andrei Paskevich committed
110 111 112 113
let dummy_tag = {
  tag_map = lazy (failwith "dummy tag");
  tag_tag = -1;
}
114

115
let tag_equal : tag -> tag -> bool = (==)
116

117
let tag_hash k = assert (k != dummy_tag); k.tag_tag
118

119
module type Weakey =
120 121
sig
  type t
Andrei Paskevich's avatar
Andrei Paskevich committed
122
  val tag : t -> tag
123 124
end

125
module Make (S : Weakey) = struct
126

127
  type key = S.t
128

Andrei Paskevich's avatar
Andrei Paskevich committed
129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148
  module H = Weak.Make (struct
    type t = S.t
    let hash k = (S.tag k).tag_tag
    let equal k1 k2 = S.tag k1 == S.tag k2
  end)

  type 'a t = {
    tbl_set : H.t;
    tbl_tag : int;
  }

  let tag_map k = Lazy.force (S.tag k).tag_map

  let find (t : 'a t) k : 'a =
    Obj.obj (Hashtbl.find (tag_map k) t.tbl_tag)

  let mem t k = Hashtbl.mem (tag_map k) t.tbl_tag

  let set (t : 'a t) k (v : 'a) =
    Hashtbl.replace (tag_map k) t.tbl_tag (Obj.repr v);
149
    ignore (H.merge t.tbl_set k)
Andrei Paskevich's avatar
Andrei Paskevich committed
150 151 152 153 154 155 156 157 158 159 160

  let remove t k =
    Hashtbl.remove (tag_map k) t.tbl_tag;
    H.remove t.tbl_set k

  let iterk fn t = H.iter fn t.tbl_set
  let foldk fn t = H.fold fn t.tbl_set

  let iter  fn t = H.iter (fun k -> fn k (find t k)) t.tbl_set
  let fold  fn t = H.fold (fun k -> fn k (find t k)) t.tbl_set

161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178
  (** This table is just a hack to keep alive the weak hashset :
      Indeed that circunvent a strange behavior/bug of weak hashset,
      when a weak hashset is garbage collected it will not anymore
      remove the dead elements from it. So during finalize or if the
      hashset is keep alive, we can acces invalid pointer...

      So to summarize we keep alive the weak hashset until we don't need them
      anymore.
  *)
  let gen_table = Hashtbl.create 5

  let tbl_final_aux t =
    iterk (fun k -> Hashtbl.remove (tag_map k) t.tbl_tag) t

  let tbl_final t =
    tbl_final_aux t;
    (** We don't need anymore the weak hashset, we can release it *)
    Hashtbl.remove gen_table t.tbl_tag
179

180 181 182 183 184 185 186 187 188 189 190
  (** All the hashweak that can be collected. When a hashweak is
      garbage collected we need to remove its tag from the key
      hashtable. Since finalisation can be triggered at anytime even
      when the key hashtable are in an inconsistent state, we need to
      delay the actual removing until it can be done safely.
      t_collected contains the delayed hashweak *)
  let t_collected = ProdConsume.create ()

  (** Do really the removing *)
  let collect () = ProdConsume.iter_remove tbl_final t_collected

191
  let create n =
Andrei Paskevich's avatar
Andrei Paskevich committed
192 193
    let t = {
      tbl_set = H.create n;
194
      tbl_tag = new_tbl_tag () }
Andrei Paskevich's avatar
Andrei Paskevich committed
195
    in
196
    Hashtbl.add gen_table t.tbl_tag t.tbl_set;
197
    Gc.finalise (fun t -> ProdConsume.add t t_collected) t;
Andrei Paskevich's avatar
Andrei Paskevich committed
198
    t
199

200 201 202
  let find x y = collect (); find x y
  let set x y z = collect (); set x y z

203
  let clear t = collect (); tbl_final_aux t; H.clear t.tbl_set
204

Andrei Paskevich's avatar
Andrei Paskevich committed
205
  let length t = H.count t.tbl_set
206

Andrei Paskevich's avatar
Andrei Paskevich committed
207
  let copy t =
208
    collect ();
Andrei Paskevich's avatar
Andrei Paskevich committed
209 210 211
    let t' = create (length t) in
    iter (set t') t;
    t'
212

Andrei Paskevich's avatar
Andrei Paskevich committed
213 214 215
  let memoize n fn =
    let h = create n in fun e ->
      try find h e
216 217
      with Not_found ->
        let v = fn e in
Andrei Paskevich's avatar
Andrei Paskevich committed
218
        set h e v;
219
        v
220

Andrei Paskevich's avatar
Andrei Paskevich committed
221
  let memoize_option n fn =
222
    let v = lazy (fn None) in
223
    let fn e = fn (Some e) in
Andrei Paskevich's avatar
Andrei Paskevich committed
224
    let fn = memoize n fn in
225 226
    function
      | Some e -> fn e
227
      | None -> Lazy.force v
228

Francois Bobot's avatar
 
Francois Bobot committed
229 230
end