(*
 * Bitset - Efficient bit sets
 * Copyright (C) 2003 Nicolas Cannasse
 *
 * This library 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; either
 * version 2.1 of the License, or (at your option) any later version,
 * with the special exception on linking described in file LICENSE.
 *
 * This library 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 library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

open ExtBytes

type intern

let bcreate : int -> intern = Obj.magic Bytes.create
external fast_get : intern -> int -> int = "%string_unsafe_get"
let fast_set : intern -> int -> int -> unit = Obj.magic Bytes.unsafe_set
external fast_bool : int -> bool = "%identity"
let fast_blit : intern -> int -> intern -> int -> int -> unit = Obj.magic Bytes.blit
let fast_fill : intern -> int -> int -> int -> unit = Obj.magic Bytes.fill
let fast_length : intern -> int= Obj.magic Bytes.length

let bget s ndx =
  assert (ndx >= 0 && ndx < fast_length s);
  fast_get s ndx

let bset s ndx v =
  assert (ndx >= 0 && ndx < fast_length s);
  fast_set s ndx v

let bblit src srcoff dst dstoff len = 
  assert (srcoff >= 0 && dstoff >= 0 && len >= 0);
  fast_blit src srcoff dst dstoff len

let bfill dst start len c = 
  assert (start >= 0 && len >= 0);
  fast_fill dst start len c

exception Negative_index of string

type t = {
  mutable data : intern;
  mutable len : int;
}

let error fname = raise (Negative_index fname)

let empty() =
  {
    data = bcreate 0;
    len = 0;
  }

let int_size = 7 (* value used to round up index *)
let log_int_size = 3 (* number of shifts *)

let create n =
  if n < 0 then error "create";
  let size = (n+int_size) lsr log_int_size in
  let b = bcreate size in
  bfill b 0 size 0;
  {
    data = b;
    len = size;
  }

let copy t =
  let b = bcreate t.len in
  bblit t.data 0 b 0 t.len;
  {
    data = b;
    len = t.len
  }

let clone = copy

let set t x =
  if x < 0 then error "set";
  let pos = x lsr log_int_size and delta = x land int_size in
  let size = t.len in
  if pos >= size then begin
    let b = bcreate (pos+1) in
    bblit t.data 0 b 0 size;
    bfill b size (pos - size + 1) 0;
    t.len <- pos + 1;
    t.data <- b;
  end;
  bset t.data pos ((bget t.data pos) lor (1 lsl delta))

let unset t x =
  if x < 0 then error "unset";
  let pos = x lsr log_int_size and delta = x land int_size in
  if pos < t.len then
    bset t.data pos ((bget t.data pos) land (0xFF lxor (1 lsl delta)))

let toggle t x =
  if x < 0 then error "toggle";
  let pos = x lsr log_int_size and delta = x land int_size in
  let size = t.len in
  if pos >= size then begin
    let b = bcreate (pos+1) in
    bblit t.data 0 b 0 size;
    bfill b size (pos - size + 1) 0;
    t.len <- pos + 1;
    t.data <- b;
  end;
  bset t.data pos ((bget t.data pos) lxor (1 lsl delta))

let put t = function
  | true -> set t
  | false -> unset t

let is_set t x =
  if x < 0 then error "is_set";
  let pos = x lsr log_int_size and delta = x land int_size in
  let size = t.len in
  if pos < size then
  fast_bool (((bget t.data pos) lsr delta) land 1)
  else
  false


exception Break_int of int

(* Find highest set element or raise Not_found *)
let find_msb t =
  (* Find highest set bit in a byte.  Does not work with zero. *)
  let byte_msb b = 
    assert (b <> 0);
    let rec loop n = 
      if b land (1 lsl n) = 0 then
        loop (n-1)
      else n in
    loop 7 in
  let n = t.len - 1
  and buf = t.data in
  try 
    for i = n downto 0 do
      let byte = bget buf i in
      if byte <> 0 then raise (Break_int ((i lsl log_int_size)+(byte_msb byte)))
    done;
    raise Not_found
  with 
    Break_int n -> n
  | _ -> raise Not_found

let compare t1 t2 =
  let some_msb b = try Some (find_msb b) with Not_found -> None in
  match (some_msb t1, some_msb t2) with
    (None, Some _) -> -1 (* 0-y -> -1 *)
  | (Some _, None) -> 1  (* x-0 ->  1 *)
  | (None, None) -> 0    (* 0-0 ->  0 *)
  | (Some a, Some b) ->  (* x-y *)
      if a < b then -1
      else if a > b then 1
      else
        begin
          (* MSBs differ, we need to scan arrays until we find a
             difference *)
          let ndx = a lsr log_int_size in 
          assert (ndx < t1.len && ndx < t2.len);
          try
            for i = ndx downto 0 do
              let b1 = bget t1.data i 
              and b2 = bget t2.data i in
              if b1 <> b2 then raise (Break_int (compare b1 b2))
            done;
            0
          with
            Break_int res -> res
        end

let equals t1 t2 =
  compare t1 t2 = 0

let partial_count t x =
  let rec nbits x =
    if x = 0 then
      0
    else if fast_bool (x land 1) then
      1 + (nbits (x lsr 1))
    else
      nbits (x lsr 1)
  in
  let size = t.len in
  let pos = x lsr log_int_size and delta = x land int_size in
  let rec loop n acc =
    if n = size then
      acc
    else
      let x = bget t.data n in
      loop (n+1) (acc + nbits x)
  in
  if pos >= size then
    0
  else
    loop (pos+1) (nbits ((bget t.data pos) lsr delta))

let count t =
  partial_count t 0

(* Find the first set bit in the bit array *)
let find_first_set b n =
  (* TODO there are many ways to speed this up.  Lookup table would be
     one way to speed this up. *)
  let find_lsb b =
    assert (b <> 0);
    let rec loop n =
      if b land (1 lsl n) <> 0 then n else loop (n+1) in
    loop 0 in

  let buf = b.data in
  let rec find_bit byte_ndx bit_offs =
    if byte_ndx >= b.len then
      None
    else
      let byte = (bget buf byte_ndx) lsr bit_offs in
      if byte = 0 then
        find_bit (byte_ndx + 1) 0
      else
        Some ((find_lsb byte) + (byte_ndx lsl log_int_size) + bit_offs) in
  find_bit (n lsr log_int_size) (n land int_size)
      
let enum t =
  let rec make n =
    let cur = ref n in
    let rec next () =
      match find_first_set t !cur with
        Some elem ->
          cur := (elem+1);
          elem
      | None ->
          raise Enum.No_more_elements in
    Enum.make
      ~next
      ~count:(fun () -> partial_count t !cur)
      ~clone:(fun () -> make !cur)
  in
  make 0

let raw_create size = 
  let b = bcreate size in
  bfill b 0 size 0;
  { data = b; len = size }

let inter a b =
  let max_size = max a.len b.len in
  let d = raw_create max_size in
  let sl = min a.len b.len in
  let abuf = a.data
  and bbuf = b.data in
  (* Note: rest of the array is set to zero automatically *)
  for i = 0 to sl-1 do
    bset d.data i ((bget abuf i) land (bget bbuf i))
  done;
  d

(* Note: rest of the array is handled automatically correct, since we
   took a copy of the bigger set. *)
let union a b = 
  let d = if a.len > b.len then copy a else copy b in
  let sl = min a.len b.len in
  let abuf = a.data
  and bbuf = b.data in
  for i = 0 to sl-1 do
    bset d.data i ((bget abuf i) lor (bget bbuf i))
  done;
  d

let diff a b = 
  let maxlen = max a.len b.len in
  let buf = bcreate maxlen in
  bblit a.data 0 buf 0 a.len;
  let sl = min a.len b.len in
  let abuf = a.data
  and bbuf = b.data in
  for i = 0 to sl-1 do
    bset buf i ((bget abuf i) land (lnot (bget bbuf i)))
  done;
  { data = buf; len = maxlen }

let sym_diff a b = 
  let maxlen = max a.len b.len in
  let buf = bcreate maxlen in
  (* Copy larger (assumes missing bits are zero) *)
  bblit (if a.len > b.len then a.data else b.data) 0 buf 0 maxlen;
  let sl = min a.len b.len in
  let abuf = a.data
  and bbuf = b.data in
  for i = 0 to sl-1 do
    bset buf i ((bget abuf i) lxor (bget bbuf i))
  done;
  { data = buf; len = maxlen }

(* TODO the following set operations can be made faster if you do the
   set operation in-place instead of taking a copy.  But be careful
   when the sizes of the bitvector strings differ. *)
let intersect t t' =
  let d = inter t t' in
  t.data <- d.data;
  t.len <- d.len

let differentiate t t' =
  let d = diff t t' in
  t.data <- d.data;
  t.len <- d.len

let unite t t' =
  let d = union t t' in
  t.data <- d.data;
  t.len <- d.len

let differentiate_sym t t' =
  let d = sym_diff t t' in
  t.data <- d.data;
  t.len <- d.len