(*

 * IO - Abstract input/output

 * Copyright (C) 2003 Nicolas Cannasse

 *

 * 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 input = {

  mutable in_read : unit -> char;

  mutable in_input : Bytes.t -> int -> int -> int;

  mutable in_close : unit -> unit;

}

type 'a output = {

  mutable out_write : char -> unit;

  mutable out_output : Bytes.t -> int -> int -> int;

  mutable out_close : unit -> 'a;

  mutable out_flush : unit -> unit;

}

exception No_more_input

exception Input_closed

exception Output_closed

(* -------------------------------------------------------------- *)

(* API *)

let default_close = (fun () -> ())

let create_in ~read ~input ~close =

  {

    in_read = read;

    in_input = input;

    in_close = close;

  }

let create_out ~write ~output ~flush ~close =

  {

    out_write = write;

    out_output = output;

    out_close = close;

    out_flush = flush;

  }

let read i = i.in_read()

let nread i n =

  if n < 0 then invalid_arg "IO.nread";

  if n = 0 then Bytes.empty

  else

  let s = Bytes.create n in

  let l = ref n in

  let p = ref 0 in

  try

    while !l > 0 do

      let r = i.in_input s !p !l in

      if r = 0 then raise No_more_input;

      p := !p + r;

      l := !l - r;

    done;

    s

  with

    No_more_input as e ->

      if !p = 0 then raise e;

      Bytes.sub s 0 !p

let nread_string i n =

  (* [nread] transfers ownership of the returned string, so

     [unsafe_to_string] is safe here *)

 Bytes.unsafe_to_string (nread i n)

let really_output o s p l' =

  let sl = Bytes.length s in

  if p + l' > sl || p < 0 || l' < 0 then invalid_arg "IO.really_output";

     let l = ref l' in

  let p = ref p in

  while !l > 0 do

    let w = o.out_output s !p !l in

    if w = 0 then raise Sys_blocked_io;

    p := !p + w;

    l := !l - w;

  done;

  l'

let input i s p l =

  let sl = Bytes.length s in

  if p + l > sl || p < 0 || l < 0 then invalid_arg "IO.input";

  if l = 0 then

    0

  else

    i.in_input s p l

let really_input i s p l' =

  let sl = Bytes.length s in

  if p + l' > sl || p < 0 || l' < 0 then invalid_arg "IO.really_input";

  let l = ref l' in

  let p = ref p in

  while !l > 0 do

    let r = i.in_input s !p !l in

    if r = 0 then raise Sys_blocked_io;

    p := !p + r;

    l := !l - r;

  done;

  l'

let really_nread i n =

  if n < 0 then invalid_arg "IO.really_nread";

  if n = 0 then Bytes.empty

  else

  let s = Bytes.create n

  in

  ignore(really_input i s 0 n);

  s

let really_nread_string i n =

  (* [really_nread] transfers ownership of the returned string,

     so [unsafe_to_string] is safe here *)

  Bytes.unsafe_to_string (really_nread i n)

let close_in i =

  let f _ = raise Input_closed in

  i.in_close();

  i.in_read <- f;

  i.in_input <- f;

  i.in_close <- f

let write o x = o.out_write x

let nwrite o s =

  let p = ref 0 in

        let l = ref (Bytes.length s) in

  while !l > 0 do

    let w = o.out_output s !p !l in

    if w = 0 then raise Sys_blocked_io;

    p := !p + w;

    l := !l - w;

  done

let nwrite_string o s =

  (* [nwrite] does not mutate or capture its [bytes] input,

     so using [Bytes.unsafe_of_string] is safe here *)

  nwrite o (Bytes.unsafe_of_string s)

let output o s p l =

  let sl = Bytes.length s in

  if p + l > sl || p < 0 || l < 0 then invalid_arg "IO.output";

  o.out_output s p l

let scanf i fmt =

  let ib = Scanf.Scanning.from_function (fun () -> try read i with No_more_input -> raise End_of_file) in

  Scanf.kscanf ib (fun _ exn -> raise exn) fmt

