Source file in.ml

type t = {
  mutable reader : unit -> chunk;
  (* Current chunk *)
  mutable buffer : buffer;
  mutable offset : int;  (** Buffer's offset *)
  mutable length : int;  (** Buffer's length *)
  (* *)
  mutable total_offset : int;  (** Total read bytes from some source *)
  (* *)
  overlap_buffer : buffer;
      (** A small buffer to resolve the data gap situation between chunks *)
}

and buffer =
  (char, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t

and chunk = buffer:buffer * offset:int * length:int

(* ===================================================================
    CONSTRUCTORS
   =================================================================== *)

let make' ?(overlap_size = 0xFF) reader =
  {
    reader;
    buffer = Bstr.empty;
    offset = 0;
    length = 0;
    total_offset = 0;
    overlap_buffer = Bstr.create overlap_size;
  }

let make ?overlap_size reader =
  let reader () =
    let buffer = reader () in
    (~buffer, ~offset:0, ~length:Bstr.(length buffer))
  in

  make' ?overlap_size reader

let of_buffer buffer =
  {
    reader = (fun () -> raise End_of_file);
    buffer;
    offset = 0;
    length = Bstr.length buffer;
    total_offset = 0;
    overlap_buffer = Bstr.empty;
  }

let of_string s = Bstr.of_string s |> of_buffer

let of_channel ?(io_buffer_size = 4096) ic =
  let buffer = Bstr.create io_buffer_size in

  let reader () =
    match In_channel.input_bigarray ic buffer 0 io_buffer_size with
    | 0 -> raise End_of_file
    | length -> (~buffer, ~offset:0, ~length)
  in

  make' reader

(* ===================================================================
    BUFFER MANIPULATION UTILITY FUNCTIONS
   =================================================================== *)

let[@inline] available_to_read in_stream = in_stream.length

let advance_offset in_stream n =
  (* assert (in_stream.length - n >= 0); *)
  in_stream.offset <- in_stream.offset + n;
  in_stream.length <- in_stream.length - n;
  in_stream.total_offset <- in_stream.total_offset + n

let[@inline] acquire_chunk in_stream = in_stream.reader ()

let set_chunk in_stream ((~buffer, ~offset, ~length) : chunk) =
  in_stream.buffer <- buffer;
  in_stream.offset <- offset;
  in_stream.length <- length

let get_chunk in_stream =
  (~buffer:in_stream.buffer, ~offset:in_stream.offset, ~length:in_stream.length)

let[@inline] position in_stream = in_stream.total_offset

let rec consume_bytes in_stream len =
  if len <> 0 then begin
    let available_bytes = available_to_read in_stream in
    let available_to_consume = min available_bytes len in

    advance_offset in_stream available_to_consume;

    if available_to_consume < len then begin
      set_chunk in_stream (acquire_chunk in_stream);
      consume_bytes in_stream (len - available_to_consume)
    end
  end

(* ===================================================================
    INPUTS
   =================================================================== *)

let gen_input
    ~(blit : Bstr.t -> src_off:int -> 'buf -> dst_off:int -> len:int -> unit)
    in_stream buffer off len =
  assert (len > 0);

  let available_bytes = available_to_read in_stream in

  if available_bytes = 0 then begin
    set_chunk in_stream (acquire_chunk in_stream)
  end;

  let batched_bytes = min len available_bytes in

  blit in_stream.buffer ~src_off:in_stream.offset buffer ~dst_off:off
    ~len:batched_bytes;

  advance_offset in_stream batched_bytes;

  batched_bytes

let rec gen_really_input ~blit in_stream buffer off len =
  if len > 0 then
    let batched_bytes = gen_input ~blit in_stream buffer off len in
    gen_really_input ~blit in_stream buffer (off + batched_bytes)
      (len - batched_bytes)

let[@inline] input in_stream buffer off len =
  gen_input ~blit:Bstr.blit in_stream buffer off len

and[@inline] really_input in_stream buffer off len =
  gen_really_input ~blit:Bstr.blit in_stream buffer off len

let[@inline] input_bytes in_stream bytes off len =
  gen_input ~blit:Bstr.blit_to_bytes in_stream bytes off len

and[@inline] really_input_bytes in_stream bytes off len =
  gen_really_input ~blit:Bstr.blit_to_bytes in_stream bytes off len

(* ===================================================================
    ENSURE MECHANISM WITH OVERLAPPING
   =================================================================== *)

let push_back in_stream chunk =
  let acquire_chunk' = in_stream.reader in

  in_stream.reader <-
    (fun () ->
      in_stream.reader <- acquire_chunk';
      chunk)

let[@inline] ensure_bytes in_stream length =
  (* assert (length <= Bstr.length in_stream.overlap_buffer); *)
  let aux in_stream available_bytes length =
    if available_bytes <> 0 then (
      really_input in_stream in_stream.overlap_buffer 0 length;
      push_back in_stream @@ get_chunk in_stream;

      set_chunk in_stream (~buffer:in_stream.overlap_buffer, ~offset:0, ~length))
    else set_chunk in_stream (acquire_chunk in_stream)
  in

  let available_bytes = (available_to_read [@inlined]) in_stream in
  if available_bytes < length then aux in_stream available_bytes length

let[@inline] ensure_bytes_at in_stream len =
  (ensure_bytes [@inlined]) in_stream len;
  let offset = in_stream.offset in
  (advance_offset [@inlined]) in_stream len;
  offset

let ensure_chunk in_stream length =
  let offset = ensure_bytes_at in_stream length in
  (~buffer:in_stream.buffer, ~offset, ~length)

let ensure_buffer in_stream len =
  let off = ensure_bytes_at in_stream len in
  Bstr.sub in_stream.buffer ~off ~len

(* ===================================================================
    COMBINATORS
   =================================================================== *)

let take n input_value in_stream = List.init n @@ fun _ -> input_value in_stream

let with_size f in_stream =
  let off = position in_stream in
  let result = f in_stream in
  let off' = position in_stream in

  (result, off' - off)

let[@tail_mod_cons] rec many input_value in_stream =
  match input_value in_stream with
  | exception End_of_file -> []
  | value -> value :: many input_value in_stream

(* ===================================================================
    INPUT INTEGER VALUES
   =================================================================== *)

let[@inline] gen_get in_stream f n =
  let off = (ensure_bytes_at [@inlined]) in_stream n in
  f in_stream.buffer off

let[@inline] input_char in_stream = gen_get in_stream Bstr.get 1
let[@inline] input_int8 in_stream = gen_get in_stream Bstr.get_int8 1
let[@inline] input_uint8 in_stream = gen_get in_stream Bstr.get_uint8 1
let[@inline] input_byte in_stream = gen_get in_stream Bstr.get_uint8 1

(* int16 *)
let[@inline] input_int16_be in_stream = gen_get in_stream Bstr.get_int16_be 2
let[@inline] input_int16_ne in_stream = gen_get in_stream Bstr.get_int16_ne 2
let[@inline] input_int16_le in_stream = gen_get in_stream Bstr.get_int16_le 2
let[@inline] input_uint16_be in_stream = gen_get in_stream Bstr.get_uint16_be 2
let[@inline] input_uint16_ne in_stream = gen_get in_stream Bstr.get_uint16_ne 2
let[@inline] input_uint16_le in_stream = gen_get in_stream Bstr.get_uint16_le 2

(* int32 *)
let[@inline] input_int32_be in_stream = gen_get in_stream Bstr.get_int32_be 4
let[@inline] input_int32_ne in_stream = gen_get in_stream Bstr.get_int32_ne 4
let[@inline] input_int32_le in_stream = gen_get in_stream Bstr.get_int32_le 4

(* int64 *)
let[@inline] input_int64_be in_stream = gen_get in_stream Bstr.get_int64_be 4
let[@inline] input_int64_ne in_stream = gen_get in_stream Bstr.get_int64_ne 4
let[@inline] input_int64_le in_stream = gen_get in_stream Bstr.get_int64_le 4

(* ===================================================================
    OTHER
   =================================================================== *)

let input_string in_stream len =
  let bytes = Bytes.create len in
  (really_input_bytes [@inlined]) in_stream bytes 0 len;
  Bytes.unsafe_to_string bytes

let input_while_into_buffer ~max_len p in_stream buffer =
  let rec aux count =
    try
      if count <= max_len then begin
        (ensure_bytes [@inlined]) in_stream 1;
        let ch = Bstr.get in_stream.buffer in_stream.offset in

        if p ch then begin
          (advance_offset [@inlined]) in_stream 1;
          Buffer.add_char buffer ch;
          aux (succ count)
        end
      end
    with _ -> ()
  in

  aux 0

let check_on_end_of_file in_steam =
  try
    let new_chunk = acquire_chunk in_steam in
    push_back in_steam new_chunk;
    false
  with End_of_file -> true

let input_while ?(max_len = Int.max_int) p in_stream =
  let[@inline] input_while_buffered max_len p in_stream =
    let[@local] buffer =
      Buffer.create (if max_len <> Int.max_int then max_len else 60)
    in

    input_while_into_buffer ~max_len p in_stream buffer;
    Buffer.contents buffer
  in

  let rec try_count_in_buffer max_len p in_stream length =
    let remaining_bytes_in_buffer = available_to_read in_stream - length in

    if remaining_bytes_in_buffer > 0 then begin
      let ch = Bstr.get in_stream.buffer (in_stream.offset + length) in
      if length <= max_len && p ch then
        try_count_in_buffer max_len p in_stream (succ length)
      else input_string in_stream length
    end
    else if check_on_end_of_file in_stream then input_string in_stream length
    else (input_while_buffered [@inlined]) max_len p in_stream
  in

  try_count_in_buffer max_len p in_stream 0

let input_while' ~max_len p in_stream =
  let string = input_while ~max_len p in_stream in
  consume_bytes in_stream (max_len - String.length string);
  string

let input_line in_stream =
  let string = input_while (( <> ) '\n') in_stream in
  match input_char in_stream with
  | exception End_of_file when "" = string -> raise End_of_file
  | (exception End_of_file) | '\n' -> string
  | _ -> failwith "impossible state of input_line"

(* ===================================================================
    TRANSFORMING
   =================================================================== *)

(* ... *)