msgpack using angstrom - a preview/WIP

msgpack_precise.ml

open Angstrom

type t =
  | Nil
  | Bool of bool
  | Int_fix_pos of int
  | Int_fix_neg of int
  | Int8 of int
  | Int16 of int
  | Int32 of int32
  | Int64 of int64
  | Uint8 of int
  | Uint16 of int
  | Uint32 of int32
  | Uint64 of int64
  | Float of float
  | Double of float
  | Str_fix of string
  | Str8 of string
  | Str16 of string
  | Str32 of string
  | Bin8 of string
  | Bin16 of string
  | Bin32 of string
  | Array_fix of t list
  | Array16 of t list
  | Array32 of t list
  | Map_fix of (t * t) list
  | Map16 of (t * t) list
  | Map32 of (t * t) list
  | Ext_fix_1 of int * string
  | Ext_fix_2 of int * string
  | Ext_fix_4 of int * string
  | Ext_fix_8 of int * string
  | Ext_fix_16 of int * string
  | Ext8 of int * string
  | Ext16 of int * string
  | Ext32 of int * string

module Be = struct
  include Be

  module Es = EndianString.BigEndian

  let get_float s = Es.get_float s 0
  let get_double s = Es.get_double s 0

  let get_int8 s = Es.get_int8 s 0
  let get_int16 s = Es.get_int16 s 0
  let get_int32 s = Es.get_int32 s 0
  let get_int64 s = Es.get_int64 s 0
  let get_uint32_as_int s =
    (* Trickery is afoot:
       The Int32 module doesn't provide a simple way to convert a Int32.t
       value to an int in an unsigned-friendly way.  So we're going to fake
       it here by padding the value we read into a (theoretically signed)
       Int64.t. *)
    let v = get_int64 ("\x00\x00\x00\x00" ^ s) in
    if Int64.compare v (Int64.of_int max_int) > 0 then (
      invalid_arg "msgpack field with length longer than max_int"
    );
    Int64.to_int v

  let get_uint8 s = Es.get_uint8 s 0
  let get_uint16 s = Es.get_uint16 s 0
  let get_uint32 s = Es.get_int32 s 0
  let get_uint64 s = Es.get_int64 s 0

  let make_es size set x =
    let buf = Bytes.create size in
    set buf 0 x;
    Bytes.to_string buf

  let of_float x =
    make_es 4 Es.set_float x
  let of_double x =
    make_es 8 Es.set_double x

  let of_int8 x =
    make_es 1 Es.set_int8 x
  let of_int16 x =
    make_es 2 Es.set_int16 x
  let of_int32 x =
    make_es 4 Es.set_int32 x
  let of_int64 x =
    make_es 8 Es.set_int64 x

  let of_uint8 x =
    make_es 1 Es.set_int8 x
  let of_uint16 x =
    make_es 2 Es.set_int16 x
  let of_uint32 x =
    make_es 4 Es.set_int32 x
  let of_uint64 x =
    make_es 8 Es.set_int64 x
  let of_uint32_as_int x =
    (* Trickery is afoot:
       The Int32 module doesn't provide a simple way to convert an int value
       to an Int32.t in an unsigned-friendly way.  So we're going to fake it
       here by padding the value we write into a (theoretically signed)
       Int64.t. *)
    let s = make_es 8 Es.set_int64 (Int64.of_int x) in
    String.sub s 4 4
end

(* Helper functions *)
let raw8 =
  Be.uint8 >>= fun length ->
  take length
let raw16 =
  Be.uint16 >>= fun length ->
  take length
let raw32 =
  take 4 >>= fun s ->
  let length = Be.get_uint32_as_int s in
  take length
let ext_fix_n n =
  Be.int8 >>= fun typ ->
  take n >>= fun content ->
  return (typ, content)
let tuple a =
  a >>= fun a' ->
  a >>= fun b' ->
  return (a', b')