let printf o fmt =

  Printf.kprintf (fun s -> nwrite_string o s) fmt

let flush o = o.out_flush()

let close_out o =

  let f _ = raise Output_closed in

  let r = o.out_close() in

  o.out_write <- f;

  o.out_output <- f;

  o.out_close <- f;

  o.out_flush <- f;

  r

let read_all i =

  let maxlen = 1024 in

  let str = ref [] in

  let pos = ref 0 in

  let rec loop() =

    let s = nread i maxlen in

    str := (s,!pos) :: !str;

    pos := !pos + Bytes.length s;

    loop()

  in

  try

    loop()

  with

    No_more_input ->

      let buf = Bytes.create !pos in

      List.iter (fun (s,p) ->

        Bytes.blit s 0 buf p (Bytes.length s)

      ) !str;

                        (* 'buf' doesn't escape, it won't be mutated again *)

      Bytes.unsafe_to_string buf

let pos_in i =

  let p = ref 0 in

  {

    in_read = (fun () ->

      let c = i.in_read() in

      incr p;

      c

    );

    in_input = (fun s sp l ->

      let n = i.in_input s sp l in

      p := !p + n;

      n

    );

    in_close = i.in_close

  } , (fun () -> !p)

let pos_out o =

  let p = ref 0 in

  {

    out_write = (fun c ->

      o.out_write c;

      incr p

    );

    out_output = (fun s sp l ->

      let n = o.out_output s sp l in

      p := !p + n;

      n

    );

    out_close = o.out_close;

    out_flush = o.out_flush;

  } , (fun () -> !p)

(* -------------------------------------------------------------- *)

(* Standard IO *)

let input_bytes s =

  let pos = ref 0 in

  let len = Bytes.length s in

  {

    in_read = (fun () ->

      if !pos >= len then raise No_more_input;

      let c = Bytes.unsafe_get s !pos in

      incr pos;

      c

    );

    in_input = (fun sout p l ->

      if !pos >= len then raise No_more_input;

      let n = (if !pos + l > len then len - !pos else l) in

      Bytes.unsafe_blit s !pos sout p n;

      pos := !pos + n;

      n

    );

    in_close = (fun () -> ());

  }

let input_string s =

  (* Bytes.unsafe_of_string is safe here as input_bytes does not

     mutate the byte sequence *)

  input_bytes (Bytes.unsafe_of_string s)

open ExtBuffer

let output_buffer close =

  let b = Buffer.create 0 in

  {

    out_write = (fun c -> Buffer.add_char b c);

    out_output = (fun s p l -> Buffer.add_subbytes b s p l; l);

    out_close = (fun () -> close b);

    out_flush = (fun () -> ());

  }

let output_string () = output_buffer Buffer.contents

let output_bytes () = output_buffer Buffer.to_bytes

let output_strings() =

  let sl = ref [] in

  let size = ref 0 in

  let b = Buffer.create 0 in

  {

    out_write = (fun c ->

      if !size = Sys.max_string_length then begin

        sl := Buffer.contents b :: !sl;

        Buffer.clear b;

        size := 0;

      end else incr size;

      Buffer.add_char b c

    );

    out_output = (fun s p l ->

      if !size + l > Sys.max_string_length then begin

        sl := Buffer.contents b :: !sl;

        Buffer.clear b;

        size := 0;

      end else size := !size + l;

      Buffer.add_subbytes b s p l;

      l

    );

    out_close = (fun () -> sl := Buffer.contents b :: !sl; List.rev (!sl));

    out_flush = (fun () -> ());

  }

let input_channel ch =

  {

    in_read = (fun () ->

      try

        input_char ch

      with

        End_of_file -> raise No_more_input

    );

    in_input = (fun s p l ->

      let n = Pervasives.input ch s p l in

      if n = 0 then raise No_more_input;

      n

    );

    in_close = (fun () -> Pervasives.close_in ch);

  }

