Source file encoder.ml

type encoder = { payload : Bytes.t; mutable pos : int }
type error = [ `No_enough_space ]

let pp_error ppf `No_enough_space = Fmt.string ppf "No enough space"

type 'err state =
  | Write of {
      buffer : string;
      off : int;
      len : int;
      continue : int -> 'err state;
    }
  | Error of 'err
  | Done

let io_buffer_size = 65536
let create () = { payload = Bytes.create io_buffer_size; pos = 0 }

exception Leave of error

let leave_with (_ : encoder) error = raise (Leave error)

let safe : (encoder -> ([> error ] as 'err) state) -> encoder -> 'err state =
 fun k encoder ->
  try k encoder with Leave (#error as err) -> Error (err :> 'err)

let flush k0 encoder =
  if encoder.pos > 0 then
    let rec k1 n =
      if n < encoder.pos then
        Write
          {
            buffer = Bytes.unsafe_to_string encoder.payload;
            off = n;
            len = encoder.pos - n;
            continue = (fun m -> k1 (n + m));
          }
      else (
        encoder.pos <- 0;
        k0 encoder)
    in
    k1 0
  else k0 encoder

let write ({ pos; payload } as encoder) s =
  let max = Bytes.length payload in
  let s_len = String.length s in
  let rem = max - pos in
  let wr_n_bytes = min rem s_len in
  Bytes.blit_string s 0 payload pos wr_n_bytes;
  encoder.pos <- pos + wr_n_bytes;
  if wr_n_bytes < s_len then leave_with encoder `No_enough_space

let blit encoder ~buf ~off ~len =
  let max = Bytes.length encoder.payload in
  let go j l encoder =
    let rem = max - encoder.pos in
    let len = if l > rem then rem else l in
    Bytes.blit_string buf (off + j) encoder.payload encoder.pos len;
    encoder.pos <- encoder.pos + len;
    if len < l then leave_with encoder `No_enough_space
  in
  go 0 len encoder