(************************************************************************)
(* This file is part of SKS.  SKS is free software; you can
   redistribute it and/or modify it under the terms of the GNU General
   Public License as published by the Free Software Foundation; either
   version 2 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
   General Public License for more details.

   You should have received a copy of the GNU 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 *)
(***********************************************************************)

open StdLabels
open MoreLabels

module Unix=UnixLabels

exception Error of string
exception LengthError of string

let width = 8

type t = { a: string;
	   bitlength: int;
	 }

let bytelength bits = 
  bits / width + (if bits mod width = 0 then 0 else 1)

let create bits = 
  let bytes = bytelength bits
  in
  { a = String.create bytes;
    bitlength = bits;
  }

let get ba bit = 
  let byte_pos = bit / width 
  and bit_pos = bit mod width in
  let intval = int_of_char (String.get ba.a byte_pos) in
  (intval lsr (width - bit_pos - 1)) land 1

let lget ba bit = get ba bit = 1

let flip ba bit = 
  let byte_pos = bit / width 
  and bit_pos = bit mod width in
  let intval = int_of_char (String.get ba.a byte_pos) in
  let new_char = char_of_int ((1 lsl (width - bit_pos - 1)) lxor intval)
  in
  String.set ba.a byte_pos new_char

let set ba bit = 
  let byte_pos = bit / width 
  and bit_pos = bit mod width in
  let intval = int_of_char (String.get ba.a byte_pos) in
  let new_char = char_of_int ((1 lsl (width - bit_pos - 1)) lor intval)
  in
  String.set ba.a byte_pos new_char

let unset ba bit = 
  let byte_pos = bit / width 
  and bit_pos = bit mod width in
  let intval = int_of_char (String.get ba.a byte_pos) in
  let new_char = char_of_int ((lnot (1 lsl (width - bit_pos - 1))) 
			      land intval)
  in
  String.set ba.a byte_pos new_char

let setval ba bit bool =
  if bool then set ba bit else unset ba bit

(************************************************************)
(* Printing and Conversions *********************************)
(************************************************************)

let print ba =
  for i = 0 to ba.bitlength - 1 do
    if get ba i = 0 
    then print_string "0" 
    else print_string "1"
  done

let hexprint ba = 
  print_string (Utils.hexstring ba.a)

let to_bool_array ba = 
  Array.init ~f:(fun i -> lget ba i) ba.bitlength

let to_string ba = 
  let string = String.create ba.bitlength in
  for i = 0 to ba.bitlength -1 do
    if get ba i = 0 then string.[i] <- '0' else string.[i] <- '1'
  done;
  string

let to_bytes ba = 
  let lastbit = (bytelength ba.bitlength)*width - 1 in
  for i = ba.bitlength to lastbit do 
    unset ba i 
  done;
  String.sub ~pos:0 ~len:(bytelength ba.bitlength) ba.a

let of_bytes string bitlength = 
  { bitlength = bitlength;
    a = String.copy string;
  }

let of_byte b =
  { bitlength = width;
    a = String.make 1 (char_of_int (b land 0xFF));
  }

let of_bytes_all string = 
  { bitlength = (String.length string) * width;
    a = String.copy string;
  }

let of_int i = 
  { bitlength = width * 4;
    a = Utils.bstring_of_int i;
  }

let of_bytes_nocopy string bitlength = 
  { bitlength = bitlength;
    a = string;
  }

let of_bytes_all_nocopy string = 
  { bitlength = (String.length string) * width;
    a = string;
  }

let to_bytes_nocopy ba = 
  let lastbit = (bytelength ba.bitlength)*8 - 1 in
  for i = ba.bitlength to lastbit do 
    unset ba i 
  done;
  ba.a

(************************************************************)
(************************************************************)
(************************************************************)

let copy ba = { ba with a = String.copy ba.a }

(** returns a copy of bitstring copied into a new bitstring of a new length.
  No guarantees are made as to the contents of the remainder of the bitstring
  if the bitstring length is extended.
 *)
let copy_len ba bitlength = 
  let bytes = bytelength bitlength in
  let str = String.create bytes in
  String.blit ~src:ba.a ~src_pos:0 
    ~dst:str ~dst_pos:0 ~len:(String.length ba.a);
  { ba with a = str; bitlength = bitlength }

(********************************************************************)
(***  Shifting  *****************************************************)
(********************************************************************)