let output_channel ch =

  {

    out_write = (fun c -> output_char ch c);

    out_output = (fun s p l -> Pervasives.output ch s p l; l);

    out_close = (fun () -> Pervasives.close_out ch);

    out_flush = (fun () -> Pervasives.flush ch);

  }

let input_enum e =

  let pos = ref 0 in

  {

    in_read = (fun () ->

      match Enum.get e with

      | None -> raise No_more_input

      | Some c ->

        incr pos;

        c

    );

    in_input = (fun s p l ->

      let rec loop p l =

        if l = 0 then

          0

        else

          match Enum.get e with

          | None -> l

          | Some c ->

            Bytes.unsafe_set s p c;

            loop (p + 1) (l - 1)

      in

      let k = loop p l in

      if k = l then raise No_more_input;

      l - k

    );

    in_close = (fun () -> ());

  }

let output_enum() =

  let b = Buffer.create 0 in

  {

    out_write = (fun x ->

      Buffer.add_char b x

    );

    out_output = (fun s p l ->

      Buffer.add_subbytes b s p l;

      l

    );

    out_close = (fun () ->

      let s = Buffer.contents b in

      ExtString.String.enum s

    );

    out_flush = (fun () -> ());

  }

let pipe() =

  let input = ref "" in

  let inpos = ref 0 in

  let output = Buffer.create 0 in

  let flush() =

    input := Buffer.contents output;

    inpos := 0;

    Buffer.reset output;

    if String.length !input = 0 then raise No_more_input

  in

  let read() =

    if !inpos = String.length !input then flush();

    let c = String.unsafe_get !input !inpos in

    incr inpos;

    c

  in

  let input s p l =

    if !inpos = String.length !input then flush();

    let r = (if !inpos + l > String.length !input then String.length !input - !inpos else l) in

    String.unsafe_blit !input !inpos s p r;

    inpos := !inpos + r;

    r

  in

  let write c =

    Buffer.add_char output c

  in

  let output s p l =

    Buffer.add_subbytes output s p l;

    l

  in

  let input = {

    in_read = read;

    in_input = input;

    in_close = (fun () -> ());

  } in

  let output = {

    out_write = write;

    out_output = output;

    out_close = (fun () -> ());

    out_flush = (fun () -> ());

  } in

  input , output

external cast_output : 'a output -> unit output = "%identity"

(* -------------------------------------------------------------- *)

(* BINARY APIs *)

exception Overflow of string

let read_byte i = int_of_char (i.in_read())

let read_signed_byte i =

  let c = int_of_char (i.in_read()) in

  if c land 128 <> 0 then

    c - 256

  else

    c

let read_string_into_buffer i =

  let b = Buffer.create 8 in

  let rec loop() =

    let c = i.in_read() in

    if c <> '\000' then begin

      Buffer.add_char b c;

      loop();

    end;

  in

  loop();

  b

let read_string i =

  Buffer.contents

    (read_string_into_buffer i)

let read_bytes i =

  Buffer.to_bytes

    (read_string_into_buffer i)

let read_line i =

  let b = Buffer.create 8 in

  let cr = ref false in

  let rec loop() =

    let c = i.in_read() in

    match c with

    | '\n' ->

      ()

    | '\r' ->

      cr := true;

      loop()

    | _ when !cr ->

      cr := false;

      Buffer.add_char b '\r';

      Buffer.add_char b c;

      loop();

    | _ ->

      Buffer.add_char b c;

      loop();

  in

  try

    loop();

    Buffer.contents b

  with

    No_more_input ->

      if !cr then Buffer.add_char b '\r';

      if Buffer.length b > 0 then

        Buffer.contents b

      else

        raise No_more_input

let read_ui16 i =

  let ch1 = read_byte i in

  let ch2 = read_byte i in

  ch1 lor (ch2 lsl 8)

let read_i16 i =

  let ch1 = read_byte i in

  let ch2 = read_byte i in

  let n = ch1 lor (ch2 lsl 8) in

  if ch2 land 128 <> 0 then

    n - 65536

  else

    n

