ptset.ml

   1 (***************************************************************************)
   2 (* Implementation for sets of positive integers implemented as deeply hash-*)
   3 (* consed Patricia trees. Provide fast set operations, fast membership as  *)
   4 (* well as fast min and max elements. Hash consing provides O(1) equality  *)
   5 (* checking                                                                *)
   6 (*                                                                         *)
   7 (***************************************************************************)
   8
   9
  10 type elt = int
  11
  12 type t = { id : int;
  13            key : int; (* hash *)
  14            node : node;
  15            }
  16 and node =
  17   | Empty
  18   | Leaf of int
  19   | Branch of int * int * t * t
  20
  21
  22 (* faster if outside of a module *)
  23 let hash_node x = match x with
  24   | Empty -> 0
  25   | Leaf i -> (i+1) land max_int
  26       (* power of 2 +/- 1 are fast ! *)
  27   | Branch (b,i,l,r) ->
  28       ((b lsl 1)+ b + i+(i lsl 4) + (l.key lsl 5)-l.key
  29        + (r.key lsl 7) - r.key) land max_int
  30
  31 module Node =
  32   struct
  33     type _t = t
  34     type t = _t
  35     external hash : t -> int = "%field1"
  36     let equal x y =
  37       if x.id == y.id || x.key == y.key || x.node == y.node then true
  38       else
  39       match (x.node,y.node) with
  40       | Empty,Empty -> true
  41       | Leaf k1, Leaf k2 when k1 == k2 -> true
  42       | Branch(p1,m1,l1,r1), Branch(p2,m2,l2,r2) when m1==m2 && p1==p2 &&
  43           (l1.id == l2.id) && (r1.id == r2.id) -> true
  44       | _ -> false
  45   end
  46
  47 module WH =Weak.Make(Node)
  48
  49 let pool = WH.create 4093
  50
  51 (* Neat trick thanks to Alain Frisch ! *)
  52
  53 let gen_uid () = Oo.id (object end)
  54
  55 let empty = { id = gen_uid ();
  56               key = 0;
  57               node = Empty }
  58
  59 let _ = WH.add pool empty
  60
  61 let is_empty s = s.id==0
  62
  63 let rec norm n =
  64   let v = { id = gen_uid ();
  65             key = hash_node n;
  66             node = n }
  67   in
  68       WH.merge pool v
  69
  70 (*  WH.merge pool *)
  71
  72 let branch  p m l r  = norm (Branch(p,m,l,r))
  73 let leaf k = norm (Leaf k)
  74
  75 (* To enforce the invariant that a branch contains two non empty sub-trees *)
  76 let branch_ne = function
  77   | (_,_,e,t) when is_empty e -> t
  78   | (_,_,t,e) when is_empty e -> t
  79   | (p,m,t0,t1)   -> branch p m t0 t1
  80
  81 (********** from here on, only use the smart constructors *************)
  82
  83 let zero_bit k m = (k land m) == 0
  84
  85 let singleton k = leaf k
  86 let is_singleton n =
  87   match n.node with Leaf _ -> true
  88     | _ -> false
  89
  90 let rec mem k n = match n.node with
  91   | Empty -> false
  92   | Leaf j -> k == j
  93   | Branch (p, _, l, r) -> if k <= p then mem k l else mem k r
  94
  95 let rec min_elt n = match n.node with
  96   | Empty -> raise Not_found
  97   | Leaf k -> k
  98   | Branch (_,_,s,_) -> min_elt s
  99
 100   let rec max_elt n = match n.node with
 101     | Empty -> raise Not_found
 102     | Leaf k -> k
 103     | Branch (_,_,_,t) -> max_elt t
 104
 105   let elements s =
 106     let rec elements_aux acc n = match n.node with
 107       | Empty -> acc
 108       | Leaf k -> k :: acc
 109       | Branch (_,_,l,r) -> elements_aux (elements_aux acc r) l
 110     in
 111     elements_aux [] s
 112
 113   let mask k m  = (k lor (m-1)) land (lnot m)
 114
 115   let naive_highest_bit x =
 116     assert (x < 256);
 117     let rec loop i =
 118       if i = 0 then 1 else if x lsr i = 1 then 1 lsl i else loop (i-1)
 119     in
 120     loop 7
 121
 122   let hbit = Array.init 256 naive_highest_bit
 123
 124   let highest_bit_32 x =
 125     let n = x lsr 24 in if n != 0 then Array.unsafe_get hbit n lsl 24
 126     else let n = x lsr 16 in if n != 0 then Array.unsafe_get hbit n lsl 16
 127     else let n = x lsr 8 in if n != 0 then Array.unsafe_get hbit n lsl 8
 128     else Array.unsafe_get hbit x
 129
 130   let highest_bit_64 x =
 131     let n = x lsr 32 in if n != 0 then (highest_bit_32 n) lsl 32
 132     else highest_bit_32 x
 133
 134   let highest_bit = match Sys.word_size with
 135     | 32 -> highest_bit_32
 136     | 64 -> highest_bit_64
 137     | _ -> assert false
 138
 139   let branching_bit p0 p1 = highest_bit (p0 lxor p1)
 140
 141   let join p0 t0 p1 t1 =
 142     let m = branching_bit p0 p1  in
 143     if zero_bit p0 m then
 144       branch (mask p0 m)  m t0 t1
 145     else
 146       branch (mask p0 m) m t1 t0
 147
 148   let match_prefix k p m = (mask k m) == p
 149
 150   let add k t =
 151     let rec ins n = match n.node with
 152       | Empty -> leaf k
 153       | Leaf j ->  if j == k then n else join k (leaf k) j n
 154       | Branch (p,m,t0,t1)  ->
 155           if match_prefix k p m then
 156             if zero_bit k m then
 157               branch p m (ins t0) t1
 158             else
 159               branch p m t0 (ins t1)
 160           else
 161             join k  (leaf k)  p n
 162     in
 163     ins t
 164
 165   let remove k t =
 166     let rec rmv n = match n.node with
 167       | Empty -> empty
 168       | Leaf j  -> if k == j then empty else n
 169       | Branch (p,m,t0,t1) ->
 170           if match_prefix k p m then
 171             if zero_bit k m then
 172               branch_ne (p, m, rmv t0, t1)
 173             else
 174               branch_ne (p, m, t0, rmv t1)
 175           else
 176             n
 177     in
 178     rmv t
 179
 180   (* should run in O(1) thanks to Hash consing *)
 181
 182   let equal a b = a==b || a.id == b.id
 183
 184   let compare a b = if a == b then 0 else a.id - b.id
 185
 186   let h_merge = Hashtbl.create 4097
 187   let com_hash x y = (x*y - (x+y)) land max_int
 188
 189   let rec merge s t =
 190     if (equal s t) (* This is cheap thanks to hash-consing *)
 191     then s
 192     else
 193     match s.node,t.node with
 194       | Empty, _  -> t
 195       | _, Empty  -> s
 196       | Leaf k, _ -> add k t
 197       | _, Leaf k -> add k s
 198       | Branch (p,m,s0,s1), Branch (q,n,t0,t1) ->
 199           if m == n && match_prefix q p m then
 200             branch p  m  (merge s0 t0) (merge s1 t1)
 201           else if m > n && match_prefix q p m then
 202             if zero_bit q m then
 203               branch p m (merge s0 t) s1
 204             else
 205               branch p m s0 (merge s1 t)
 206           else if m < n && match_prefix p q n then
 207             if zero_bit p n then
 208               branch q n (merge s t0) t1
 209             else
 210               branch q n t0 (merge s t1)
 211           else
 212             (* The prefixes disagree. *)
 213             join p s q t
 214
 215
 216
 217
 218   let rec subset s1 s2 = (equal s1 s2) ||
 219     match (s1.node,s2.node) with
 220       | Empty, _ -> true
 221       | _, Empty -> false
 222       | Leaf k1, _ -> mem k1 s2
 223       | Branch _, Leaf _ -> false
 224       | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
 225           if m1 == m2 && p1 == p2 then
 226             subset l1 l2 && subset r1 r2
 227           else if m1 < m2 && match_prefix p1 p2 m2 then
 228             if zero_bit p1 m2 then
 229               subset l1 l2 && subset r1 l2
 230             else
 231               subset l1 r2 && subset r1 r2
 232           else
 233             false
 234
 235
 236
 237
 238   let union s1 s2 = merge s1 s2
 239
 240   let rec inter s1 s2 =
 241     if equal s1 s2
 242     then s1
 243     else
 244       match (s1.node,s2.node) with
 245         | Empty, _ -> empty
 246         | _, Empty -> empty
 247         | Leaf k1, _ -> if mem k1 s2 then s1 else empty
 248         | _, Leaf k2 -> if mem k2 s1 then s2 else empty
 249         | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
 250             if m1 == m2 && p1 == p2 then
 251               merge (inter l1 l2)  (inter r1 r2)
 252             else if m1 > m2 && match_prefix p2 p1 m1 then
 253               inter (if zero_bit p2 m1 then l1 else r1) s2
 254             else if m1 < m2 && match_prefix p1 p2 m2 then
 255               inter s1 (if zero_bit p1 m2 then l2 else r2)
 256             else
 257               empty
 258
 259   let rec diff s1 s2 =
 260     if equal s1 s2
 261     then empty
 262     else
 263       match (s1.node,s2.node) with
 264         | Empty, _ -> empty
 265         | _, Empty -> s1
 266         | Leaf k1, _ -> if mem k1 s2 then empty else s1
 267         | _, Leaf k2 -> remove k2 s1
 268         | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
 269             if m1 == m2 && p1 == p2 then
 270               merge (diff l1 l2) (diff r1 r2)
 271             else if m1 > m2 && match_prefix p2 p1 m1 then
 272               if zero_bit p2 m1 then
 273                 merge (diff l1 s2) r1
 274               else
 275                 merge l1 (diff r1 s2)
 276             else if m1 < m2 && match_prefix p1 p2 m2 then
 277               if zero_bit p1 m2 then diff s1 l2 else diff s1 r2
 278             else
 279           s1
 280
 281
 282
 283
 284 (*s All the following operations ([cardinal], [iter], [fold], [for_all],
 285     [exists], [filter], [partition], [choose], [elements]) are
 286     implemented as for any other kind of binary trees. *)
 287
 288 let rec cardinal n = match n.node with
 289   | Empty -> 0
 290   | Leaf _ -> 1
 291   | Branch (_,_,t0,t1) -> cardinal t0 + cardinal t1
 292
 293 let rec iter f n = match n.node with
 294   | Empty -> ()
 295   | Leaf k -> f k
 296   | Branch (_,_,t0,t1) -> iter f t0; iter f t1
 297
 298 let rec fold f s accu = match s.node with
 299   | Empty -> accu
 300   | Leaf k -> f k accu
 301   | Branch (_,_,t0,t1) -> fold f t0 (fold f t1 accu)
 302
 303 let rec for_all p n = match n.node with
 304   | Empty -> true
 305   | Leaf k -> p k
 306   | Branch (_,_,t0,t1) -> for_all p t0 && for_all p t1
 307
 308 let rec exists p n = match n.node with
 309   | Empty -> false
 310   | Leaf k -> p k
 311   | Branch (_,_,t0,t1) -> exists p t0 || exists p t1
 312
 313 let rec filter pr n = match n.node with
 314   | Empty -> empty
 315   | Leaf k -> if pr k then n else empty
 316   | Branch (p,m,t0,t1) -> branch_ne (p, m, filter pr t0, filter pr t1)
 317
 318 let partition p s =
 319   let rec part (t,f as acc) n = match n.node with
 320     | Empty -> acc
 321     | Leaf k -> if p k then (add k t, f) else (t, add k f)
 322     | Branch (_,_,t0,t1) -> part (part acc t0) t1
 323   in
 324   part (empty, empty) s
 325
 326 let rec choose n = match n.node with
 327   | Empty -> raise Not_found
 328   | Leaf k -> k
 329   | Branch (_, _,t0,_) -> choose t0   (* we know that [t0] is non-empty *)
 330
 331
 332 let split x s =
 333   let coll k (l, b, r) =
 334     if k < x then add k l, b, r
 335     else if k > x then l, b, add k r
 336     else l, true, r
 337   in
 338   fold coll s (empty, false, empty)
 339
 340
 341
 342 let rec dump n =
 343   Printf.eprintf "{ id = %i; key = %i ; node=" n.id n.key;
 344   match n.node with
 345     | Empty -> Printf.eprintf "Empty; }\n"
 346     | Leaf k -> Printf.eprintf "Leaf %i; }\n" k
 347     | Branch (p,m,l,r) ->
 348         Printf.eprintf "Branch(%i,%i,id=%i,id=%i); }\n"
 349           p m l.id r.id;
 350         dump l;
 351         dump r
 352
 353 (*i*)
 354 let make l = List.fold_left (fun acc e -> add e acc ) empty l
 355 (*i*)
 356
 357 (*s Additional functions w.r.t to [Set.S]. *)
 358
 359 let rec intersect s1 s2 = (equal s1 s2) ||
 360   match (s1.node,s2.node) with
 361   | Empty, _ -> false
 362   | _, Empty -> false
 363   | Leaf k1, _ -> mem k1 s2
 364   | _, Leaf k2 -> mem k2 s1
 365   | Branch (p1,m1,l1,r1), Branch (p2,m2,l2,r2) ->
 366       if m1 == m2 && p1 == p2 then
 367         intersect l1 l2 || intersect r1 r2
 368       else if m1 < m2 && match_prefix p2 p1 m1 then
 369         intersect (if zero_bit p2 m1 then l1 else r1) s2
 370       else if m1 > m2 && match_prefix p1 p2 m2 then
 371         intersect s1 (if zero_bit p1 m2 then l2 else r2)
 372       else
 373         false
 374
 375
 376 let hash s = s.key
 377
 378 let from_list l = List.fold_left (fun acc i -> add i acc) empty l
 379
 380 type int_vector
 381
 382 external int_vector_alloc : int -> int_vector = "caml_int_vector_alloc"
 383 external int_vector_set : int_vector -> int -> int -> unit = "caml_int_vector_set"
 384 external int_vector_length : int_vector -> int  = "caml_int_vector_length"
 385 external int_vector_empty : unit -> int_vector = "caml_int_vector_empty"
 386
 387 let empty_vector = int_vector_empty ()
 388
 389 let to_int_vector_ext s =
 390   let l = cardinal s in
 391   let v = int_vector_alloc l in
 392   let i = ref 0 in
 393     iter (fun e -> int_vector_set v !i e; incr i) s;
 394     v
 395
 396 let hash_vectors = Hashtbl.create 4097
 397
 398 let to_int_vector s =
 399   try
 400     Hashtbl.find hash_vectors s.key
 401   with
 402       Not_found ->
 403         let v = to_int_vector_ext s in
 404           Hashtbl.add hash_vectors s.key v;
 405           v
 406
 407