(* 'diskzip' command for intelligently compressing disk images.
   (C) Copyright 2007 Richard W.M. Jones, Red Hat Inc.
   http://libvirt.org/

   This program 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 *)

open Int63.Operators

type t = int63 * string array

(* NB: Sys.max_string_length is by no means a 'round' number. *)
let maxlen = Int63.of_int Sys.max_string_length

let create size =
  (* Get the size in bytes & round up. *)
  let sizebytes = (size +^ ~^7) >^> 3 in

  let n = Int63.to_int (sizebytes /^ maxlen) in
  let overflow = Int63.to_int (sizebytes %^ maxlen) in
  let zerochar = Char.chr 0 in
  let array =
    if overflow <> 0 then
      Array.init (n+1)
	(function
	 | i when i = n -> String.make overflow zerochar
	 | _ -> String.make Sys.max_string_length zerochar)
    else
      Array.init n (fun _ -> String.make Sys.max_string_length zerochar) in
  size, array

let mask7 = ~^7

let set (size, t) offset =
  if offset < Int63.zero || offset >= size then invalid_arg "bitmap";
  let offset_bytes, offset_bits =
    offset >^> 3,
    Int63.to_int (offset &^ mask7) in
  let offset_strs, offset_in_str =
    Int63.to_int (offset_bytes /^ maxlen),
    Int63.to_int (offset_bytes %^ maxlen) in
  let c = Char.code t.(offset_strs).[offset_in_str] in
  t.(offset_strs).[offset_in_str] <- Char.chr (c lor (1 lsl offset_bits))

let clear (size, t) offset =
  if offset < Int63.zero || offset >= size then invalid_arg "bitmap";
  let offset_bytes, offset_bits =
    offset >^> 3,
    Int63.to_int (offset &^ mask7) in
  let offset_strs, offset_in_str =
    Int63.to_int (offset_bytes /^ maxlen),
    Int63.to_int (offset_bytes %^ maxlen) in
  let c = Char.code t.(offset_strs).[offset_in_str] in
  t.(offset_strs).[offset_in_str] <-
    Char.chr (c land (lnot (1 lsl offset_bits)))

let get (size, t) offset =
  if offset < Int63.zero || offset >= size then invalid_arg "bitmap";
  let offset_bytes, offset_bits =
    offset >^> 3,
    Int63.to_int (offset &^ mask7) in
  let offset_strs, offset_in_str =
    Int63.to_int (offset_bytes /^ maxlen),
    Int63.to_int (offset_bytes %^ maxlen) in
  let c = Char.code t.(offset_strs).[offset_in_str] in
  c land (1 lsl offset_bits) <> 0

let set_bool t offset v =
  (if v then set else clear) t offset

let set_int t offset v =
  (if v <> 0 then set else clear) t offset

let iter_set f (size, t) =
  let n = ref Int63.zero in
  Array.iter (
    fun str ->
      let m = ref 0x80 in
      for i = 0 to String.length str - 1 do
	let c = Char.code str.[i] in
	for j = 0 to 7 do
	  let n' = !n in
	  if n' < size then (
	    let b = c land !m <> 0 in
	    let b = f n' b in
	    str.[i] <- Char.chr (if b then c lor !m else c land lnot !m)
	  );
	  n := Int63.succ n';
	  m := !m lsr 1
	done
      done
  ) t

let iter f t =
  let f i b = f i b; b in
  iter_set f t