let sign_bit_i32 = lnot 0x7FFF_FFFF

let read_32 ~i31 ch =

  let ch1 = read_byte ch in

  let ch2 = read_byte ch in

  let ch3 = read_byte ch in

  let ch4 = read_byte ch in

  if ch4 land 128 <> 0 then begin

    if i31 && ch4 land 64 = 0 then raise (Overflow "read_i31");

    ch1 lor (ch2 lsl 8) lor (ch3 lsl 16) lor ((ch4 land 127) lsl 24) lor sign_bit_i32

  end else begin

    if i31 && ch4 land 64 <> 0 then raise (Overflow "read_i31");

    ch1 lor (ch2 lsl 8) lor (ch3 lsl 16) lor (ch4 lsl 24)

  end

let read_i31 ch = read_32 ~i31:true ch

let read_i32_as_int ch = read_32 ~i31:false ch

let read_i32 = read_i31

let read_real_i32 ch =

  let ch1 = read_byte ch in

  let ch2 = read_byte ch in

  let ch3 = read_byte ch in

  let base = Int32.of_int (ch1 lor (ch2 lsl 8) lor (ch3 lsl 16)) in

  let big = Int32.shift_left (Int32.of_int (read_byte ch)) 24 in

  Int32.logor base big

let read_i64 ch =

  let ch1 = read_byte ch in

  let ch2 = read_byte ch in

  let ch3 = read_byte ch in

  let ch4 = read_byte ch in

  let base = Int64.of_int (ch1 lor (ch2 lsl 8) lor (ch3 lsl 16)) in

  let small = Int64.logor base (Int64.shift_left (Int64.of_int ch4) 24) in

  let big = Int64.of_int32 (read_real_i32 ch) in

  Int64.logor (Int64.shift_left big 32) small

let read_float32 ch =

  Int32.float_of_bits (read_real_i32 ch)

let read_double ch =

  Int64.float_of_bits (read_i64 ch)

