stdlib/bitSet.ml

   1 (*
   2  * Bitset - Efficient bit sets
   3  * Copyright (C) 2003 Nicolas Cannasse
   4  *
   5  * This library is free software; you can redistribute it and/or
   6  * modify it under the terms of the GNU Lesser General Public
   7  * License as published by the Free Software Foundation; either
   8  * version 2.1 of the License, or (at your option) any later version,
   9  * with the special exception on linking described in file LICENSE.
  10  *
  11  * This library is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with this library; if not, write to the Free Software
  18  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA     02111-1307      USA
  19  *)
  20
  21 type intern
  22
  23 external bcreate : int -> intern = "caml_create_string"
  24 external fast_get : intern -> int -> int = "%string_unsafe_get"
  25 external fast_set : intern -> int -> int -> unit = "%string_unsafe_set"
  26 external fast_bool : int -> bool = "%identity"
  27 external fast_blit : intern -> int -> intern -> int -> int -> unit =  "caml_blit_string" "noalloc"
  28 external fast_fill : intern -> int -> int -> int -> unit = "caml_fill_string" "noalloc"
  29 external fast_length : intern -> int = "%string_length"
  30
  31 let bget s ndx =
  32   assert (ndx >= 0 && ndx < fast_length s);
  33   fast_get s ndx
  34
  35 let bset s ndx v =
  36   assert (ndx >= 0 && ndx < fast_length s);
  37   fast_set s ndx v
  38
  39 let bblit src srcoff dst dstoff len =
  40   assert (srcoff >= 0 && dstoff >= 0 && len >= 0);
  41   fast_blit src srcoff dst dstoff len
  42
  43 let bfill dst start len c =
  44   assert (start >= 0 && len >= 0);
  45   fast_fill dst start len c
  46
  47 exception Negative_index of string
  48
  49 type t = {
  50         mutable data : intern;
  51         mutable len : int;
  52 }
  53
  54 let error fname = raise (Negative_index fname)
  55
  56 let empty() =
  57         {
  58                 data = bcreate 0;
  59                 len = 0;
  60         }
  61
  62 let int_size = 7 (* value used to round up index *)
  63 let log_int_size = 3 (* number of shifts *)
  64
  65 let create n =
  66         if n < 0 then error "create";
  67         let size = (n+int_size) lsr log_int_size in
  68         let b = bcreate size in
  69         bfill b 0 size 0;
  70         {
  71                 data = b;
  72                 len = size;
  73         }
  74
  75 let copy t =
  76         let b = bcreate t.len in
  77         bblit t.data 0 b 0 t.len;
  78         {
  79                 data = b;
  80                 len = t.len
  81         }
  82
  83 let clone = copy
  84
  85 let set t x =
  86         if x < 0 then error "set";
  87         let pos = x lsr log_int_size and delta = x land int_size in
  88         let size = t.len in
  89         if pos >= size then begin
  90                 let b = bcreate (pos+1) in
  91                 bblit t.data 0 b 0 size;
  92                 bfill b size (pos - size + 1) 0;
  93                 t.len <- pos + 1;
  94                 t.data <- b;
  95         end;
  96         bset t.data pos ((bget t.data pos) lor (1 lsl delta))
  97
  98 let unset t x =
  99         if x < 0 then error "unset";
 100         let pos = x lsr log_int_size and delta = x land int_size in
 101         if pos < t.len then
 102                 bset t.data pos ((bget t.data pos) land (0xFF lxor (1 lsl delta)))
 103
 104 let toggle t x =
 105         if x < 0 then error "toggle";
 106         let pos = x lsr log_int_size and delta = x land int_size in
 107         let size = t.len in
 108         if pos >= size then begin
 109                 let b = bcreate (pos+1) in
 110                 bblit t.data 0 b 0 size;
 111                 bfill b size (pos - size + 1) 0;
 112                 t.len <- pos + 1;
 113                 t.data <- b;
 114         end;
 115         bset t.data pos ((bget t.data pos) lxor (1 lsl delta))
 116
 117 let put t = function
 118         | true -> set t
 119         | false -> unset t
 120
 121 let is_set t x =
 122   if x < 0 then error "is_set";
 123   let pos = x lsr log_int_size and delta = x land int_size in
 124   let size = t.len in
 125   if pos < size then
 126         fast_bool (((bget t.data pos) lsr delta) land 1)
 127   else
 128         false
 129
 130
 131 exception Break_int of int
 132
 133 (* Find highest set element or raise Not_found *)
 134 let find_msb t =
 135   (* Find highest set bit in a byte.  Does not work with zero. *)
 136   let byte_msb b =
 137     assert (b <> 0);
 138     let rec loop n =
 139       if b land (1 lsl n) = 0 then
 140         loop (n-1)
 141       else n in
 142     loop 7 in
 143   let n = t.len - 1
 144   and buf = t.data in
 145   try
 146     for i = n downto 0 do
 147       let byte = bget buf i in
 148       if byte <> 0 then raise (Break_int ((i lsl log_int_size)+(byte_msb byte)))
 149     done;
 150     raise Not_found
 151   with
 152     Break_int n -> n
 153   | _ -> raise Not_found
 154
 155 let compare t1 t2 =
 156   let some_msb b = try Some (find_msb b) with Not_found -> None in
 157   match (some_msb t1, some_msb t2) with
 158     (None, Some _) -> -1 (* 0-y -> -1 *)
 159   | (Some _, None) -> 1  (* x-0 ->  1 *)
 160   | (None, None) -> 0    (* 0-0 ->  0 *)
 161   | (Some a, Some b) ->  (* x-y *)
 162       if a < b then -1
 163       else if a > b then 1
 164       else
 165         begin
 166           (* MSBs differ, we need to scan arrays until we find a
 167              difference *)
 168           let ndx = a lsr log_int_size in
 169           assert (ndx < t1.len && ndx < t2.len);
 170           try
 171             for i = ndx downto 0 do
 172               let b1 = bget t1.data i
 173               and b2 = bget t2.data i in
 174               if b1 <> b2 then raise (Break_int (compare b1 b2))
 175             done;
 176             0
 177           with
 178             Break_int res -> res
 179         end
 180
 181 let equals t1 t2 =
 182         compare t1 t2 = 0
 183
 184 let partial_count t x =
 185         let rec nbits x =
 186                 if x = 0 then
 187                         0
 188                 else if fast_bool (x land 1) then
 189                         1 + (nbits (x lsr 1))
 190                 else
 191                         nbits (x lsr 1)
 192         in
 193         let size = t.len in
 194         let pos = x lsr log_int_size and delta = x land int_size in
 195         let rec loop n acc =
 196                 if n = size then
 197                         acc
 198                 else
 199                         let x = bget t.data n in
 200                         loop (n+1) (acc + nbits x)
 201         in
 202         if pos >= size then
 203                 0
 204         else
 205                 loop (pos+1) (nbits ((bget t.data pos) lsr delta))
 206
 207 let count t =
 208   partial_count t 0
 209
 210 (* Find the first set bit in the bit array *)
 211 let find_first_set b n =
 212   (* TODO there are many ways to speed this up.  Lookup table would be
 213      one way to speed this up. *)
 214   let find_lsb b =
 215     assert (b <> 0);
 216     let rec loop n =
 217       if b land (1 lsl n) <> 0 then n else loop (n+1) in
 218     loop 0 in
 219
 220   let buf = b.data in
 221   let rec find_bit byte_ndx bit_offs =
 222     if byte_ndx >= b.len then
 223       None
 224     else
 225       let byte = (bget buf byte_ndx) lsr bit_offs in
 226       if byte = 0 then
 227         find_bit (byte_ndx + 1) 0
 228       else
 229         Some ((find_lsb byte) + (byte_ndx lsl log_int_size) + bit_offs) in
 230   find_bit (n lsr log_int_size) (n land int_size)
 231 (*
 232 let enum t =
 233   let rec make n =
 234     let cur = ref n in
 235     let rec next () =
 236       match find_first_set t !cur with
 237         Some elem ->
 238           cur := (elem+1);
 239           elem
 240       | None ->
 241           raise Enum.No_more_elements in
 242     Enum.make
 243       ~next
 244       ~count:(fun () -> partial_count t !cur)
 245       ~clone:(fun () -> make !cur)
 246   in
 247   make 0
 248 *)
 249 let raw_create size =
 250   let b = bcreate size in
 251   bfill b 0 size 0;
 252   { data = b; len = size }
 253
 254 let inter a b =
 255   let max_size = max a.len b.len in
 256   let d = raw_create max_size in
 257   let sl = min a.len b.len in
 258   let abuf = a.data
 259   and bbuf = b.data in
 260   (* Note: rest of the array is set to zero automatically *)
 261   for i = 0 to sl-1 do
 262     bset d.data i ((bget abuf i) land (bget bbuf i))
 263   done;
 264   d
 265
 266 (* Note: rest of the array is handled automatically correct, since we
 267    took a copy of the bigger set. *)
 268 let union a b =
 269   let d = if a.len > b.len then copy a else copy b in
 270   let sl = min a.len b.len in
 271   let abuf = a.data
 272   and bbuf = b.data in
 273   for i = 0 to sl-1 do
 274     bset d.data i ((bget abuf i) lor (bget bbuf i))
 275   done;
 276   d
 277
 278 let diff a b =
 279   let maxlen = max a.len b.len in
 280   let buf = bcreate maxlen in
 281   bblit a.data 0 buf 0 a.len;
 282   let sl = min a.len b.len in
 283   let abuf = a.data
 284   and bbuf = b.data in
 285   for i = 0 to sl-1 do
 286     bset buf i ((bget abuf i) land (lnot (bget bbuf i)))
 287   done;
 288   { data = buf; len = maxlen }
 289
 290 let sym_diff a b =
 291   let maxlen = max a.len b.len in
 292   let buf = bcreate maxlen in
 293   (* Copy larger (assumes missing bits are zero) *)
 294   bblit (if a.len > b.len then a.data else b.data) 0 buf 0 maxlen;
 295   let sl = min a.len b.len in
 296   let abuf = a.data
 297   and bbuf = b.data in
 298   for i = 0 to sl-1 do
 299     bset buf i ((bget abuf i) lxor (bget bbuf i))
 300   done;
 301   { data = buf; len = maxlen }
 302
 303 (* TODO the following set operations can be made faster if you do the
 304    set operation in-place instead of taking a copy.  But be careful
 305    when the sizes of the bitvector strings differ. *)
 306 let intersect t t' =
 307   let d = inter t t' in
 308   t.data <- d.data;
 309   t.len <- d.len
 310
 311 let differentiate t t' =
 312   let d = diff t t' in
 313   t.data <- d.data;
 314   t.len <- d.len
 315
 316 let unite t t' =
 317   let d = union t t' in
 318   t.data <- d.data;
 319   t.len <- d.len
 320
 321 let differentiate_sym t t' =
 322   let d = sym_diff t t' in
 323   t.data <- d.data;
 324   t.len <- d.len