let msgpack any =
  any_char >>= function
    (* Untagged constants *)
  | '\xc0' -> return @@ Nil
  | '\xc2' -> return @@ Bool false
  | '\xc3' -> return @@ Bool true
  (* Untagged integers *)
  | c when Char.code c land 0x80 = 0 ->
    return @@ Int_fix_pos (Be.get_uint8 (String.make 1 c))
  | c when Char.code c land 0xe0 = 0xe0 ->
    return @@ Int_fix_neg (Be.get_int8 (String.make 1 c))
  (* Tagged integers *)
  | '\xcc' -> Be.uint8  >>| fun i -> Uint8 i
  | '\xcd' -> Be.uint16 >>| fun i -> Uint16 i
  | '\xce' -> Be.uint32 >>| fun i -> Uint32 i
  | '\xcf' -> Be.uint64 >>| fun i -> Uint64 i
  | '\xd0' -> Be.int8   >>| fun i -> Int8 i
  | '\xd1' -> Be.int16  >>| fun i -> Int16 i
  | '\xd2' -> Be.int32  >>| fun i -> Int32 i
  | '\xd3' -> Be.int64  >>| fun i -> Int64 i
  (* Tagged floating point values *)
  | '\xca' -> Be.float  >>| fun f -> Float f
  | '\xcb' -> Be.double >>| fun d -> Double d
  (* Strings *)
  | c when Char.code c land 0xe0 = 0xa0 ->
    let length = Be.get_uint8 (String.make 1 c) land 0x1f in
    take length     >>| fun x -> Str_fix x
  | '\xd9' -> raw8  >>| fun s -> Str8 s
  | '\xda' -> raw16 >>| fun s -> Str16 s
  | '\xdb' -> raw32 >>| fun s -> Str32 s
  (* Binary blobs *)
  | '\xc4' -> raw8  >>| fun b -> Bin8 b
  | '\xc5' -> raw16 >>| fun b -> Bin16 b
  | '\xc6' -> raw32 >>| fun b -> Bin32 b
  (* Extension types *)
  | '\xd4' -> ext_fix_n 1  >>| fun (typ, x) -> Ext_fix_1  (typ, x)
  | '\xd5' -> ext_fix_n 2  >>| fun (typ, x) -> Ext_fix_2  (typ, x)
  | '\xd6' -> ext_fix_n 4  >>| fun (typ, x) -> Ext_fix_4  (typ, x)
  | '\xd7' -> ext_fix_n 8  >>| fun (typ, x) -> Ext_fix_8  (typ, x)
  | '\xd8' -> ext_fix_n 16 >>| fun (typ, x) -> Ext_fix_16 (typ, x)
  | '\xc7' ->
    Be.uint8 >>= fun length ->
    Be.int8 >>= fun typ ->
    take length >>= fun content ->
    return @@ Ext8 (typ, content)
  | '\xc8' ->
    Be.uint16 >>= fun length ->
    Be.int8 >>= fun typ ->
    take length >>= fun content ->
    return @@ Ext16 (typ, content)
  | '\xc9' ->
    take 4 >>= fun s ->
    let length = Be.get_uint32_as_int s in
    Be.int8 >>= fun typ ->
    take length >>= fun content ->
    return @@ Ext32 (typ, content)
  (* Arrays *)
  | c when Char.code c land 0xf0 = 0x90 ->
    let length = Be.get_uint8 (String.make 1 c) land 0x0f in
    count length any >>| fun x -> Array_fix x
  | '\xdc' ->
    Be.uint16 >>= fun length ->
    count length any >>| fun x -> Array16 x
  | '\xdd' ->
    take 4 >>= fun s ->
    let length = Be.get_uint32_as_int s in
    count length any >>| fun x -> Array32 x
  (* Maps *)
  | c when Char.code c land 0xf0 = 0x80 ->
    let length = Be.get_uint8 (String.make 1 c) land 0x0f in
    count length (tuple any) >>| fun x -> Map_fix x
  | '\xde' ->
    Be.uint16 >>= fun length ->
    count length (tuple any) >>| fun x -> Map16 x
  | '\xdf' ->
    take 4 >>= fun s ->
    let length = Be.get_uint32_as_int s in
    count length (tuple any) >>| fun x -> Map32 x
  (* Anything else is invalid *)
  | _ -> fail "msgpack"

let msgpack = fix (fun v -> msgpack v)

let msgpacks =
  many msgpack