let write_byte o n =

  (* doesn't test bounds of n in order to keep semantics of Pervasives.output_byte *)

  write o (Char.unsafe_chr (n land 0xFF))

let write_string o s =

  nwrite_string o s;

  write o '\000'

let write_bytes o s =

  nwrite o s;

  write o '\000'

let write_line o s =

  nwrite_string o s;

  write o '\n'

let write_ui16 ch n =

  if n < 0 || n > 0xFFFF then raise (Overflow "write_ui16");

  write_byte ch n;

  write_byte ch (n lsr 8)

let write_i16 ch n =

  if n < -0x8000 || n > 0x7FFF then raise (Overflow "write_i16");

  if n < 0 then

    write_ui16 ch (65536 + n)

  else

    write_ui16 ch n

let write_32 ch n =

  write_byte ch n;

  write_byte ch (n lsr 8);

  write_byte ch (n lsr 16);

  write_byte ch (n asr 24)

let write_i31 ch n =

#ifndef WORD_SIZE_32

  if n < -0x4000_0000 || n > 0x3FFF_FFFF then raise (Overflow "write_i31");

#endif

  write_32 ch n

let write_i32 ch n =

#ifndef WORD_SIZE_32

  if n < -0x8000_0000 || n > 0x7FFF_FFFF then raise (Overflow "write_i32");

#endif

  write_32 ch n

let write_real_i32 ch n =

  let base = Int32.to_int n in

  let big = Int32.to_int (Int32.shift_right_logical n 24) in

  write_byte ch base;

  write_byte ch (base lsr 8);

  write_byte ch (base lsr 16);

  write_byte ch big

let write_i64 ch n =

  write_real_i32 ch (Int64.to_int32 n);

  write_real_i32 ch (Int64.to_int32 (Int64.shift_right_logical n 32))

let write_float32 ch f =

  write_real_i32 ch (Int32.bits_of_float f)

let write_double ch f =

  write_i64 ch (Int64.bits_of_float f)

(* -------------------------------------------------------------- *)

(* Big Endians *)

module BigEndian = struct

let read_ui16 i =

  let ch2 = read_byte i in

  let ch1 = read_byte i in

  ch1 lor (ch2 lsl 8)

let read_i16 i =

  let ch2 = read_byte i in

  let ch1 = read_byte i in

  let n = ch1 lor (ch2 lsl 8) in

  if ch2 land 128 <> 0 then

    n - 65536

  else

    n

let sign_bit_i32 = lnot 0x7FFF_FFFF

let read_32 ~i31 ch =

  let ch4 = read_byte ch in

  let ch3 = read_byte ch in

  let ch2 = read_byte ch in

  let ch1 = read_byte ch in

  if ch4 land 128 <> 0 then begin

    if i31 && ch4 land 64 = 0 then raise (Overflow "read_i31");

    ch1 lor (ch2 lsl 8) lor (ch3 lsl 16) lor ((ch4 land 127) lsl 24) lor sign_bit_i32

  end else begin

    if i31 && ch4 land 64 <> 0 then raise (Overflow "read_i31");

    ch1 lor (ch2 lsl 8) lor (ch3 lsl 16) lor (ch4 lsl 24)

  end

let read_i31 ch = read_32 ~i31:true ch

let read_i32_as_int ch = read_32 ~i31:false ch

let read_i32 = read_i31

let read_real_i32 ch =

  let big = Int32.shift_left (Int32.of_int (read_byte ch)) 24 in

  let ch3 = read_byte ch in

  let ch2 = read_byte ch in

  let ch1 = read_byte ch in

  let base = Int32.of_int (ch1 lor (ch2 lsl 8) lor (ch3 lsl 16)) in

  Int32.logor base big

let read_i64 ch =

  let big = Int64.of_int32 (read_real_i32 ch) in

  let ch4 = read_byte ch in

  let ch3 = read_byte ch in

  let ch2 = read_byte ch in

  let ch1 = read_byte ch in

  let base = Int64.of_int (ch1 lor (ch2 lsl 8) lor (ch3 lsl 16)) in

  let small = Int64.logor base (Int64.shift_left (Int64.of_int ch4) 24) in

  Int64.logor (Int64.shift_left big 32) small

let read_float32 ch =

  Int32.float_of_bits (read_real_i32 ch)

let read_double ch =

  Int64.float_of_bits (read_i64 ch)

let write_ui16 ch n =

  if n < 0 || n > 0xFFFF then raise (Overflow "write_ui16");

  write_byte ch (n lsr 8);

  write_byte ch n

let write_i16 ch n =

  if n < -0x8000 || n > 0x7FFF then raise (Overflow "write_i16");

  if n < 0 then

    write_ui16 ch (65536 + n)

  else

    write_ui16 ch n

let write_32 ch n =

  write_byte ch (n asr 24);

  write_byte ch (n lsr 16);

  write_byte ch (n lsr 8);

  write_byte ch n

let write_i31 ch n =

#ifndef WORD_SIZE_32

  if n < -0x4000_0000 || n > 0x3FFF_FFFF then raise (Overflow "write_i31");

#endif

  write_32 ch n

let write_i32 ch n =

#ifndef WORD_SIZE_32

  if n < -0x8000_0000 || n > 0x7FFF_FFFF then raise (Overflow "write_i32");

#endif

  write_32 ch n

let write_real_i32 ch n =

  let base = Int32.to_int n in

  let big = Int32.to_int (Int32.shift_right_logical n 24) in

  write_byte ch big;

  write_byte ch (base lsr 16);

  write_byte ch (base lsr 8);

  write_byte ch base

let write_i64 ch n =

  write_real_i32 ch (Int64.to_int32 (Int64.shift_right_logical n 32));

  write_real_i32 ch (Int64.to_int32 n)

let write_float32 ch f =

  write_real_i32 ch (Int32.bits_of_float f)

let write_double ch f =

  write_i64 ch (Int64.bits_of_float f)

end

(* -------------------------------------------------------------- *)

(* Bits API *)

type 'a bc = {

  ch : 'a;

  mutable nbits : int;

  mutable bits : int;

}

