(*

 * Unzip - inflate format decompression algorithm

 * Copyright (C) 2004 Nicolas Cannasse

 * Compliant with RFC 1950 and 1951

 *

 * 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 huffman =

  | Found of int

  | NeedBit of huffman * huffman

  | NeedBits of int * huffman array

type adler32 = {

  mutable a1 : int;

  mutable a2 : int;

}

type window = {

  mutable wbuffer : Bytes.t;

  mutable wpos : int;

  wcrc : adler32;

}

type state =

  | Head

  | Block

  | CData

  | Flat

  | Crc

  | Dist

  | DistOne

  | Done

type t = {

  mutable znbits : int;

  mutable zbits : int;

  mutable zstate : state;

  mutable zfinal : bool;

  mutable zhuffman : huffman;

  mutable zhuffdist : huffman option;

  mutable zlen : int;

  mutable zdist : int;

  mutable zneeded : int;

  mutable zoutput : Bytes.t;

  mutable zoutpos : int;

  zinput : IO.input;

  zlengths : int array;

  zwindow : window;

}

type error_msg =

  | Invalid_huffman

  | Invalid_data

  | Invalid_crc

  | Truncated_data

  | Unsupported_dictionary

exception Error of error_msg

let error msg = raise (Error msg)

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

(* HUFFMAN TREES *)

let rec tree_depth = function

  | Found _ -> 0

  | NeedBits _ -> assert false

  | NeedBit (a,b) ->

    1 + min (tree_depth a) (tree_depth b)

let rec tree_compress t =

  match tree_depth t with

  | 0 -> t

  | 1 ->

    (match t with

    | NeedBit (a,b) -> NeedBit (tree_compress a,tree_compress b)

    | _ -> assert false)

  | d ->

    let size = 1 lsl d in

    let tbl = Array.make size (Found (-1)) in

    tree_walk tbl 0 0 d t;

    NeedBits (d,tbl)

and tree_walk tbl p cd d = function

  | NeedBit (a,b) when d > 0 ->

    tree_walk tbl p (cd + 1) (d-1) a;

    tree_walk tbl (p lor (1 lsl cd)) (cd + 1) (d-1) b;

  | t ->

    Array.set tbl p (tree_compress t)

let make_huffman lengths pos nlengths maxbits =

  let counts = Array.make maxbits 0 in

  for i = 0 to nlengths - 1 do

    let p = Array.unsafe_get lengths (i + pos) in

    if p >= maxbits then error Invalid_huffman;

    Array.unsafe_set counts p (Array.unsafe_get counts p + 1);

  done;

  let code = ref 0 in

  let tmp = Array.make maxbits 0 in

  for i = 1 to maxbits - 2 do

    code := (!code + Array.unsafe_get counts i) lsl 1;

    Array.unsafe_set tmp i !code;

  done;

  let bits = Hashtbl.create 0 in

  for i = 0 to nlengths - 1 do

    let l = Array.unsafe_get lengths (i + pos) in

    if l <> 0 then begin

      let n = Array.unsafe_get tmp (l - 1) in

      Array.unsafe_set tmp (l - 1) (n + 1);

      Hashtbl.add bits (n,l) i;

    end;

  done;

  let rec tree_make v l =

    if l > maxbits then error Invalid_huffman;

    try

      Found (Hashtbl.find bits (v,l))

    with

      Not_found ->

        NeedBit (tree_make (v lsl 1) (l + 1) , tree_make (v lsl 1 lor 1) (l + 1))

  in

  tree_compress (NeedBit (tree_make 0 1 , tree_make 1 1))

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

(* ADLER32 (CRC) *)

let adler32_create() = {

  a1 = 1;

  a2 = 0;

}

let adler32_update a s p l =

  let p = ref p in

  for i = 0 to l - 1 do

    let c = int_of_char (Bytes.unsafe_get s !p) in

    a.a1 <- (a.a1 + c) mod 65521;

    a.a2 <- (a.a2 + a.a1) mod 65521;

    incr p;

  done

let adler32_read ch =

  let a2a = IO.read_byte ch in

  let a2b = IO.read_byte ch in

  let a1a = IO.read_byte ch in

  let a1b = IO.read_byte ch in

  {

    a1 = (a1a lsl 8) lor a1b;

    a2 = (a2a lsl 8) lor a2b;

  }

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

(* WINDOW *)

let window_size = 1 lsl 15

let buffer_size = 1 lsl 16

let window_create size = {

    wbuffer = Bytes.create buffer_size;

    wpos = 0;

    wcrc = adler32_create()

  }

let window_slide w = 

  adler32_update w.wcrc w.wbuffer 0 window_size;

  let b = Bytes.create buffer_size in

  w.wpos <- w.wpos - window_size;

  Bytes.unsafe_blit w.wbuffer window_size b 0 w.wpos;

  w.wbuffer <- b

let window_add_bytes w s p len =

  if w.wpos + len > buffer_size then window_slide w;

  Bytes.unsafe_blit s p w.wbuffer w.wpos len;

  w.wpos <- w.wpos + len

let window_add_char w c =

  if w.wpos = buffer_size then window_slide w;

  Bytes.unsafe_set w.wbuffer w.wpos c;

  w.wpos <- w.wpos + 1

let window_get_last_char w =

  Bytes.unsafe_get w.wbuffer (w.wpos - 1)

let window_available w =

  w.wpos

let window_checksum w =

  adler32_update w.wcrc w.wbuffer 0 w.wpos;

  w.wcrc

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

let len_extra_bits_tbl = [|0;0;0;0;0;0;0;0;1;1;1;1;2;2;2;2;3;3;3;3;4;4;4;4;5;5;5;5;0;-1;-1|]

let len_base_val_tbl = [|3;4;5;6;7;8;9;10;11;13;15;17;19;23;27;31;35;43;51;59;67;83;99;115;131;163;195;227;258|]

let dist_extra_bits_tbl = [|0;0;0;0;1;1;2;2;3;3;4;4;5;5;6;6;7;7;8;8;9;9;10;10;11;11;12;12;13;13;-1;-1|]

let dist_base_val_tbl = [|1;2;3;4;5;7;9;13;17;25;33;49;65;97;129;193;257;385;513;769;1025;1537;2049;3073;4097;6145;8193;12289;16385;24577|]

let code_lengths_pos = [|16;17;18;0;8;7;9;6;10;5;11;4;12;3;13;2;14;1;15|]

let fixed_huffman = make_huffman (Array.init 288 (fun n ->

                  if n <= 143 then 8

                  else if n <= 255 then 9

                  else if n <= 279 then 7

                  else 8

                )) 0 288 10

let get_bits z n =

  while z.znbits < n do

    z.zbits <- z.zbits lor ((IO.read_byte z.zinput) lsl z.znbits);

    z.znbits <- z.znbits + 8;

  done;

  let b = z.zbits land (1 lsl n - 1) in

  z.znbits <- z.znbits - n;

  z.zbits <- z.zbits lsr n;

  b

let get_bit z =

  if z.znbits = 0 then begin

    z.znbits <- 8;

    z.zbits <- IO.read_byte z.zinput;

  end;

  let b = z.zbits land 1 = 1 in

  z.znbits <- z.znbits - 1;

  z.zbits <- z.zbits lsr 1;

  b

let rec get_rev_bits z n =

  if n = 0 then

    0

  else if get_bit z then

    (1 lsl (n - 1)) lor (get_rev_bits z (n-1))

  else

    get_rev_bits z (n-1)

let reset_bits z =

  z.zbits <- 0;

  z.znbits <- 0

let add_bytes z s p l =

  window_add_bytes z.zwindow s p l;

  Bytes.unsafe_blit s p z.zoutput z.zoutpos l;

  z.zneeded <- z.zneeded - l;

  z.zoutpos <- z.zoutpos + l

let add_char z c =

  window_add_char z.zwindow c;

  Bytes.unsafe_set z.zoutput z.zoutpos c;

  z.zneeded <- z.zneeded - 1;

  z.zoutpos <- z.zoutpos + 1

let add_dist_one z n =

  let c = window_get_last_char z.zwindow in

  let s = Bytes.make n c in

  add_bytes z s 0 n

let add_dist z d l =

  add_bytes z z.zwindow.wbuffer (z.zwindow.wpos - d) l

let rec apply_huffman z = function

  | Found n -> n

  | NeedBit (a,b) -> apply_huffman z (if get_bit z then b else a)

  | NeedBits (n,t) -> apply_huffman z (Array.unsafe_get t (get_bits z n))

let inflate_lengths z a max =

  let i = ref 0 in

  let prev = ref 0 in

  while !i < max do

    match apply_huffman z z.zhuffman with

    | n when n <= 15 ->

      prev := n;

      Array.unsafe_set a !i n;

      incr i

    | 16 ->

      let n = 3 + get_bits z 2 in

      if !i + n > max then error Invalid_data;

      for k = 0 to n - 1 do

        Array.unsafe_set a !i !prev;

        incr i;

      done;

    | 17 ->

      let n = 3 + get_bits z 3 in

      i := !i + n;

      if !i > max then error Invalid_data;

    | 18 ->

      let n = 11 + get_bits z 7 in

      i := !i + n;

      if !i > max then error Invalid_data;

    | _ ->

      error Invalid_data

  done

let rec inflate_loop z =

  match z.zstate with

  | Head ->

    let cmf = IO.read_byte z.zinput in

    let cm = cmf land 15 in

    let cinfo = cmf lsr 4 in

    if cm <> 8 || cinfo <> 7 then error Invalid_data;

    let flg = IO.read_byte z.zinput in

    (*let fcheck = flg land 31 in*)

    let fdict = flg land 32 <> 0 in

    (*let flevel = flg lsr 6 in*)

    if (cmf lsl 8 + flg) mod 31 <> 0 then error Invalid_data;

    if fdict then error Unsupported_dictionary;

    z.zstate <- Block;

    inflate_loop z

  | Crc ->

    let calc = window_checksum z.zwindow in

    let crc = adler32_read z.zinput in

    if calc <> crc then error Invalid_crc;

    z.zstate <- Done;

    inflate_loop z

  | Done ->

    ()

  | Block ->

    z.zfinal <- get_bit z;

    let btype = get_bits z 2 in

    (match btype with

    | 0 -> (* no compression *)

      z.zlen <- IO.read_ui16 z.zinput;

      let nlen = IO.read_ui16 z.zinput in

      if nlen <> 0xffff - z.zlen then error Invalid_data;

      z.zstate <- Flat;

      inflate_loop z;

      reset_bits z

    | 1 -> (* fixed Huffman *)

      z.zhuffman <- fixed_huffman;

      z.zhuffdist <- None;

      z.zstate <- CData;

      inflate_loop z

    | 2 -> (* dynamic Huffman *)

      let hlit = get_bits z 5 + 257 in

      let hdist = get_bits z 5 + 1 in

      let hclen = get_bits z 4 + 4 in

      for i = 0 to hclen - 1 do

        Array.unsafe_set z.zlengths (Array.unsafe_get code_lengths_pos i) (get_bits z 3);

      done;

      for i = hclen to 18 do

        Array.unsafe_set z.zlengths (Array.unsafe_get code_lengths_pos i) 0;

      done;

      z.zhuffman <- make_huffman z.zlengths 0 19 8;

      let lengths = Array.make (hlit + hdist) 0 in

      inflate_lengths z lengths (hlit + hdist);

      z.zhuffdist <- Some (make_huffman lengths hlit hdist 16);

      z.zhuffman <- make_huffman lengths 0 hlit 16;      

      z.zstate <- CData;

      inflate_loop z

    | _ ->

      error Invalid_data)

  | Flat ->

    let rlen = min z.zlen z.zneeded in

    let str = IO.nread z.zinput rlen in

    let len = Bytes.length str in

    z.zlen <- z.zlen - len;

    add_bytes z str 0 len;

    if z.zlen = 0 then z.zstate <- (if z.zfinal then Crc else Block);

    if z.zneeded > 0 then inflate_loop z

  | DistOne ->

    let len = min z.zlen z.zneeded in

    add_dist_one z len;

    z.zlen <- z.zlen - len;

    if z.zlen = 0 then z.zstate <- CData;

    if z.zneeded > 0 then inflate_loop z

  | Dist ->

    while z.zlen > 0 && z.zneeded > 0 do

      let len = min z.zneeded (min z.zlen z.zdist) in

      add_dist z z.zdist len;

      z.zlen <- z.zlen - len;

    done;

    if z.zlen = 0 then z.zstate <- CData;

    if z.zneeded > 0 then inflate_loop z

  | CData ->

    match apply_huffman z z.zhuffman with

    | n when n < 256 ->

      add_char z (Char.unsafe_chr n);

      if z.zneeded > 0 then inflate_loop z

    | 256 ->

      z.zstate <- if z.zfinal then Crc else Block;

      inflate_loop z

    | n ->

      let n = n - 257 in

      let extra_bits = Array.unsafe_get len_extra_bits_tbl n in

      if extra_bits = -1 then error Invalid_data;

      z.zlen <- (Array.unsafe_get len_base_val_tbl n) + (get_bits z extra_bits);

      let dist_code = (match z.zhuffdist with None -> get_rev_bits z 5 | Some h -> apply_huffman z h) in

      let extra_bits = Array.unsafe_get dist_extra_bits_tbl dist_code in

      if extra_bits = -1 then error Invalid_data;

      z.zdist <- (Array.unsafe_get dist_base_val_tbl dist_code) + (get_bits z extra_bits);

      if z.zdist > window_available z.zwindow then error Invalid_data;

      z.zstate <- (if z.zdist = 1 then DistOne else Dist);

      inflate_loop z

let inflate_data z s pos len =

  if pos < 0 || len < 0 || pos + len > Bytes.length s then invalid_arg "inflate_data";

  z.zneeded <- len;

  z.zoutpos <- pos;

  z.zoutput <- s;

  try

    if len > 0 then inflate_loop z;

    len - z.zneeded

  with

    IO.No_more_input -> error Truncated_data

let inflate_init ?(header=true) ch = 

  {

    zfinal = false;

    zhuffman = fixed_huffman;

    zhuffdist = None;

    zlen = 0;

    zdist = 0;

    zstate = (if header then Head else Block);

    zinput = ch;

    zbits = 0;

    znbits = 0;

    zneeded = 0;

    zoutput = Bytes.empty;

    zoutpos = 0;

    zlengths = Array.make 19 (-1);

    zwindow = window_create (1 lsl 15)

  }

let inflate ?(header=true) ch =

  let z = inflate_init ~header ch in

  let s = Bytes.create 1 in

  IO.create_in

    ~read:(fun() ->

      let l = inflate_data z s 0 1 in

      if l = 1 then Bytes.unsafe_get s 0 else raise IO.No_more_input

    )

    ~input:(fun s p l ->

      let n = inflate_data z s p l in

      if n = 0 then raise IO.No_more_input;

      n

    )

    ~close:(fun () ->

      IO.close_in ch

    )