Source file bigstring.ml

include Shexp_bigstring.Std.Bigstring

external unsafe_write
  :  Unix.file_descr
  -> t
  -> pos:int
  -> len:int
  -> int
  = "shexp_bigstring_io_write"

external unsafe_read
  :  Unix.file_descr
  -> t
  -> pos:int
  -> len:int
  -> int
  = "shexp_bigstring_io_read"

let write fd buf ~pos ~len =
  check_pos_len_exn ~pos ~len ~length:(length buf);
  unsafe_write fd buf ~pos ~len
;;

let read fd buf ~pos ~len =
  check_pos_len_exn ~pos ~len ~length:(length buf);
  unsafe_read fd buf ~pos ~len
;;

let read_all fd =
  fold_temporary ~size:4096 ~init:0 ~f:(fun buf pos ->
    match read fd buf ~pos ~len:(length buf - pos) with
    | 0 -> Return (sub_string buf ~pos:0 ~len:pos)
    | n ->
      let pos = pos + n in
      let size = length buf in
      let new_size =
        if pos = size
        then (
          let size = size * 2 in
          if size < 0 then max_int else size)
        else size
      in
      Resize { new_size; state = pos }
    | exception Unix.Unix_error (EINTR, _, _) -> Continue { state = pos })
;;

let write_all fd s =
  let len = String.length s in
  with_temporary ~size:len ~f:(fun buf ->
    blit_string_t ~src:s ~src_pos:0 ~dst:buf ~dst_pos:0 ~len;
    let rec loop pos =
      if pos = len
      then ()
      else (
        match write fd buf ~pos ~len:(len - pos) with
        | n -> loop (pos + n)
        | exception Unix.Unix_error (EINTR, _, _) -> loop pos)
    in
    loop 0)
;;

let read_exactly fd len =
  with_temporary ~size:len ~f:(fun buf ->
    let rec loop pos =
      if pos = len
      then sub_string buf ~pos:0 ~len
      else (
        match read fd buf ~pos ~len:(len - pos) with
        | n -> loop (pos + n)
        | exception Unix.Unix_error (EINTR, _, _) -> loop pos)
    in
    loop 0)
;;

type read_all_interruptible_result =
  { interrupted : bool
  ; collected : string
  }

let read_all_interruptible ?(delay = 1.0) fd ~stop =
  fold_temporary ~size:4096 ~init:0 ~f:(fun buf pos ->
    let size = length buf in
    match Unix.select [ fd ] [] [] delay with
    | [], [], [] ->
      if stop ()
      then Return { interrupted = true; collected = sub_string buf ~pos:0 ~len:pos }
      else Continue { state = pos }
    | _ ->
      (match read fd buf ~pos ~len:(size - pos) with
       | 0 -> Return { interrupted = false; collected = sub_string buf ~pos:0 ~len:pos }
       | n ->
         let pos = pos + n in
         let new_size =
           if pos = size
           then (
             let size = size * 2 in
             if size < 0 then max_int else size)
           else size
         in
         Resize { new_size; state = pos }))
;;

type separator =
  | End_of_line
  | Char of char

let fold_gen fd ~sep ~init ~f =
  fold_temporary ~size:4096 ~init:(0, init) ~f:(fun buf (pos, acc) ->
    match read fd buf ~pos ~len:(length buf - pos) with
    | 0 ->
      if pos = 0 then Return acc else Return (f acc (sub_string buf ~pos:0 ~len:pos))
    | n ->
      let rec loop acc ~start ~pos ~stop : (_, _) fold_temporary_result =
        let char =
          match sep with
          | End_of_line -> '\n'
          | Char c -> c
        in
        match index buf ~pos ~len:(stop - pos) ~char with
        | None ->
          if start > 0
          then blit ~src:buf ~dst:buf ~src_pos:start ~dst_pos:0 ~len:(stop - start);
          let pos = stop - start in
          let size = length buf in
          let new_size =
            if pos = size
            then (
              let size = size * 2 in
              if size < 0 then max_int else size)
            else size
          in
          Resize { new_size; state = pos, acc }
        | Some end_of_chunk ->
          let chunk_len =
            if sep = End_of_line && end_of_chunk > 0 && buf.{end_of_chunk - 1} = '\r'
            then end_of_chunk - 1 - start
            else end_of_chunk - start
          in
          let chunk = sub_string buf ~pos:start ~len:chunk_len in
          let acc = f acc chunk in
          loop acc ~start:(end_of_chunk + 1) ~pos:(end_of_chunk + 1) ~stop
      in
      loop acc ~start:0 ~pos ~stop:(pos + n)
    | exception Unix.Unix_error (EINTR, _, _) -> Continue { state = pos, acc })
;;

let fold_lines fd ~init ~f = fold_gen fd ~sep:End_of_line ~init ~f
let fold_chunks fd ~sep ~init ~f = fold_gen fd ~sep:(Char sep) ~init ~f