Lwt friendly mmap-like I/O

lwt_indexed_io.ml

let ( >>= ) = Lwt.( >>= )
let ( >|= ) = Lwt.( >|= )

external identity : 'a -> 'a = "%identity"

module type Src_sig = sig
  type t

  val read_bytes : t -> size:int -> offset:int -> string option Lwt.t
  val write_bytes : t -> string -> offset:int -> unit Lwt.t
end

module type Conv_sig = sig
  val get_int32 : string -> int -> int32
  val set_int32 : string -> int -> int32 -> unit
  val get_int64 : string -> int -> int64
  val set_int64 : string -> int -> int64 -> unit
end

module Make(Src : Src_sig)(Es : Conv_sig) = struct
  let read fd ~size ~offset ~get ~conv =
    Src.read_bytes fd ~size ~offset >|= fun maybe_bytes ->
    match maybe_bytes with
    | None -> None
    | Some bytes ->
      Some (
        get bytes 0
        |> conv
      )

  let write fd x ~size ~offset ~set ~conv =
    let bytes = String.make size ' ' in
    set bytes size (conv x);
    Src.write_bytes fd bytes ~offset

  let read_four_byte fd offset conv =
    read fd ~size:4 ~offset ~get:Es.get_int32 ~conv

  let write_four_byte fd offset conv =
    write fd ~size:4 ~offset ~set:Es.set_int32 ~conv

  let four_byte r w =
    (fun fd offset -> read fd ~size:4 ~offset ~get:Es.get_int32 ~conv:r),
    (fun fd offset -> write fd ~size:4 ~offset ~set:Es.set_int32 ~conv:w)

  let eight_byte r w =
    (fun fd offset -> read fd ~size:8 ~offset ~get:Es.get_int64 ~conv:r),
    (fun fd offset -> write fd ~size:8 ~offset ~set:Es.set_int64 ~conv:w)

  let read_float32, write_float32 =
    four_byte Int32.float_of_bits Int32.bits_of_float
  let read_int32, write_int32 =
    four_byte identity identity

  let read_float64, write_float64 =
    eight_byte Int64.float_of_bits Int64.bits_of_float
  let read_int64, write_int64 =
    eight_byte identity identity
end

module Fd = struct
  module Src = struct
    type t = Lwt_unix.file_descr

    let read_bytes fd ~size ~offset =
      let buffer = String.make size ' ' in
      let byte_offset = offset * size in
      Lwt_unix.lseek fd byte_offset Unix.SEEK_SET >>= fun _ ->
      Lwt_unix.read fd buffer 0 size >>= fun bytes_read ->
      let result = if bytes_read = size then Some buffer else None in
      Lwt.return result

    let write_bytes fd buffer ~offset =
      let size = String.length buffer in
      let byte_offset = offset * size in
      Lwt_unix.lseek fd byte_offset Unix.SEEK_SET >>= fun _ ->
      let rec write_loop offset remaining_bytes =
        if remaining_bytes = 0 then
          Lwt.return_unit
        else (
          Lwt_unix.write fd buffer offset remaining_bytes
          >>= fun bytes_written ->
          write_loop (offset + bytes_written) (remaining_bytes - bytes_written)
        )
      in
      write_loop 0 size
  end

  module LittleEndian = Make(Src)(EndianString.LittleEndian)
  module BigEndian = Make(Src)(EndianString.BigEndian)
end

module Io = struct
  module Src = struct
    type t = Lwt_io.input_channel * Lwt_io.output_channel

    let read_bytes (input, _) ~size ~offset =
      let buffer = String.make size '.' in
      let byte_offset = offset * size in
      Lwt_io.set_position input (Int64.of_int byte_offset) >>= fun () ->
      Lwt_io.read_into input buffer 0 size >>= fun bytes_read ->
      let result = if bytes_read = size then Some buffer else None in
      Lwt.return result

    let write_bytes (_, output) buffer ~offset =
      let size = String.length buffer in
      let byte_offset = offset * size in
      Lwt_io.set_position output (Int64.of_int byte_offset) >>= fun () ->
      Lwt_io.write_from_exactly output buffer 0 size
  end

  module LittleEndian = Make(Src)(EndianString.LittleEndian)
  module BigEndian = Make(Src)(EndianString.BigEndian)
end

test_lwt_indexed_io.ml

open Batteries
open Bigarray
let ( >>= ) = Lwt.( >>= )

let () =
  let value_count = 1_000 in
  (* Create some test values *)
  let file = "test.data" in
  File.with_file_out ~mode:[`create] file (
    fun fout ->
      for i = 0 to value_count do
        IO.write_float fout (float_of_int i)
      done
  );
  (* File descriptor base *)
  Lwt_main.run (
    Lwt_unix.openfile file [Unix.O_RDONLY] 0o644 >>= fun fd ->
    let rec loop i =
      if i >= 0 then (
        Lwt_indexed_io.Fd.LittleEndian.read_float32 fd i >>= fun maybe ->
        match maybe with
        | None -> Lwt_io.printl "File descriptor interface - BAD!"
        | Some x ->
          if float_of_int i = x then (
            loop (i - 1)
          )
          else (
            Lwt_io.printl "File descriptor interface - NOT GOOD!"
          )
      )
      else
        Lwt_io.printl "File descriptor interface - ok"
    in
    loop value_count >>= fun () ->
    Lwt_unix.close fd
  );
  (* Lwt_io-based *)
  Lwt_main.run (
    Lwt_io.with_file ~mode:Lwt_io.input file (
      fun input ->
        let output =
          Lwt_io.make ~mode:Lwt_io.output (fun _ _ _ -> Lwt.return 0)
        in
        let fd = input, output in
        let rec loop i =
          if i >= 0 then (
            Lwt_indexed_io.Io.LittleEndian.read_float32 fd i >>= fun maybe ->
            match maybe with
            | None -> Lwt_io.printl "IO interface - BAD!"
            | Some x ->
              if float_of_int i = x then (
                loop (i - 1)
              )
              else (
                Lwt_io.printl "IO interface - NOT GOOD!"
              )
          )
          else
            Lwt_io.printl "IO interface - ok"
        in
        loop value_count
    )
  )