Source file IO_in.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
open Common_

class type t =
  object
    method input : bytes -> int -> int -> int
    (** Read into the slice. Returns [0] only if the
        stream is closed. *)

    method close : unit -> unit
    (** Close the input. Must be idempotent. *)
  end

let create ?(close = ignore) ~input () : t =
  object
    method close = close
    method input = input
  end

let empty : t =
  object
    method close () = ()
    method input _ _ _ = 0
  end

let of_bytes ?(off = 0) ?len (b : bytes) : t =
  (* i: current position in [b] *)
  let i = ref off in

  let len =
    match len with
    | Some n ->
      if n > Bytes.length b - off then invalid_arg "Iostream.In.of_bytes";
      n
    | None -> Bytes.length b - off
  in
  let end_ = off + len in

  object
    method input b_out i_out len_out =
      let n = min (end_ - !i) len_out in
      Bytes.blit b !i b_out i_out n;
      i := !i + n;
      n

    method close () = i := end_
  end

let of_string ?off ?len s : t = of_bytes ?off ?len (Bytes.unsafe_of_string s)

(** Read into the given slice.
      @return the number of bytes read, [0] means end of input. *)
let[@inline] input (self : #t) buf i len = self#input buf i len

(** Close the channel. *)
let[@inline] close self : unit = self#close ()

let rec really_input (self : #t) buf i len =
  if len > 0 then (
    let n = input self buf i len in
    if n = 0 then raise End_of_file;
    (really_input [@tailrec]) self buf (i + n) (len - n)
  )

let really_input_string self n : string =
  let buf = Bytes.create n in
  really_input self buf 0 n;
  Bytes.unsafe_to_string buf

let copy_into ?(buf = Bytes.create _default_buf_size) (ic : #t) (oc : IO_out.t)
    : unit =
  let continue = ref true in
  while !continue do
    let len = input ic buf 0 (Bytes.length buf) in
    if len = 0 then
      continue := false
    else
      IO_out.output oc buf 0 len
  done

let concat (l0 : t list) : t =
  let l = ref l0 in
  let rec input b i len : int =
    match !l with
    | [] -> 0
    | ic :: tl ->
      let n = ic#input b i len in
      if n > 0 then
        n
      else (
        l := tl;
        input b i len
      )
  in
  let close () = List.iter close l0 in
  create ~close ~input ()

let input_all ?(buf = Bytes.create 128) (self : #t) : string =
  let buf = ref buf in
  let i = ref 0 in

  let[@inline] full_ () = !i = Bytes.length !buf in

  let grow_ () =
    let old_size = Bytes.length !buf in
    let new_size = min Sys.max_string_length (old_size + (old_size / 4) + 10) in
    if old_size = new_size then
      failwith "input_all: maximum input size exceeded";
    let new_buf = Bytes.extend !buf 0 (new_size - old_size) in
    buf := new_buf
  in

  let rec loop () =
    if full_ () then grow_ ();
    let available = Bytes.length !buf - !i in
    let n = input self !buf !i available in
    if n > 0 then (
      i := !i + n;
      (loop [@tailrec]) ()
    )
  in
  loop ();

  if full_ () then
    Bytes.unsafe_to_string !buf
  else
    Bytes.sub_string !buf 0 !i

let of_unix_fd ?(close_noerr = false) ?(buf = Bytes.create _default_buf_size)
    (fd : Unix.file_descr) : t =
  let buf_len = ref 0 in
  let buf_off = ref 0 in

  let refill () =
    buf_off := 0;
    buf_len := IO.read fd buf 0 (Bytes.length buf)
  in

  object
    method input b i len : int =
      if !buf_len = 0 then refill ();
      let n = min len !buf_len in
      if n > 0 then (
        Bytes.blit buf !buf_off b i n;
        buf_off := !buf_off + n;
        buf_len := !buf_len - n
      );
      n

    method close () =
      if close_noerr then (
        try Unix.close fd with _ -> ()
      ) else
        Unix.close fd
  end