Source file lzo.ml

module Bigarray = Bigarray_compat (* XXX(dinosaure): MirageOS compatibility. *)
type bigstring =
  (char, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t

type error =
  [ `Malformed of string
  | `Invalid_argument of string
  | `Invalid_dictionary ]

let pf = Format.fprintf

let pp_error ppf = function
  | `Malformed err -> pf ppf "%s" err
  | `Invalid_argument err -> pf ppf "%s" err
  | `Invalid_dictionary -> pf ppf "Invalid dictionary"

let ( <.> ) f g = fun x -> f (g x)

let bigstring_length x = Bigarray.Array1.dim x [@@inline]
let bigstring_create l = Bigarray.Array1.create Bigarray.char Bigarray.c_layout l
let bigstring_sub buf off len = Bigarray.Array1.sub buf off len
let bigstring_empty = bigstring_create 0

(* XXX(dinosaure): we want to control which exception is raised if we want to have
   a bad access. In this case, [Out_of_bound]. *)

exception Out_of_bound

external unsafe_get_char : bigstring -> int -> char = "%caml_ba_unsafe_ref_1"

let get_char buf ofs =
  if ofs < 0 || ofs > bigstring_length buf - 1
  then raise Out_of_bound ;
  unsafe_get_char buf ofs

external unsafe_get_int16 : bigstring -> int -> int = "%caml_bigstring_get16u"

let get_int16 buf ofs =
  if ofs < 0 || ofs > bigstring_length buf - 2
  then raise Out_of_bound ;
  unsafe_get_int16 buf ofs

external unsafe_get_int8  : bigstring -> int -> int = "%caml_ba_unsafe_ref_1"

let get_int8 buf ofs =
  if ofs < 0 || ofs > bigstring_length buf - 1
  then raise Out_of_bound ;
  unsafe_get_int8 buf ofs

external unsafe_set_int8  : bigstring -> int -> int -> unit = "%caml_ba_unsafe_set_1"

let set_int8 buf ofs x =
  if ofs < 0 || ofs > bigstring_length buf - 1
  then raise Out_of_bound ;
  unsafe_set_int8 buf ofs x

external unsafe_get_int32 : bigstring -> int -> int32 = "%caml_bigstring_get32"
external unsafe_get_int64 : bigstring -> int -> int64 = "%caml_bigstring_get64"

let get_int32 buf ofs =
  if ofs < 0 || ofs > bigstring_length buf - 4
  then raise Out_of_bound ;
  unsafe_get_int32 buf ofs

let get_int64 buf ofs =
  if ofs < 0 || ofs > bigstring_length buf - 8
  then raise Out_of_bound ;
  unsafe_get_int64 buf ofs

external unsafe_set_int32 : bigstring -> int -> int32 -> unit = "%caml_bigstring_set32"

let set_int32 buf ofs x =
  if ofs < 0 || ofs > bigstring_length buf - 4
  then raise Out_of_bound ;
  unsafe_set_int32 buf ofs x

external swap16 : int -> int = "%bswap16"

(* XXX(dinosaure): assume that LZO does need [memcpy] behaviour. *)
let unsafe_blit src src_off dst dst_off len =
  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_int32 src (src_off + i) in
    unsafe_set_int32 dst (dst_off + i) v
  done ;

  for i = 0 to len0 - 1 do
    let i = len1 * 4 + i in
    let v = unsafe_get_int8 src (src_off + i) in
    unsafe_set_int8 dst (dst_off + i) v
  done

let blit src src_off dst dst_off len =
  if len < 0 || src_off < 0 || src_off > bigstring_length src - len
             || dst_off < 0 || dst_off > bigstring_length dst - len
  then raise Out_of_bound ;
  unsafe_blit src src_off dst dst_off len

external bytes_unsafe_set_int32 : bytes -> int -> int32 -> unit = "%caml_bytes_set32u"
external bytes_unsafe_set_int8 : bytes -> int -> int -> unit = "%bytes_unsafe_set"

let unsafe_bigstring_to_string buf ofs len =
  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_int32 buf (ofs + i) in
    bytes_unsafe_set_int32 res i v
  done ;

  for i = 0 to len0 - 1 do
    let i = len1 * 4 + i in
    let v = unsafe_get_int8 buf (ofs + i) in
    bytes_unsafe_set_int8 res i v
  done ;

  Bytes.unsafe_to_string res

let bigstring_to_string buf ofs len =
  if ofs < 0 || len < 0 || ofs > bigstring_length buf - len
  then raise Out_of_bound ;
  unsafe_bigstring_to_string buf ofs len

let unsafe_get_int16_le =
  if Sys.big_endian
  then fun buf off -> swap16 (unsafe_get_int16 buf off)
  else fun buf off -> unsafe_get_int16 buf off

let unsafe_get_int16_be =
  if Sys.big_endian
  then fun buf off -> unsafe_get_int16 buf off
  else fun buf off -> swap16 (unsafe_get_int16 buf off)

let unsafe_get_int16 buf ofs = function
  | `LE -> unsafe_get_int16_le buf ofs
  | `BE -> unsafe_get_int16_be buf ofs

let get_int16 buf ofs endian =
  if ofs < 0 || ofs > bigstring_length buf - 2
  then raise Out_of_bound ;
  unsafe_get_int16 buf ofs endian

let kstrf k fmt = Format.kasprintf k fmt

type sub = { off : int; len : int; }

module State : sig
  type t = private int

  val of_int : int -> t

  val _0 : t
  val _3 : t
end = struct
  type t = int

  let of_int x = x

  let _0 = 0
  let _3 = 3
end

type ('a, 'error) t =
  | Return   : 'a -> ('a, 'error) t
  | Bind     : ('a, 'error) t * ('a -> ('b, 'error) t) -> ('b, 'error) t
  | State    : (state, 'error) t
  | Transmit : int * state -> (unit, 'error) t
  | Count    : (int, 'error) t
  | Copy     : sub * state -> (unit, 'error) t
  | Fail     : 'error -> ('a, 'error) t
  | Peek     : 'a value -> ('a, 'error) t
  | Junk     : 'a value -> (unit, 'error) t
  | Fix      : (('a, 'error) t -> ('a, 'error) t) -> ('a, 'error) t
  | Lazy     : ('a, 'error) t Lazy.t -> ('a, 'error) t
and 'a value =
  | Byte : char value
  | Short : [ `LE | `BE ] -> int value
and state = State.t

type v =
  { i : bigstring
  ; mutable i_pos : int
  ; mutable o : bigstring
  ; mutable o_pos : int
  ; mutable state : State.t }

let transmit v len =
  blit v.i v.i_pos v.o v.o_pos len ;
  v.i_pos <- v.i_pos + len ;
  v.o_pos <- v.o_pos + len ;
  Ok ()

let copy v ~off ~len =
  blit v.o (v.o_pos - off) v.o v.o_pos len ;
  v.o_pos <- v.o_pos + len ;
  Ok ()

let transmit_to_buffer buf v len =
  let res = bigstring_to_string v.i v.i_pos len in
  Buffer.add_string buf res ;
  v.i_pos <- v.i_pos + len ;
  v.o_pos <- v.o_pos + len ;
  Ok ()

let copy_to_buffer buf v ~off ~len =
  let rec go off len =
    if len = 0 then Ok ()
    else
      let pos = Buffer.length buf - off in
      let rem = Buffer.length buf - pos in
      let cpy = min rem len in
      let res = Buffer.sub buf pos cpy in
      Buffer.add_string buf res ;
      go (off + cpy) (len - cpy) in
  if Buffer.length buf >= off
  then go off len
  else Error `Invalid_dictionary

let count t =
  let res = ref 0 in
  let idx = ref t.i_pos in
  let max = bigstring_length t.i in

  while not (!idx > max - 4) && get_int32 t.i !idx = 0l
  do idx := !idx + 4 ; res := !res + 4 done ;
  while !idx + 1 <= max && get_int8 t.i !idx = 0
  do incr idx ; incr res done ;

  if !idx < max
  then ( t.i_pos <- !idx + 1
       ; Ok (!res * 255 + get_int8 t.i !idx) )
  else Error (`Malformed "Invalid input")

type ('a, 'b) k =
  | Ok of 'a | Error of 'b | Kontinuation of (unit -> ('a, 'b) t)

let ( >>= ) : ('a, 'err) result -> ('a -> ('b, 'err) result) -> ('b, 'err) result = fun  x f -> match x with
  | Ok x -> f x
  | Error err -> Error err

let copy_done ~transmit t =
  let state = (t.state :> int) in
  transmit t (state land 3)

let run
  :  transmit:(v -> int -> (unit, 'error) result)
  -> copy:(v -> off:int -> len:int -> (unit, 'error) result)
  -> ('a, 'error) t -> v -> ('a, 'error) result
  = fun ~transmit ~copy fiber t ->
    let rec go
      : type a. v -> (a, ([> `Malformed of string ] as 'error)) t -> (a, 'error) k
      = fun t -> function
      | Fail err -> Error err
      | Return v -> Ok v
      | Peek Byte -> Ok (get_char t.i t.i_pos)
      | Junk Byte ->
        if t.i_pos < bigstring_length t.i
        then ( t.i_pos <- t.i_pos + 1 ; Ok () ) else raise Out_of_bound
      | Peek (Short endian) ->
        Ok (get_int16 t.i t.i_pos endian)
      | Junk (Short _) ->
        if t.i_pos + 1 < bigstring_length t.i
        then ( t.i_pos <- t.i_pos + 2 ; Ok () ) else raise Out_of_bound
      | Bind (x, f) ->
        ( match go t x with
        | Ok v -> go t (f v)
        | Error _ as err -> err
        | Kontinuation x -> go t (Bind (x (), f)) )
      | Fix fix ->
        (* XXX(dinosaure): [Kontinuation] exists to break the stack-overflow with [js_of_ocaml] but
           it was not implemented yet. *)
        let rec m = lazy (fix r)
        and r = Lazy m in
        go t r
      | State -> Ok t.state
      | Count ->
        ( match count t with Ok v -> Ok v | Error err -> Error err )
      | Transmit (len, state) ->
        t.state <- state ;
        ( match transmit t len with Ok v -> Ok v | Error err -> Error err )
      | Lazy m -> go t (Lazy.force m)
      | Copy ({ off; len; }, state) ->
        t.state <- state ;
        let fiber =
          copy t ~off ~len:(len + 2) >>= fun () ->
          copy_done ~transmit t in
        match fiber with
        | Ok v -> Ok v
        | Error err -> Error err in
    let rec unroll t fiber : _ result = match go t fiber with
      | Ok v -> Ok v
      | Error err -> Error err
      | Kontinuation fiber -> unroll t (fiber ()) in
    unroll t fiber

module DSL = struct
  let return x = Return x
  let bind x f = Bind (x, f)
  let ( >>= ) x f = Bind (x, f)
  let peek v = Peek v
  let junk v = Junk v
  let byte = Byte
  let state = State
  let transmit ~len state = Transmit (len, state)
  let count = Count
  let copy ~off ~len state = Copy ({ off; len; }, state)
  let leshort = Short `LE
  let beshort = Short `BE
  let end_of_lzo = Return ()
  let fix f = Fix f
  let malformedf fmt = kstrf (fun s -> Fail (`Malformed s)) fmt
  let read v = peek v >>= fun r -> junk v >>= fun () -> return r
end

let fiber : (unit, [> error ]) t =
  let open DSL in
  fix @@ fun m ->
  read byte >>= fun chr ->
  state >>= fun state ->
  match chr, (state :> int) land 3 with
  | '\001' .. '\015', 0 ->
    transmit (Char.code chr + 3) state >>= fun () -> m
  | '\000', 0 ->
    count >>= fun count ->
    let len = 3 + 15 + count in
    transmit ~len State._3 >>= fun () -> m
  | '\000' .. '\015', (1 | 2 | 3) ->
    let d, state =
      Char.code chr lsr 2,
      State.of_int (Char.code chr) in
    read byte >>= fun h ->
    let off = (Char.code h lsl 2) + d + 1 in
    copy ~off ~len:2 state >>= fun () -> m
  | '\016' .. '\031', _ ->
    let length = Char.code chr land 0b111 in
    let with_length len =
      read leshort >>= fun s ->
      let off =
        let h = (Char.code chr land 8) lsr 3 in
        16384 + (h lsl 14) + (s lsr 2) in
      let state = State.of_int (s land 0xff) in
      if off = 16384 then end_of_lzo else copy ~off ~len state >>= fun () -> m in
    if length = 0
    then count >>= fun count -> with_length (7 + count)
    else with_length length
  | '\032' .. '\063', _ ->
    let with_length len =
      read leshort >>= fun s ->
      let state = State.of_int (s land 0xff) in
      let off = succ (s lsr 2) in
      copy ~off ~len state >>= fun () -> m in
    let length = Char.code chr land 0b11111 in
    if length = 0
    then count >>= fun count -> with_length (31 + count)
    else with_length length
  | '\064' .. '\255', _ ->
    let state = State.of_int (Char.code chr) in
    let len, d =
      (Char.code chr lsr 5) - 1, (* t = (t >> 5) - 1 *)
      ((Char.code chr lsr 2) land 7) (* m_pos = (t >> 2) & 7 *)in
    read byte >>= fun h ->
    let off = (Char.code h lsl 3) + d + 1 in
    copy ~off ~len state >>= fun () -> m
  | _ -> assert false (* TODO: replace it by something else to ensure exhaustive pattern-matching. *)

let fiber : (unit, [> error ]) t =
  let open DSL in
  peek byte >>= fun chr ->
  match chr with
  | '\016' -> malformedf "No dictionary at offset 0 available"
  | '\000' .. '\017' -> fiber
  | '\018' -> junk byte >>= fun () -> transmit ~len:1 State._0 >>= fun () -> fiber
  | '\019' -> junk byte >>= fun () -> transmit ~len:2 State._0 >>= fun () -> fiber
  | '\020' -> junk byte >>= fun () -> transmit ~len:3 State._0 >>= fun () -> fiber
  | '\021' -> junk byte >>= fun () -> transmit ~len:4 State._0 >>= fun () -> fiber
  | '\022' .. '\255' as chr ->
    let len = Char.code chr - 17 in
    junk byte >>= fun () -> transmit ~len State._0

let uncompress input output : (bigstring, [> error ]) result =
  let v = { i= input
          ; i_pos= 0
          ; o= output
          ; o_pos= 0
          ; state= State._0 } in
  match run ~transmit ~copy fiber v with
  | Ok () -> Ok (bigstring_sub output 0 v.o_pos)
               (* TODO(dinosaure): we can replace it by [unsafe_sub]. *)
  | Error (#error as err)-> Error err
  | exception Out_of_bound -> Error (`Invalid_argument "Input is malformed or output is not large enough")

let uncompress_with_buffer ?(chunk= 0x1000) input : (string, [> error ]) result =
  let v = { i= input
          ; i_pos= 0
          ; o= bigstring_empty
          ; o_pos= 0
          ; state= State._0 } in
  let buf = Buffer.create chunk in
  let transmit v len = transmit_to_buffer buf v len in
  let copy v ~off ~len = copy_to_buffer buf v ~off ~len in
  match run ~transmit ~copy fiber v with
  | Ok () -> Ok (Buffer.contents buf)
  | Error (#error as err) -> Error err
  | exception Out_of_bound -> Error (`Malformed "Malformed input")

(* inflate *)

let _m3_marker = 32
let _m4_marker = 16

let _m2_max_len = 8
let _m3_max_len = 33
let _m4_max_len = 9

let _m2_max_offset = 0x0800
let _m3_max_offset = 0x4000

let (.%[]) buf ofs = get_int8 buf ofs
let (.%[]<-) buf ofs v = set_int8 buf ofs v

external swap32 : int32 -> int32 = "caml_int32_bswap"

let index =
  [|  0;  1;  2; 53;  3;  7; 54; 27;  4; 38; 41;  8; 34; 55; 48; 28;
     62;  5; 39; 46; 44; 42; 22;  9; 24; 35; 59; 56; 49; 18; 29; 11;
     63; 52;  6; 26; 37; 40; 33; 47; 61; 45; 43; 21; 23; 58; 17; 10;
     51; 25; 36; 32; 60; 20; 57; 16; 50; 31; 19; 15; 30; 14; 13; 12; |]

let ctz v =
  let neg = Int64.neg in
  let ( land ) = Int64.logand in
  let ( *    ) = Int64.mul in
  let ( >>   ) = Int64.shift_right_logical in
  let idx = ((v land (neg v)) * 0x022fdd63cc95386dL) >> 58 in
  index.(Int64.to_int idx)

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

let record_match ~off ~len out_data _anchor out_pos =
  let out_pos = ref out_pos in

  if len <= _m2_max_len && off <= _m2_max_offset
  then ( let off = off - 1 in
         out_data.%[!out_pos] <- ((len - 1) lsl 5) lor ((off land 7) lsl 2) ; incr out_pos ;
         out_data.%[!out_pos] <- off asr 3 ; incr out_pos )
  else if off <= _m3_max_offset
  then ( let off = off - 1 in
         if len <= _m3_max_len
         then ( out_data.%[!out_pos] <- _m3_marker lor (len - 2) ; incr out_pos )
         else
           ( let len = ref (len - _m3_max_len) in
             out_data.%[!out_pos] <- _m3_marker lor 0 ; incr out_pos
           ; while !len > 255
             do len := !len - 255 ; out_data.%[!out_pos] <- 0 ; incr out_pos done ;
             out_data.%[!out_pos] <- !len ; incr out_pos ) ;
         out_data.%[!out_pos] <- off lsl 2 ; incr out_pos ;
         out_data.%[!out_pos] <- off asr 6 ; incr out_pos )
  else
    ( let off = off - 0x4000 in
      if len <= _m4_max_len
      then ( out_data.%[!out_pos] <- _m4_marker lor ((off asr 11) land 8) lor (len - 2) ; incr out_pos )
      else ( let len = ref (len - _m4_max_len) in
             out_data.%[!out_pos] <- _m4_marker lor ((off asr 11) land 8) ; incr out_pos ;
             while !len > 255
             do len := !len - 255 ; out_data.%[!out_pos] <- 0 ; incr out_pos done ;
             out_data.%[!out_pos] <- !len ; incr out_pos ) ;
      out_data.%[!out_pos] <- off lsl 2 ; incr out_pos ;
      out_data.%[!out_pos] <- off asr 6 ; incr out_pos ) ;
  !out_pos

let record_literals ~off ~len in_data out_data _anchor out_pos =
  let out_pos = ref out_pos in
  let in_pos = ref off in

  if len > 0
  then ( if len <= 3
         then ( out_data.%[!out_pos - 2] <- out_data.%[!out_pos - 2] lor len
              ; unsafe_blit in_data off out_data !out_pos 4
              ; out_pos := !out_pos + len )
         else if len <= 16
         then ( out_data.%[!out_pos] <- (len - 3)
              ; incr out_pos
              ; unsafe_blit in_data off out_data !out_pos 8
              ; unsafe_blit in_data (off + 8) out_data (!out_pos + 8) 8
              ; out_pos := !out_pos + len )
         else
           ( if len <= 18
             then ( out_data.%[!out_pos] <- len - 3 ; incr out_pos )
             else ( let len' = ref (len - 18) in
                    out_data.%[!out_pos] <- 0 ; incr out_pos
                  ; while !len' > 255
                    do len' := !len' - 255 ; out_data.%[!out_pos] <- 0 ; incr out_pos done
                  ; out_data.%[!out_pos] <- !len' ; incr out_pos )
           ; unsafe_blit in_data off out_data !out_pos len
           ; out_pos := !out_pos + len
           ; in_pos := !in_pos + len ) ) ;
  !out_pos, !in_pos

let record_trailer ~off ~len in_data out_data out_pos =
  let out_pos = ref out_pos in

  if len > 0
  then ( if !out_pos = 0 && len < 238
         then ( out_data.%[!out_pos] <- 17 + len ; incr out_pos )
         else if len <= 3
         then ( out_data.%[!out_pos - 2] <- out_data.%[!out_pos - 2] lor len )
         else if len <= 18
         then ( out_data.%[!out_pos] <- len - 3 ; incr out_pos )
         else
           ( let len' = ref (len - 18) in
             out_data.%[!out_pos] <- 0 ; incr out_pos ;
             while !len' > 255 do len' := !len' - 255 ; out_data.%[!out_pos] <- 0 ; incr out_pos done ;
             out_data.%[!out_pos] <- !len' ; incr out_pos ) ;
         unsafe_blit in_data off out_data !out_pos len ) ;
  out_pos := !out_pos + len ;
  out_data.%[!out_pos] <- _m4_marker lor 1 ; incr out_pos ;
  out_data.%[!out_pos] <- 0 ; incr out_pos ;
  out_data.%[!out_pos] <- 0 ; incr out_pos ;
  !out_pos

let compress in_data in_pos in_len out_data out_pos out_len t wrkmem =
  let idx_end = max 0 (in_len - 20) in

  let rec literal idx0 idx1 op t =
    (* literal: *)
    let idx0 = idx0 + (1 + ((idx0 - idx1) asr 5)) in (* TODO: check [lsr]. *)
    next idx0 idx1 op t
  and next idx0 idx1 op t =
    (* next: *)
    if idx0 - in_pos >= idx_end
    then (* break *)
      ( let idx1 = idx1 - t and t = 0 in
        (in_len - (idx1 - in_pos - t), op) )
    else
      let v =
        get_int32_le in_data idx0 in
      let index = Int32.(logand (shift_right (mul 0x1824429dl v) (32 - 14)) (sub (shift_left 1l 14) 1l)) in
      let index = Int32.to_int index in
      let reference = Bigarray.Array1.get wrkmem index + in_pos in
      Bigarray.Array1.set wrkmem index (idx0 - in_pos) ;

      if v <> get_int32_le in_data reference
      then literal idx0 idx1 op t
      else
        ( let idx1 = idx1 - t in let t = 0 in
          let unrecorded = idx0 - idx1 in
          let op, idx1 = record_literals ~off:idx1 ~len:unrecorded in_data out_data out_pos op in

          let len = ref 4 in
          while idx0 + !len - in_pos < idx_end
                && get_int64 in_data (idx0 + !len) = get_int64 in_data (reference + !len)
          do len := !len + 8 done
          (* XXX(dinosaure): it seems that [minilzo] does not call [ctz] at the end of [progl].
             May be we do an unsafe access! TODO! *)
        ; if idx0 + !len - in_pos < in_len
          then len := !len + (ctz (Int64.logxor (get_int64 in_data (idx0 + !len)) (get_int64 in_data (reference + !len))) / 8)
        ; let op = record_match ~off:(idx0 - reference) ~len:!len out_data out_pos op in
          next (idx0 + !len) (idx0 + !len) op t ) in
  let idx0 = in_pos + (if t < 4 then 4 - t else 0) in
  literal idx0 in_pos out_pos t

type wrkmem = (int, Bigarray.int16_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t

let make_wrkmem () = Bigarray.Array1.create Bigarray.int16_unsigned Bigarray.c_layout (1 lsl 14)

module Wrkmem = struct
  let memset t v = Bigarray.Array1.fill t v
end

let compress in_data in_len out_data out_len wrkmem =
  let rec go idx len out_pos t =
    if len <= 20
    then trailer idx len out_pos t
    else (* len > 20 *)
      ( let ll = min len 49152 in
        if (t + ll) lsr 5 <= 0 then trailer idx len out_pos t
        else
          ( Wrkmem.memset wrkmem 0
          ; let t, out_pos = compress in_data idx ll out_data out_pos out_len t wrkmem in
            go (idx + ll) (len - ll) out_pos t ) )
  and trailer idx len out_pos t =
    let t = t + len in
    let out_pos = record_trailer ~off:(in_len - t) ~len:t in_data out_data out_pos in
    out_pos in
  try go 0 in_len 0 0
  with Out_of_bound -> invalid_arg "lzo: output is not large enough"

let compress in_data out_data wrkmem =
  Wrkmem.memset wrkmem 0 ;
  compress
    in_data (bigstring_length in_data)
    out_data (bigstring_length out_data)
    wrkmem