Source file `bytebuffer.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
155
156
157
open! Core

type t =
  { mutable buf : (Bigstring.t[@sexp.opaque])
  ; mutable pos_read : int
  ; mutable pos_fill : int
  ; max_buffer_size : int
  }
[@@deriving sexp_of]

let create ?max_buffer_size size =
  let max_buffer_size =
    match max_buffer_size with
    | None -> Int.max_value
    | Some s -> s
  in
  if size > max_buffer_size
  then
    raise_s
      [%message
        "Invalid buffer size"
          ~requested_size:(size : int)
          ~max_buffer_size:(max_buffer_size : int)];
  let buf = Bigstring.create size in
  { buf; pos_read = 0; pos_fill = 0; max_buffer_size }
;;

let unsafe_buf t = t.buf
let pos t = t.pos_read

let compact t =
  if t.pos_read > 0
  then (
    let len = t.pos_fill - t.pos_read in
    Bigstring.blit ~src:t.buf ~dst:t.buf ~src_pos:t.pos_read ~dst_pos:0 ~len;
    t.pos_read <- 0;
    t.pos_fill <- len)
;;

let length t = t.pos_fill - t.pos_read
let can_reclaim_space t = t.pos_read > 0
let capacity t = Bigstring.length t.buf
let available_to_write t = Bigstring.length t.buf - t.pos_fill

let maybe_grow_buffer t new_length =
  if new_length > t.max_buffer_size
  then
    raise_s
      [%message
        "Cannot grow Bytebuffer" ~t:(t : t) ~new_length_requested:(new_length : int)];
  let len = Int.min t.max_buffer_size (Int.ceil_pow2 new_length) in
  t.buf <- Bigstring.unsafe_destroy_and_resize t.buf ~len
;;

let drop t len =
  if len < 0 || len > length t then invalid_arg "Bytebuffer.drop: Index out of bounds";
  t.pos_read <- t.pos_read + len
;;

let read fd t =
  let count =
    Bigstring_unix.read fd t.buf ~pos:t.pos_fill ~len:(Bigstring.length t.buf - t.pos_fill)
  in
  if count > 0 then t.pos_fill <- t.pos_fill + count;
  count
;;

let write fd t =
  let count = Bigstring_unix.write fd t.buf ~pos:t.pos_read ~len:(length t) in
  if count > 0 then t.pos_read <- t.pos_read + count;
  count
;;

let read_assume_fd_is_nonblocking fd t =
  let res =
    Bigstring_unix.read_assume_fd_is_nonblocking
      fd
      t.buf
      ~pos:t.pos_fill
      ~len:(Bigstring.length t.buf - t.pos_fill)
  in
  if Core_unix.Syscall_result.Int.is_ok res
  then (
    let count = Core_unix.Syscall_result.Int.ok_exn res in
    if count > 0 then t.pos_fill <- t.pos_fill + count);
  res
;;

let write_assume_fd_is_nonblocking fd t =
  let res =
    Bigstring_unix.write_assume_fd_is_nonblocking fd t.buf ~pos:t.pos_read ~len:(length t)
  in
  if res > 0 then t.pos_read <- t.pos_read + res;
  res
;;

module Fill = struct
  let char t ch =
    if available_to_write t < 1 then maybe_grow_buffer t (Bigstring.length t.buf + 1);
    Bigstring.set t.buf t.pos_fill ch;
    t.pos_fill <- t.pos_fill + 1
  ;;

  let add_gen t ?(pos = 0) ?len ~total_length ~blit str =
    let len =
      match len with
      | Some i -> i
      | None -> total_length - pos
    in
    Ordered_collection_common.check_pos_len_exn ~pos ~len ~total_length;
    if available_to_write t < len then maybe_grow_buffer t (Bigstring.length t.buf + len);
    blit ~src:str ~src_pos:pos ~dst:t.buf ~dst_pos:t.pos_fill ~len;
    t.pos_fill <- t.pos_fill + len
  ;;

  let string t ?pos ?len str =
    add_gen
      t
      ?pos
      ?len
      ~total_length:(String.length str)
      ~blit:Bigstring.From_string.blit
      str
  ;;

  let bytes t ?pos ?len str =
    add_gen
      t
      ?pos
      ?len
      ~total_length:(Bytes.length str)
      ~blit:Bigstring.From_bytes.blit
      str
  ;;

  let bigstring t ?pos ?len str =
    add_gen t ?pos ?len ~total_length:(Bigstring.length str) ~blit:Bigstring.blit str
  ;;

  let bytebuffer t buf = bigstring t ~pos:buf.pos_read ~len:(length buf) buf.buf
end

module Consume = struct
  let stringo ?pos ?len t =
    let pos, len =
      Ordered_collection_common.get_pos_len_exn ?pos ?len () ~total_length:(length t)
    in
    let res = Bigstring.To_string.sub t.buf ~pos:(pos + t.pos_read) ~len in
    drop t len;
    res
  ;;
end

let index t ch =
  let idx = Bigstring.unsafe_find t.buf ch ~pos:t.pos_read ~len:(length t) in
  if idx < 0 then -1 else idx - t.pos_read
;;