let of_string s =
  parse_only msgpack (`String s)

let of_string_exn s =
  match of_string s with
  | Result.Ok msg -> msg
  | Result.Error e -> invalid_arg e

let msgs_of_string s =
  parse_only msgpacks (`String s)

let msgs_of_string_exn s =
  match msgs_of_string s with
  | Result.Ok msg -> msg
  | Result.Error e -> invalid_arg e

let to_buffer buf msg =
  let bc c = Buffer.add_char buf (Char.chr c) in
  let bs s = Buffer.add_string buf s in
  let rec add m =
    match m with
    | Nil -> bc 0xc0
    | Bool false -> bc 0xc2
    | Bool true -> bc 0xc3
    | Int_fix_pos i -> bs (Be.of_uint8 i)
    | Int_fix_neg i -> bs (Be.of_int8 i)
    | Int8 i -> bc 0xd0; bs (Be.of_int8 i)
    | Int16 i -> bc 0xd1; bs (Be.of_int16 i)
    | Int32 i -> bc 0xd2; bs (Be.of_int32 i)
    | Int64 i -> bc 0xd3; bs (Be.of_int64 i)
    | Uint8 i -> bc 0xcc; bs (Be.of_uint8 i)
    | Uint16 i -> bc 0xcd; bs (Be.of_uint16 i)
    | Uint32 i -> bc 0xce; bs (Be.of_uint32 i)
    | Uint64 i -> bc 0xcf; bs (Be.of_uint64 i)
    | Float f -> bc 0xca; bs (Be.of_float f)
    | Double d -> bc 0xcb; bs (Be.of_double d)
    | Str_fix s -> bc (0xa0 lor String.length s); bs s
    | Str8 s ->
      bc 0xd9;
      bs (Be.of_uint8 (String.length s));
      bs s
    | Str16 s ->
      bc 0xda;
      bs (Be.of_uint16 (String.length s));
      bs s
    | Str32 s ->
      bc 0xdb;
      bs (Be.of_uint32_as_int (String.length s));
      bs s
    | Bin8 b ->
      bc 0xc4;
      bs (Be.of_uint8 (String.length b));
      bs b
    | Bin16 b ->
      bc 0xc5;
      bs (Be.of_uint16 (String.length b));
      bs b
    | Bin32 b ->
      bc 0xc6;
      bs (Be.of_uint32_as_int (String.length b));
      bs b
    | Array_fix a -> bc (0x90 lor List.length a); List.iter add a
    | Array16 a ->
      bc 0xdc;
      bs (Be.of_uint16 (List.length a));
      List.iter add a
    | Array32 a ->
      bc 0xdd;
      bs (Be.of_uint32_as_int (List.length a));
      List.iter add a
    | Map_fix m ->
      bc (0x80 lor List.length m);
      List.iter (fun (k, v) -> add k; add v) m
    | Map16 m ->
      bc 0xde;
      bs (Be.of_uint16 (List.length m));
      List.iter (fun (k, v) -> add k; add v) m
    | Map32 m ->
      bc 0xdf;
      bs (Be.of_uint32_as_int (List.length m));
      List.iter (fun (k, v) -> add k; add v) m
    | Ext_fix_1 (typ, s) ->
      bc 0xd4;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_2 (typ, s) ->
      bc 0xd5;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_4 (typ, s) ->
      bc 0xd6;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_8 (typ, s) ->
      bc 0xd7;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_16 (typ, s) ->
      bc 0xd8;
      bs (Be.of_int8 typ);
      bs s
    | Ext8 (typ, s) ->
      bc 0xc7;
      bs (Be.of_uint8 (String.length s));
      bs (Be.of_int8 typ);
      bs s
    | Ext16 (typ, s) ->
      bc 0xc8;
      bs (Be.of_uint16 (String.length s));
      bs (Be.of_int8 typ);
      bs s
    | Ext32 (typ, s) ->
      bc 0xc9;
      bs (Be.of_uint32_as_int (String.length s));
      bs (Be.of_int8 typ);
      bs s
  in
  add msg

let to_string msg =
  let buf = Buffer.create 1_024 in
  to_buffer buf msg;
  Buffer.contents buf

let msgs_to_string msgs =
  let buf = Buffer.create 1_024 in
  List.iter (to_buffer buf) msgs;
  Buffer.contents buf