type in_bits = input bc

type out_bits = unit output bc

exception Bits_error

let input_bits ch =

  {

    ch = ch;

    nbits = 0;

    bits = 0;

  }

let output_bits ch =

  {

    ch = cast_output ch;

    nbits = 0;

    bits = 0;

  }

let rec read_bits b n =

  if b.nbits >= n then begin

    let c = b.nbits - n in

    let k = (b.bits asr c) land ((1 lsl n) - 1) in

    b.nbits <- c;

    k

  end else begin

    let k = read_byte b.ch in

    if b.nbits >= 24 then begin

      if n >= 31 then raise Bits_error;

      let c = 8 + b.nbits - n in

      let d = b.bits land ((1 lsl b.nbits) - 1) in

      let d = (d lsl (8 - c)) lor (k lsr c) in

      b.bits <- k;

      b.nbits <- c;

      d

    end else begin

      b.bits <- (b.bits lsl 8) lor k;

      b.nbits <- b.nbits + 8;

      read_bits b n;

    end

  end

let drop_bits b =

  b.nbits <- 0

let rec write_bits b ~nbits x =

  let n = nbits in

  if n + b.nbits >= 32 then begin

    if n > 31 then raise Bits_error;

    let n2 = 32 - b.nbits - 1 in

    let n3 = n - n2 in

    write_bits b ~nbits:n2 (x asr n3);

    write_bits b ~nbits:n3 (x land ((1 lsl n3) - 1));

  end else begin

    if n < 0 then raise Bits_error;

    if (x < 0 || x > (1 lsl n - 1)) && n <> 31 then raise Bits_error;

    b.bits <- (b.bits lsl n) lor x;

    b.nbits <- b.nbits + n;

    while b.nbits >= 8 do

      b.nbits <- b.nbits - 8;

      write_byte b.ch (b.bits asr b.nbits)

    done

  end

let flush_bits b =

  if b.nbits > 0 then write_bits b (8 - b.nbits) 0

(* -------------------------------------------------------------- *)

(* Generic IO *)

class in_channel ch =

  object

  method input s pos len = input ch s pos len

  method close_in() = close_in ch

  end

class out_channel ch =

  object

  method output s pos len = output ch s pos len

  method flush() = flush ch

  method close_out() = ignore(close_out ch)

  end

class in_chars ch =

  object

  method get() = try read ch with No_more_input -> raise End_of_file

  method close_in() = close_in ch

  end

class out_chars ch =

  object

  method put t = write ch t

  method flush() = flush ch

  method close_out() = ignore(close_out ch)

  end

let from_in_channel ch =

  let cbuf = Bytes.create 1 in

  let read() =

    try

      if ch#input cbuf 0 1 = 0 then raise Sys_blocked_io;

      Bytes.unsafe_get cbuf 0

    with

      End_of_file -> raise No_more_input

  in

  let input s p l =

    ch#input s p l

  in

  create_in

    ~read

    ~input

    ~close:ch#close_in

let from_out_channel ch =

  let cbuf = Bytes.create 1 in

  let write c =

    Bytes.unsafe_set cbuf 0 c;

    if ch#output cbuf 0 1 = 0 then raise Sys_blocked_io;

  in

  let output s p l =

    ch#output s p l

  in

  create_out

    ~write

    ~output

    ~flush:ch#flush

    ~close:ch#close_out

let from_in_chars ch =

  let input s p l =

    let i = ref 0 in

    try

      while !i < l do

        Bytes.unsafe_set s (p + !i) (ch#get());

        incr i

      done;

      l

    with

      End_of_file when !i > 0 ->

        !i

  in

  create_in

    ~read:ch#get

    ~input

    ~close:ch#close_in

let from_out_chars ch =

  let output s p l =

    for i = p to p + l - 1 do

      ch#put (Bytes.unsafe_get s i)

    done;

    l

  in

  create_out