let shift_pair_left c1 c2 bits= 
  let i1 = int_of_char c1
  and i2 = int_of_char c2  in
  let shifted_int = 
    (i1 lsl bits) lor (i2 lsr (width - bits))
  in 
  char_of_int (shifted_int land 0xFF)

let shift_pair_right c1 c2 bits = 
  let i1 = int_of_char c1
  and i2 = int_of_char c2 in
  let shifted_int = 
    (i1 lsl (width - bits)) lor (i2 lsr bits)
  in 
  char_of_int (shifted_int land 0xFF)
    
(**********************************)

let shift_left_small ba bits = 
  if bits > 0 then
    let bytes = bytelength ba.bitlength in
    for i = 0 to bytes-2 do
      ba.a.[i] <- shift_pair_left ba.a.[i] ba.a.[i+1] bits
    done;
    ba.a.[bytes-1] <- shift_pair_left ba.a.[bytes-1] '\000' bits

let shift_right_small ba bits = 
  if bits > 0 then
    let bytes = bytelength ba.bitlength in
    for i = bytes-1 downto 1 do
      ba.a.[i] <- shift_pair_right ba.a.[i-1] ba.a.[i] bits
    done;
    ba.a.[0] <-  shift_pair_right '\000' ba.a.[0] bits

(**********************************)

let rec shift_left ba bits =
  if bits < 0 then
    shift_right ba (-bits)
  else
  let bytelength = bytelength ba.bitlength 
  and bytes = bits / width
  and bits = bits mod width in
  if bytes > 0 
  then 
    begin
      for i = 0 to bytelength - 1 - bytes do 
	ba.a.[i] <- ba.a.[i+bytes];
      done;
      for i = bytelength - bytes to bytelength - 1 do
	ba.a.[i] <- '\000'
      done
    end;
  shift_left_small ba bits

and shift_right ba bits =
  if bits < 0 then
    shift_left ba (-bits)
  else
    let bytelength = bytelength ba.bitlength
    and bytes = bits / width
    and bits = bits mod width in
    if bytes > 0 
    then 
      begin
	for i = bytelength - 1 downto bytes do 
	  ba.a.[i] <- ba.a.[i-bytes];
	done;
	for i = bytes - 1 downto 0 do
	  ba.a.[i] <- '\000'
	done
      end;
    shift_right_small ba bits
    
let num_bits ba = ba.bitlength
let num_bytes ba = bytelength ba.bitlength

(********************************************************************)
(********************************************************************)
(********************************************************************)

let rmasks = 
  Array.init width ~f:(fun i -> 0xFF lsl (width - i))

(* Later, extend to have optional initial-position arguments *)
let blit ~src ~dst ~len = 
  (* these tests are probably redundant, since they'll cause 
     exceptions deeper in.  OCaml's lousy traceback features, however, make
     it somewhat useful to have these here. *)
  if len < 0 
  then raise (Invalid_argument "Bitstring.blit: negative len");
  if dst.bitlength < len 
  then raise (Invalid_argument "Bitstring.blit: dst too short");
  if src.bitlength < len 
  then raise (Invalid_argument "Bitstring.blit: src too short");
  let bytelen = len / width
  and bitlen = len mod width in
  String.blit 
    ~src:src.a ~src_pos:0 
    ~dst:dst.a ~dst_pos:0 ~len:bytelen;
  if bitlen > 0 then
    let srcval = int_of_char (String.get src.a bytelen)
    and dstval = int_of_char (String.get dst.a bytelen) in
    let newdst = (rmasks.(bitlen) land srcval) lor 
		 ((lnot rmasks.(bitlen)) land dstval)
    in
    dst.a.[bytelen] <- char_of_int newdst

  
(* let full_blit ~src ~src_pos ~dst ~dst_pos ~len =  *)
  

let zero_out bs = 
  String.fill bs.a ~pos:0 ~len:(String.length bs.a) '\000'

(*
let extract bs ~pos ~len = 
  let first_bit = pos % 8
  let first_byte = pos / 8 in
  let last_byte = (pos + len) / 8 + 
		  (if (pos + len) % 8 > 0 then 1 else 0) in
  let byte_len =  last_byte - first_byte + 1 in
  let newbs = Bitstring.create len in
  String.blit 
    ~src:bs.a ~src_pos:src_first_byte 
    ~dst:newbs.a ~dst_pos:0 ~len:byte_len;
  shift_left newbs first_bit;
*)


(*
let concat bs1 bs2 = 
  let newbs = create (bs1.bits + bs2.bits) in
  blit ~src:bs1 ~dst:newbs ~len:(bs1.bits);
*)