Source file zl.ml

let io_buffer_size = 65536
let kstrf k fmt = Format.kasprintf k fmt
let invalid_arg fmt = Format.kasprintf invalid_arg fmt

module Bigarray = Bigarray_compat
type bigstring =
  (char, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t
type window = De.window

let bigstring_create l = Bigarray.Array1.create Bigarray.char Bigarray.c_layout l
let bigstring_empty = Bigarray.Array1.create Bigarray.char Bigarray.c_layout 0
let bigstring_length x = Bigarray.Array1.dim x [@@inline]

external unsafe_get_uint8 : bigstring -> int -> int = "%caml_ba_ref_1"
external unsafe_get_uint16 : bigstring -> int -> int = "%caml_bigstring_get16"
external unsafe_get_uint32 : bigstring -> int -> int32 = "%caml_bigstring_get32"

external unsafe_set_uint8 : bigstring -> int -> int -> unit = "%caml_ba_set_1"
external unsafe_set_uint16 : bigstring -> int -> int -> unit = "%caml_bigstring_set16"
external unsafe_set_uint32 : bigstring -> int -> int32 -> unit = "%caml_bigstring_set32"

external swap16 : int -> int = "%bswap16"
external swap32 : int32 -> int32 = "caml_int32_bswap"

let string_unsafe_get_uint8 : string -> int -> int =
  fun buf off -> Char.code (String.get buf off)

external string_unsafe_get_uint32 : string -> int -> int32 = "%caml_string_get32"

let bytes_unsafe_get_uint8 : bytes -> int -> int =
  fun buf off -> Char.code (Bytes.get buf off)

external bytes_unsafe_get_uint32 : bytes -> int -> int32 = "%caml_bytes_get32"

let bytes_unsafe_set_uint8 : bytes -> int -> int -> unit =
  fun buf off v -> Bytes.set buf off (Char.unsafe_chr (v land 0xff))

external bytes_unsafe_set_uint32 : bytes -> int -> int32 -> unit = "%caml_bytes_set32"

let input_bigstring ic buf off len =
  let tmp = Bytes.create len in
  let res = input ic tmp 0 len in

  let len0 = res land 3 in
  let len1 = res asr 2 in

  for i = 0 to len1 - 1
  do
    let i = i * 4 in
    let v = bytes_unsafe_get_uint32 tmp i in
    unsafe_set_uint32 buf (off + i) v
  done ;

  for i = 0 to len0 - 1
  do
    let i = len1 * 4 + i in
    let v = bytes_unsafe_get_uint8 tmp i in
    unsafe_set_uint8 buf (off + i) v
  done ; res

let bigstring_to_string v =
  let len = bigstring_length v in
  let res = Bytes.create len in
  let len0 = len land 3 in
  let len1 = len asr 2 in

  for i = 0 to len1 - 1
  do
    let i = i * 4 in
    let v = unsafe_get_uint32 v i in
    bytes_unsafe_set_uint32 res i v
  done ;

  for i = 0 to len0 - 1
  do
    let i = len1 * 4 + i in
    let v = unsafe_get_uint8 v i in
    bytes_unsafe_set_uint8 res i v
  done ;

  Bytes.unsafe_to_string res

let output_bigstring oc buf off len =
  (* XXX(dinosaure): stupidly slow! *)
  let v = Bigarray.Array1.sub buf off len in
  let v = bigstring_to_string v in
  output_string oc v

let bigstring_of_string v =
  let len = String.length v in
  let res = bigstring_create len in
  let len0 = len land 3 in
  let len1 = len asr 2 in

  for i = 0 to len1 - 1
  do
    let i = i * 4 in
    let v = string_unsafe_get_uint32 v i in
    unsafe_set_uint32 res i v
  done ;

  for i = 0 to len0 - 1
  do
    let i = len1 * 4 + i in
    let v = string_unsafe_get_uint8 v i in
    unsafe_set_uint8 res i v
  done ; res

let unsafe_get_uint32_be =
  if Sys.big_endian
  then fun buf off -> unsafe_get_uint32 buf off
  else fun buf off -> swap32 (unsafe_get_uint32 buf off)

let unsafe_set_uint32_be =
  if Sys.big_endian
  then fun buf off v -> unsafe_set_uint32 buf off v
  else fun buf off v -> unsafe_set_uint32 buf off (swap32 v)

let unsafe_set_uint16_be =
  if Sys.big_endian
  then fun buf off v -> unsafe_set_uint16 buf off v
  else fun buf off v -> unsafe_set_uint16 buf off (swap16 v)

let invalid_bounds off len = invalid_arg "Out of bounds (off: %d, len: %d)" off len

let _deflated = 8 (* Compression method *)

module Inf = struct
  type src = [ `Channel of in_channel | `String of string | `Manual ]

  (* XXX(dinosaure): immutable style. *)
  type decoder =
    { src : De.Inf.src
    ; i : bigstring
    ; i_pos : int
    ; i_len : int
    ; f : bool
    ; wr : int
    ; hd : int
    ; dd : dd
    ; fdict : bool
    ; flevel : int
    ; cinfo : int
    ; allocate : int -> De.window
    ; t : bigstring
    ; t_need : int
    ; t_len : int
    ; k : decoder -> signal }
  and dd =
    | Dd of { state : De.Inf.decoder
            ; window : De.window
            ; o : De.bigstring }
    | Hd of { o : De.bigstring }
  and signal =
    [ `Await of decoder
    | `Flush of decoder
    | `End of decoder
    | `Malformed of string ]

  let malformedf fmt = kstrf (fun s -> `Malformed s) fmt

  let err_unexpected_end_of_input _ =
    malformedf "Unexpected end of input"

  let err_invalid_checksum has expect _ =
    malformedf "Invalid checksum (expect:%04lx, has:%04lx)" expect (Optint.to_int32 has)

  let err_invalid_header _ =
    malformedf "Invalid Zlib header"

  (* remaining bytes to read [d.i] *)
  let i_rem d = d.i_len - d.i_pos + 1
  [@@inline]

  (* End of input [eoi] is signalled by [d.i_pos = 0] and [d.i_len = min_int]
     which implies [i_rem d < 0] is [true]. *)

  let eoi d =
    { d with i= bigstring_empty
           ; i_pos= 0
           ; i_len= min_int }

  let refill k d = match d.dd, d.src with
    | Dd { state; _ }, `String _ ->
      De.Inf.src state bigstring_empty 0 0 ;
      k (eoi d)
    | Dd { state; _ }, `Channel ic ->
      let res = input_bigstring ic d.i 0 (bigstring_length d.i) in
      De.Inf.src state d.i 0 res ; k d
    | (Dd _ | Hd _), `Manual ->
      `Await { d with k }
    | Hd _, `String _ ->
      k (eoi d)
    | Hd _, `Channel ic ->
      let res = input_bigstring ic d.i 0 (bigstring_length d.i) in
      if res == 0
      then k (eoi d)
      else k { d with i_pos= 0; i_len= res - 1 }

  let flush k d = `Flush { d with k }

  let blit src ~src_off dst ~dst_off ~len =
    let a = Bigarray.Array1.sub src src_off len in
    let b = Bigarray.Array1.sub dst dst_off len in
    Bigarray.Array1.blit a b

  let rec t_fill k d =
    let blit d len =
      blit d.i ~src_off:d.i_pos d.t ~dst_off:d.t_len ~len ;
      { d with i_pos= d.i_pos + len
             ; t_len= d.t_len + len } in
    let rem = i_rem d in
    if rem < 0 then malformedf "Unexpected end of input"
    else
      let need = d.t_need - d.t_len in
      if rem < need
      then let d = blit d rem in refill (t_fill k) d
      else let d = blit d need in k { d with t_need= 0 }

  let t_need n d = { d with t_need= n }

  let checksum d =
    let k d = match d.dd with
      | Dd { state; _ } ->
        let a = De.Inf.checksum state in
        let b = unsafe_get_uint32_be d.t 0 in

        if Optint.to_int32 a = b (* FIXME: Optint.equal a (Optint.of_int32 b) bugs! *)
        then `End d
        else err_invalid_checksum a b d
      | Hd _ -> assert false in

    t_fill k (t_need 4 d)

  let rec header d =
    let k d =
      let[@warning "-8"] Hd { o; } = d.dd in
      let cmf = unsafe_get_uint16 d.t 0 in
      let cm = cmf land 0b1111 in
      let cinfo = (cmf lsr 4) land 0b1111 in
      let flg = cmf lsr 8 in
      let fdict = (flg lsr 5) land 0b1 in
      let flevel = (flg lsr 6) land 0b11 in
      let window = d.allocate (cinfo + 8) in
      let state = De.Inf.decoder `Manual ~o ~w:window in
      let dd = Dd { state; window; o; } in
      if ((cmf land 0xff) lsl 8 + (cmf lsr 8)) mod 31 != 0
      || cm != _deflated
      then err_invalid_header d
      else
        ( if i_rem d > 0 then De.Inf.src state d.i d.i_pos (i_rem d)
        ; decode { d with hd= unsafe_get_uint16 d.t 0
                        ; k= decode
                        ; dd
                        ; t_need= 0; t_len= 0
                        ; fdict= fdict == 1; flevel; cinfo } ) in
    if i_rem d >= 2
    then ( unsafe_set_uint16 d.t 0 (unsafe_get_uint16 d.i d.i_pos)
         ; k { d with i_pos= d.i_pos + 2 } )
    else ( if i_rem d < 0
           then err_unexpected_end_of_input d
           else if i_rem d == 0 then refill header d
           else t_fill k (t_need 2 d) )

  and decode d = match d.dd with
    | Hd _ -> header d
    | Dd { state; o; _ } ->
      match De.Inf.decode state with
      | `Flush ->
        if d.f
        then flush decode d
        else
          let len = bigstring_length o - De.Inf.dst_rem state in
          flush decode { d with wr= d.wr + len; f= true; }
      | `Await ->
        let len = i_rem d - De.Inf.src_rem state in
        refill decode { d with i_pos= d.i_pos + len }
      | `End ->
        if d.f
        then flush decode d
        else
          let len = bigstring_length o - De.Inf.dst_rem state in
          if len > 0
          then
            flush decode
              { d with i_pos= d.i_pos + (i_rem d - De.Inf.src_rem state)
                     ; wr= d.wr + len
                     ; f= true }
          else checksum
              { d with i_pos= d.i_pos + (i_rem d - De.Inf.src_rem state) }
      | `Malformed err -> `Malformed err

  let src d s j l =
    if (j < 0 || l < 0 || j + l > bigstring_length s)
    then invalid_bounds j l ;
    let d =
      if (l == 0)
      then eoi d
      else
        { d with i= s
               ; i_pos= j
               ; i_len= j + l - 1 } in
    match d.dd with
    | Dd { state; _ } ->
      De.Inf.src state s j l ; d
    | Hd _ -> d

  let flush d = match d.dd with
    | Hd _ -> { d with f= false }
    | Dd { state; _ } ->
      De.Inf.flush state ; { d with f= false }

  let dst_rem d = match d.dd with
    | Hd _ -> 0
    | Dd { state; _ } -> De.Inf.dst_rem state

  let src_rem d = i_rem d

  let write { wr; _ } = wr

  let decoder src ~o ~allocate =
    let i, i_pos, i_len = match src with
      | `Manual -> bigstring_empty, 1, 0
      | `String x -> bigstring_of_string x, 0, String.length x - 1
      | `Channel _ -> bigstring_create io_buffer_size, 1, 0 in
    { i; i_pos; i_len
    ; src
    ; f= false
    ; wr= 0
    ; hd= 0
    ; dd= Hd { o; }
    ; fdict= false
    ; flevel= 2
    ; cinfo= 8
    ; allocate
    ; t= bigstring_create 4
    ; t_need= 0
    ; t_len= 0
    ; k= decode }

  let reset d =
    let i, i_pos, i_len = match d.src with
      | `Manual -> bigstring_empty, 1, 0
      | `String x -> bigstring_of_string x, 0, String.length x - 1
      | `Channel _ -> bigstring_create io_buffer_size, 1, 0 in
    let o = match d.dd with
      | Hd { o; } -> o
      | Dd { o; _ } -> o in
    { i; i_pos; i_len
    ; src= d.src
    ; f= false
    ; wr= 0
    ; hd= 0
    ; dd= Hd { o; }
    ; fdict= false
    ; flevel= 2
    ; cinfo= 8
    ; allocate= d.allocate
    ; t= d.t
    ; t_need= 0
    ; t_len= 0
    ; k= decode }

  let decode d = d.k d
end

module Def = struct
  type src = [ `Channel of in_channel | `String of string | `Manual ]
  type dst = [ `Channel of out_channel | `Buffer of Buffer.t | `Manual ]

  type encoder =
    { src : src
    ; dst : dst
    ; level : int
    ; i : bigstring
    ; i_pos : int
    ; i_len : int
    ; o : bigstring
    ; o_pos : int
    ; o_len : int
    ; q : De.Queue.t
    ; s : De.Lz77.state
    ; e : De.Def.encoder
    ; w : De.window
    ; state : state
    ; k : encoder -> [ `Await of encoder | `Flush of encoder | `End of encoder ] }
  and state = Hd (* header process *)| Dd (* DEFLATE process *)

  type ret = [ `Await of encoder | `End of encoder | `Flush of encoder ]

  let o_rem e = e.o_len - e.o_pos + 1
  let i_rem s = s.i_len - s.i_pos + 1

  let eoi e =
    De.Lz77.src e.s bigstring_empty 0 0 ;
    { e with i= bigstring_empty
           ; i_pos= 0
           ; i_len= min_int }

  let src e s j l =
    if (j < 0 || l < 0 || j + l > bigstring_length s)
    then invalid_bounds j l ;
    De.Lz77.src e.s s j l ;
    if (l == 0) then eoi e
    else { e with i= s; i_pos= j; i_len= j + l - 1 }

  let dst e s j l =
    if (j < 0 || l < 0 || j + l > bigstring_length s)
    then invalid_bounds j l ;
    ( ( match e.state with
          | Hd -> ()
          | Dd -> De.Def.dst e.e s j l )
    ; { e with o= s; o_pos= j; o_len= j + l - 1 } )

  let refill k e = match e.src with
    | `String _ -> k (eoi e)
    | `Channel ic ->
      let res = input_bigstring ic e.i 0 (bigstring_length e.i) in
      k (src e e.i 0 res)
    | `Manual -> `Await { e with k }

  let flush k e = match e.dst with
    | `Buffer b ->
      let len = bigstring_length e.o - o_rem e in
      for i = 0 to len - 1
      do Buffer.add_char b e.o.{i} done ;
      k (dst e e.o 0 (bigstring_length e.o))
    | `Channel oc ->
      output_bigstring oc e.o 0 (bigstring_length e.o - o_rem e) ;
      k (dst e e.o 0 (bigstring_length e.o))
    | `Manual -> `Flush { e with k }

  let identity e = `End e

  let rec checksum e =
    let k e =
      let checksum = Optint.to_int32 (De.Lz77.checksum e.s) in
      unsafe_set_uint32_be e.o e.o_pos checksum ;
      flush identity { e with o_pos= e.o_pos + 4 } in
    if o_rem e >= 4 then k e else flush checksum e

  let make_block ?(last= false) e =
    if last = false
    then
      let literals = De.Lz77.literals e.s in
      let distances = De.Lz77.distances e.s in
      let dynamic = De.Def.dynamic_of_frequencies ~literals ~distances in
      { De.Def.kind= De.Def.Dynamic dynamic; last; }
    else { De.Def.kind= De.Def.Fixed; last; }

  let rec encode e = match e.state with
    | Hd ->
      let k e =
        let header = (_deflated + ((De.window_bits e.w - 8) lsl 4)) lsl 8 in
        let header = header lor (e.level lsl 6) in
        let header = header + (31 - (header mod 31)) in
        unsafe_set_uint16_be e.o e.o_pos header ;
        if i_rem e > 0 then De.Lz77.src e.s e.i e.i_pos (i_rem e) ;
        (* XXX(dinosaure): we need to protect [e.s] against EOI signal. *)
        De.Def.dst e.e e.o (e.o_pos + 2) (o_rem e - 2) ;
        encode { e with state= Dd; o_pos= e.o_pos + 2 } in
      if o_rem e >= 2 then k e else flush encode e
    | Dd ->
      let rec partial k e =
        k e (De.Def.encode e.e `Await)
      and compress e =
        match De.Lz77.compress e.s with
        | `Await ->
          refill compress { e with i_pos= e.i_pos + (i_rem e - De.Lz77.src_rem e.s) }
        | `Flush ->
          encode_deflate e (De.Def.encode e.e `Flush)
        | `End ->
          De.Queue.push_exn e.q De.Queue.eob ;
          let block = make_block ~last:true e in
          trailing e (De.Def.encode e.e (`Block block))
      and encode_deflate e = function
        | `Partial ->
          let len = o_rem e - De.Def.dst_rem e.e in
          flush (partial encode_deflate) { e with o_pos= e.o_pos + len }
        | `Ok ->
          compress e
        | `Block ->
          let block = make_block e in
          encode_deflate e (De.Def.encode e.e (`Block block))
      and trailing e = function
        | `Partial ->
          let len = o_rem e - De.Def.dst_rem e.e in
          flush (partial trailing) { e with o_pos= e.o_pos + len }
        | `Ok ->
          let len = o_rem e - De.Def.dst_rem e.e in
          checksum { e with o_pos= e.o_pos + len }
        | `Block -> assert false (* XXX(dinosaure): should never occur! *) in

      compress e

  let src_rem = i_rem
  let dst_rem = o_rem

  let encoder src dst ~q ~w ~level =
    let i, i_pos, i_len = match src with
      | `Manual -> bigstring_empty, 1, 0
      | `String x -> bigstring_of_string x, 0, String.length x - 1
      | `Channel _ -> bigstring_create io_buffer_size, 1, 0 in
    let o, o_pos, o_len = match dst with
      | `Manual -> bigstring_empty, 1, 0
      | `Buffer _
      | `Channel _ -> bigstring_create io_buffer_size, 0, io_buffer_size - 1 in
    if level < 0 || level > 3
    then invalid_arg "Invalid compression level %d (must be in the range 0...3)" level ;
    { src
    ; dst
    ; i; i_pos; i_len
    ; o; o_pos; o_len
    ; level
    ; e= De.Def.encoder `Manual ~q
    ; s= De.Lz77.state `Manual ~q ~w
    ; q
    ; w
    ; state= Hd
    ; k= encode }

  let encode e = e.k e
end

module Higher = struct
  let compress ?(level= 0) ~w ~q ~i ~o ~refill ~flush =
    let encoder = Def.encoder `Manual `Manual ~q ~w ~level in
    let rec go encoder = match Def.encode encoder with
      | `Await encoder ->
        let len = refill i in
        go (Def.src encoder i 0 len)
      | `Flush encoder ->
        let len = bigstring_length o - Def.dst_rem encoder in
        flush o len ; go (Def.dst encoder o 0 (bigstring_length o))
      | `End encoder ->
        let len = bigstring_length o - Def.dst_rem encoder in
        if len > 0 then flush o len in
    go (Def.dst encoder o 0 (bigstring_length o))

  let uncompress ~allocate ~i ~o ~refill ~flush =
    let decoder = Inf.decoder `Manual ~allocate ~o in
    let rec go decoder = match Inf.decode decoder with
      | `Await decoder ->
        let len = refill i in
        go (Inf.src decoder i 0 len)
      | `Flush decoder ->
        let len = bigstring_length o - Inf.dst_rem decoder in
        flush o len ; go (Inf.flush decoder)
      | `End decoder ->
        let len = bigstring_length o - Inf.dst_rem decoder in
        if len > 0 then flush o len ; Ok ()
      | `Malformed err -> Error (`Msg err) in
    go decoder
end