redis_eio.ml

open Eio.Std

module Buf = Eio.Buf_read
module Flow = Eio.Flow


open struct
  let spf = Printf.sprintf
end

module Message = struct
  type t =
    | Simple of string
    | Bulk of string
    | Int of int
    | Error of string
    | Array of t list

  let rec pp out : t -> unit = function
    | Simple s -> Format.fprintf out "(simple %S)" s
    | Bulk c -> Format.fprintf out "(bulk %S)" c
    | Int i -> Format.fprintf out "%d" i
    | Error e -> Format.fprintf out "(error %S)" e
    | Array l ->
      Format.fprintf out "[@[%a@]]"
        (Format.pp_print_list ~pp_sep:(fun out () -> Format.fprintf out ";@ ") pp) l

  let write (out:Flow.sink) (self:t) : unit =
    let buf = Cstruct.create 4096 in
    let i = ref 0 in

    let maybe_write ?(force=false) () =
      if !i = Cstruct.length buf || (force && !i > 0) then (
        Flow.copy (Flow.cstruct_source [Cstruct.sub buf 0 !i]) out;
        i := 0;
      )
    in

    let write_char c =
      Cstruct.set_char buf !i c;
      incr i;
      maybe_write()
    in

    let rec write_substring s off len : unit =
      let n = min (Cstruct.length buf - !i) len in
      Cstruct.blit_from_string s off buf !i n;
      i := !i + n;
      maybe_write();
      let len' = len - n in
      if len' > 0 then
        write_substring s (off + n) len'
    in
    let write_string s = write_substring s 0 (String.length s) in

    let rec loop m =
      match m with
      | Simple s ->
        write_char '+';
        write_string s;
        write_string "\r\n";
      | Error s ->
        write_char '-';
        write_string s;
        write_string "\r\n";
      | Int i ->
        write_char ':';
        write_string (string_of_int i);
        write_string "\r\n";
      | Bulk s ->
        write_char '$';
        write_string (string_of_int @@ String.length s);
        write_string "\r\n";
        write_string s;
        write_string "\r\n";
      | Array l ->
        write_char '*';
        write_string (string_of_int @@ List.length l);
        write_string "\r\n";
        List.iter loop l
    in
    loop self; maybe_write ~force:true ()

  let parse (buf:Buf.t) : t option =
    let open Buf.Syntax in
    let rec loop buf =
      let c = Buf.any_char buf in
      match c with
      | '+' ->
        let s = Buf.line buf in
        Simple s
      | '-' ->
        let s = Buf.line buf in
        Error (String.trim s)
      | '*' ->
        let s = Buf.line buf in
        let n = int_of_string @@ String.trim s in
        Array (List.init n (fun _ -> loop buf))
      | '$' ->
        let s = Buf.line buf in
        let n = int_of_string @@ String.trim s in
        let content = Buf.take n buf in
        let trail = String.trim @@ Buf.line buf in
        if trail <> "" then failwith "bulk string not ended by crlf";
        Bulk content
      | ':' ->
        let s = Buf.line buf in
        Int (int_of_string @@ String.trim s)
      | c -> failwith (spf "unknown prefix %C" c)
    in
    if Buf.at_end_of_input buf then None else Some (loop buf)

  let parse_exn buf =
    match parse buf with
    | Some m -> m
    | None -> raise End_of_file

end