Source file cborl.ml

(*
 * SPDX-FileCopyrightText: 2022 pukkamustard <pukkamustard@posteo.net>
 *
 * SPDX-License-Identifier: AGPL-3.0-or-later
 *)

module Item = struct
  type item =
    | Null
    | Undefined
    | Bool of bool
    | Integer of Z.t
    | ByteString of string
    | TextString of string
    | Array of item list
    | Map of (item * item) list
    | Tag of Z.t * item
    | Simple of int
    | Float of float
    | Double of float

  let rec pp ppf = function
    | Null -> Fmt.pf ppf "null"
    | Undefined -> Fmt.pf ppf "undefined"
    | Bool false -> Fmt.pf ppf "false"
    | Bool true -> Fmt.pf ppf "true"
    | Integer n -> Fmt.pf ppf "%a" Z.pp_print n
    | ByteString s ->
        Fmt.pf ppf "@[<2>h'%a'@]" Fmt.(on_string @@ octets ~sep:nop ()) s
    | TextString s -> Fmt.pf ppf "@[<1>\"%s\"@]" s
    | Array l -> Fmt.pf ppf "@[<1>[%a]@]" (Fmt.list ~sep:Fmt.comma pp) l
    | Map m ->
        Fmt.(
          pf ppf "@[<1>{%a}@]"
            (list ~sep:semi (pair ~sep:(any ":@ ") (styled `Yellow pp) pp))
            m)
    | Tag (number, content) ->
        Fmt.(
          pf ppf "%a(@[<2>%a@])" (styled `Yellow Z.pp_print) number pp content)
    | Simple number -> Fmt.(pf ppf "@[simple(%d)@]" number)
    | Float f -> Fmt.(pf ppf "%f" f)
    | Double d -> Fmt.(pf ppf "%f" d)

  let rec equal a b =
    match (a, b) with
    | Null, Null -> true
    | Undefined, Undefined -> true
    | Bool a, Bool b -> Bool.equal a b
    | Integer a, Integer b -> Z.equal a b
    | ByteString a, ByteString b -> String.equal a b
    | TextString a, TextString b -> String.equal a b
    | Array a, Array b -> List.equal equal a b
    | Map a, Map b ->
        List.equal (fun (ka, va) (kb, vb) -> equal ka kb && equal va vb) a b
    | Tag (ta, ca), Tag (tb, cb) -> Z.equal ta tb && equal ca cb
    | Simple a, Simple b -> Int.equal a b
    | Float a, Float b -> Float.equal a b
    | Double a, Double b -> Float.equal a b
    | _ -> false
end

include Item

module Signal = struct
  type t =
    | Null
    | Undefined
    | Bool of bool
    | Integer of Z.t
    | ByteString of int
    | TextString of int
    | Byte of char
    | Array of int option
    | Break
    | Map of int option
    | Tag of Z.t
    | Simple of int
    | Float of float
    | Double of float

  let pp ppf = function
    | Null -> Fmt.pf ppf "Null"
    | Undefined -> Fmt.pf ppf "Undefined"
    | Bool false -> Fmt.pf ppf "Bool false"
    | Bool true -> Fmt.pf ppf "Bool true"
    | Integer n -> Fmt.pf ppf "Integer %a" Z.pp_print n
    | ByteString l -> Fmt.pf ppf "ByteString %d" l
    | TextString l -> Fmt.pf ppf "TextString %d" l
    | Byte c -> Fmt.pf ppf "Byte %a" Fmt.(octets ()) (1, fun _ -> c)
    | Array c_opt ->
        Fmt.pf ppf "Array %a" Fmt.(option ~none:(any "-") int) c_opt
    | Map c_opt -> Fmt.pf ppf "Map %a" Fmt.(option ~none:(any "-") int) c_opt
    | Tag t -> Fmt.pf ppf "Tag %a" Z.pp_print t
    | Simple n -> Fmt.pf ppf "Simple %d" n
    | Break -> Fmt.pf ppf "Break"
    | Float f -> Fmt.(pf ppf "Float %f" f)
    | Double f -> Fmt.(pf ppf "Double %f" f)

  (* Returns the number of byte requires to represent z. *)
  let z_numbytes z =
    let z = if Z.sign z >= 0 then z else Z.(~- (~$1) - z) in
    let numbits = Z.numbits z |> Float.of_int in
    numbits /. 8.0 |> ceil |> int_of_float

  let rec of_item item =
    match item with
    | Item.Null -> Seq.return Null
    | Item.Undefined -> Seq.return Undefined
    | Item.Bool b -> Seq.return (Bool b)
    | Item.Integer z when z_numbytes z <= 8 -> Seq.return (Integer z)
    | Item.Integer z ->
        let len = z_numbytes z in
        if Z.sign z >= 0 then
          let bytes = Z.to_bits z in
          Seq.(
            append
              (cons
                 (* tag 2 for positive bignum *)
                 (Tag Z.(~$2))
                 (* ByteString with length len *)
                 (return @@ ByteString len))
              (init len (fun i -> Byte bytes.[len - 1 - i])))
        else
          let z = Z.(~- (~$1) - z) in
          let bytes = Z.to_bits z in
          Seq.(
            append
              (cons
                 (* tag 3 for negative bignum *)
                 (Tag Z.(~$3))
                 (* ByteString with length len *)
                 (return @@ ByteString len))
              (init len (fun i -> Byte bytes.[len - 1 - i])))
    | Item.ByteString s ->
        Seq.cons
          (ByteString (String.length s))
          (Seq.map (fun b -> Byte b) @@ String.to_seq s)
    | Item.TextString s ->
        Seq.cons
          (TextString (String.length s))
          (Seq.map (fun b -> Byte b) @@ String.to_seq s)
    | Item.Array l ->
        Seq.cons
          (Array (Some (List.length l)))
          Seq.(concat_map of_item @@ List.to_seq l)
    | Item.Map l ->
        Seq.cons
          (Map (Some (List.length l)))
          Seq.(
            concat_map (fun (key, value) ->
                Seq.append (of_item key) (of_item value))
            @@ List.to_seq l)
    | Item.Tag (t, c) -> Seq.cons (Tag t) (of_item c)
    | Item.Float f -> Seq.return (Float f)
    | Item.Simple n -> Seq.return (Simple n)
    | Item.Double f -> Seq.return (Double f)

  let indefinite_length_array items =
    Seq.(cons (Array None) (append (concat_map of_item items) (return Break)))

  let indefinite_length_map key_values =
    Seq.(
      cons (Map None)
        (append
           (concat_map
              (fun (key, value) -> Seq.append (of_item key) (of_item value))
              key_values)
           (return Break)))

  exception NotWellFormed of string

  let not_well_formed msg = raise @@ NotWellFormed msg

  module Parser = struct
    let items_to_map items =
      let result = ref [] in
      let n = List.length items in
      for i = 0 to (n / 2) - 1 do
        result :=
          (List.nth items (2 * i), List.nth items ((2 * i) + 1)) :: !result
      done;
      List.rev !result

    let rec items k result n seq =
      match n with
      | 0 -> k (List.rev result) seq
      | n ->
          item
            (fun item seq -> items k (item :: result) (n - 1) seq)
            (fun _ ->
              not_well_formed "unexpected break encountered while parsing array")
            (fun () ->
              not_well_formed "unexpected end of input while parsing array")
            seq

    and items_until_break k result seq =
      item
        (fun item seq -> items_until_break k (item :: result) seq)
        (fun seq -> k (List.rev result) seq)
        (fun () ->
          not_well_formed
            "unexpected end of input while parsing indefinite-length array")
        seq

    and bytes k buffer n seq =
      match n with
      | 0 -> k (Buffer.contents buffer) seq
      | n -> (
          match Seq.uncons seq with
          | Some (Byte b, seq) ->
              Buffer.add_char buffer b;
              bytes k buffer (n - 1) seq
          | Some _ -> not_well_formed "expecting a byte"
          | None ->
              not_well_formed "unexpected end of input while reading bytes")

    and item k k_break k_end seq =
      match Seq.uncons seq with
      | Some (Null, seq) -> k Item.Null seq
      | Some (Undefined, seq) -> k Item.Undefined seq
      | Some (Bool b, seq) -> k (Item.Bool b) seq
      | Some (Integer z, seq) -> k (Item.Integer z) seq
      | Some (ByteString len, seq) ->
          bytes
            (fun bytes seq -> k (Item.ByteString bytes) seq)
            (Buffer.create len) len seq
      | Some (TextString len, seq) ->
          bytes
            (fun bytes seq -> k (Item.TextString bytes) seq)
            (Buffer.create len) len seq
      | Some (Byte _, _) -> not_well_formed "unexpected byte"
      | Some (Array (Some n), seq) ->
          items (fun items seq -> k (Item.Array items) seq) [] n seq
      | Some (Array None, seq) ->
          items_until_break (fun items seq -> k (Item.Array items) seq) [] seq
      | Some (Map (Some n), seq) ->
          items
            (fun items seq -> k (Item.Map (items_to_map items)) seq)
            [] (n * 2) seq
      | Some (Map None, seq) ->
          items_until_break
            (fun items seq -> k (Item.Map (items_to_map items)) seq)
            [] seq
      | Some (Tag t, seq) when Z.(equal t ~$2) || Z.(equal t ~$3) ->
          item
            (fun item seq ->
              match item with
              | Item.ByteString bytes ->
                  let len = String.length bytes in
                  let bytes_le =
                    String.init len (fun i -> bytes.[len - 1 - i])
                  in
                  let z =
                    if Z.(equal t ~$2) then Z.(of_bits bytes_le)
                    else Z.(~- (~$1) - of_bits bytes_le)
                  in
                  k (Item.Integer z) seq
              | _ -> k (Item.Tag (t, item)) seq)
            (fun _ -> not_well_formed "unexpected break encountered")
            (fun () -> not_well_formed "unexpected end of input")
            seq
      | Some (Tag t, seq) ->
          item
            (fun item seq -> k (Item.Tag (t, item)) seq)
            (fun _ -> not_well_formed "unexpected break encountered")
            (fun () -> not_well_formed "unexpected end of input")
            seq
      | Some (Simple n, seq) -> k (Item.Simple n) seq
      | Some (Break, seq) -> k_break seq
      | Some (Float f, seq) -> k (Item.Float f) seq
      | Some (Double f, seq) -> k (Item.Double f) seq
      | None -> k_end ()
  end

  let rec to_items seq =
    Parser.item
      (fun item seq -> Seq.cons item (fun () -> to_items seq ()))
      (fun _ -> not_well_formed "unexpected break encountered")
      (fun () -> Seq.empty)
      seq

  let to_item seq =
    Parser.item
      (fun item seq -> (item, seq))
      (fun _ -> not_well_formed "unexpected break encountered")
      (fun () -> not_well_formed "unexpected end of input")
      seq

  module Writer = struct
    let initial_byte major_type ai =
      assert (major_type >= 0 && major_type < 8);
      assert (ai >= 0 && ai < 32);
      char_of_int @@ ((major_type lsl 5) lor ai)

    let head_z major_type arg =
      let argument_numbytes = z_numbytes arg in

      let argument_bytes size =
        let s = Z.to_bits arg in
        Seq.(
          append
            (* pad with zeroes *)
            (take (size - argument_numbytes) @@ repeat (Char.chr 0))
            (* binary representation of arg in big-endian (reveresed) *)
            (init argument_numbytes (fun i -> s.[argument_numbytes - 1 - i])))
      in

      if Z.(lt arg ~$24) then
        Seq.return (initial_byte major_type (Z.to_int arg))
      else if argument_numbytes = 1 then
        Seq.cons (initial_byte major_type 24) (argument_bytes 1)
      else if argument_numbytes = 2 then
        Seq.cons (initial_byte major_type 25) (argument_bytes 2)
      else if argument_numbytes <= 4 then
        Seq.cons (initial_byte major_type 26) (argument_bytes 4)
      else if argument_numbytes <= 8 then
        Seq.cons (initial_byte major_type 27) (argument_bytes 8)
      else failwith "argument too big for head"

    let integer n =
      if Z.sign n >= 0 then head_z 0 n else head_z 1 Z.(~- (~$1) - n)

    let array size_opt =
      match size_opt with
      | Some l -> head_z 4 Z.(of_int l)
      | None -> Seq.return @@ initial_byte 4 31

    let map size_opt =
      match size_opt with
      | Some l -> head_z 5 Z.(of_int l)
      | None -> Seq.return @@ initial_byte 5 31

    let float f =
      let bytes = Bytes.create 4 in
      Bytes.set_int32_be bytes 0 (Int32.bits_of_float f);
      Seq.cons (initial_byte 7 26) (Bytes.to_seq bytes)

    let double f =
      let bytes = Bytes.create 8 in
      Bytes.set_int64_be bytes 0 (Int64.bits_of_float f);
      Seq.cons (initial_byte 7 27) (Bytes.to_seq bytes)

    let signal signal =
      match signal with
      | Bool false -> Seq.return @@ initial_byte 7 20
      | Bool true -> Seq.return @@ initial_byte 7 21
      | Null -> Seq.return @@ initial_byte 7 22
      | Undefined -> Seq.return @@ initial_byte 7 23
      | ByteString len -> head_z 2 (Z.of_int len)
      | TextString len -> head_z 3 (Z.of_int len)
      | Byte b -> Seq.return b
      | Integer i -> integer i
      | Array size_opt -> array size_opt
      | Map size_opt -> map size_opt
      | Tag t -> head_z 6 t
      | Simple n when n < 20 || (32 <= n && n <= 255) -> head_z 7 (Z.of_int n)
      | Simple _ -> not_well_formed "will not encode reserved simple value"
      | Float f -> float f
      | Double f -> double f
      | Break -> Seq.return @@ initial_byte 7 31
  end

  let write signals = Seq.concat_map Writer.signal signals

  module Reader = struct
    let initial_byte seq =
      match Seq.uncons seq with
      | Some (initial_byte, seq) ->
          let b = Char.code initial_byte in
          Some ((b lsr 5, b land 0b11111), seq)
      | None -> None

    let rec read_reversed_bytes n ?(result = []) seq =
      if n = 0 then (List.to_seq result |> String.of_seq, seq)
      else
        match Seq.uncons seq with
        | Some (v, seq) -> read_reversed_bytes (n - 1) ~result:(v :: result) seq
        | None ->
            not_well_formed
              "unexpected end of input while reading reversed bytes"

    let read_bytes n seq =
      let r, seq = read_reversed_bytes n seq in
      let l = String.length r in
      (String.init l (fun i -> r.[l - 1 - i]), seq)

    let read_z_argument ai seq =
      let reversed_bytes, seq =
        if ai = 24 then read_reversed_bytes 1 seq
        else if ai = 25 then read_reversed_bytes 2 seq
        else if ai = 26 then read_reversed_bytes 4 seq
        else if ai = 27 then read_reversed_bytes 8 seq
        else failwith "invalid additional information"
      in
      (* In CBOR integers are always stored in big-endian, whereas Z expects little-endian *)
      (Z.of_bits reversed_bytes, seq)

    let unsigned_int k ai seq =
      if ai < 24 then Seq.cons (Integer (Z.of_int ai)) (k seq)
      else if ai < 28 then
        let int, seq = read_z_argument ai seq in
        Seq.cons (Integer int) (k seq)
      else if ai < 31 then
        not_well_formed
          "reserved values encountered while handling major type 0"
      else
        not_well_formed
          "unexpected additional information while handling major type 0"

    let negative_int k ai seq =
      if ai < 24 then Seq.cons (Integer Z.(neg @@ succ @@ of_int ai)) (k seq)
      else if ai < 28 then
        let uint, seq = read_z_argument ai seq in
        let int = Z.(neg @@ succ @@ uint) in
        Seq.cons (Integer int) (k seq)
      else if ai < 31 then
        not_well_formed
          "reserved values encountered while handling major type 1"
      else
        not_well_formed
          "unexpected additional information while handling major type 1"

    let rec bytes k n seq =
      match n with
      | 0 -> k seq
      | n -> (
          match Seq.uncons seq with
          | Some (b, seq) ->
              Seq.cons (Byte b) (fun () -> bytes k (n - 1) seq ())
          | None ->
              not_well_formed "unexpected end of input while reading bytes")

    let bytestring k ai seq =
      if ai < 24 then Seq.cons (ByteString ai) (bytes k ai seq)
      else if ai < 28 then
        let len_z, seq = read_z_argument ai seq in
        let len = Z.to_int len_z in
        Seq.cons (ByteString len) (bytes k len seq)
      else if ai < 31 then
        not_well_formed
          "reserved values encountered while handling major type 2"
      else if ai = 31 then failwith "indefinite-length bytestring not supported"
      else
        not_well_formed
          "unexpected additional information while handling major type 2"

    let textstring k ai seq =
      if ai < 24 then Seq.cons (TextString ai) (bytes k ai seq)
      else if ai < 28 then
        let len_z, seq = read_z_argument ai seq in
        let len = Z.to_int len_z in
        Seq.cons (TextString len) (bytes k len seq)
      else if ai < 31 then
        not_well_formed
          "reserved values encountered while handling major type 3"
      else if ai = 31 then failwith "indefinite-length textstring not supported"
      else
        not_well_formed
          "unexpected additional information while handling major type 3"

    let array k ai seq =
      if ai < 24 then Seq.cons (Array (Some ai)) (k seq)
      else if ai < 28 then
        let len_z, seq = read_z_argument ai seq in
        let len = Z.to_int len_z in
        Seq.cons (Array (Some len)) (k seq)
      else if ai < 31 then
        not_well_formed
          "reserved values encountered while handling major type 4"
      else if ai = 31 then Seq.cons (Array None) (k seq)
      else
        not_well_formed
          "unexpected additional information while handling major type 4"

    let map k ai seq =
      if ai < 24 then Seq.cons (Map (Some ai)) (k seq)
      else if ai < 28 then
        let len_z, seq = read_z_argument ai seq in
        let len = Z.to_int len_z in
        Seq.cons (Map (Some len)) (k seq)
      else if ai < 31 then
        not_well_formed
          "reserved values encountered while handling major type 5"
      else if ai = 31 then Seq.cons (Map None) (k seq)
      else
        not_well_formed
          "unexpected additional information while handling major type 5"

    let tag k ai seq =
      if ai < 24 then Seq.cons (Tag (Z.of_int ai)) (k seq)
      else if ai < 28 then
        let tag, seq = read_z_argument ai seq in
        Seq.cons (Tag tag) (k seq)
      else if ai < 31 then
        not_well_formed
          "reserved values encountered while handling major type 6"
      else if ai = 31 then Seq.cons (Array None) (k seq)
      else
        not_well_formed
          "unexpected additional information while handling major type 6"

    let float k seq =
      let bytes, seq = read_bytes 4 seq in
      let f =
        Int32.float_of_bits @@ Bytes.get_int32_be (Bytes.of_string bytes) 0
      in
      Seq.cons (Float f) (k seq)

    let double k seq =
      let bytes, seq = read_bytes 8 seq in
      let f =
        Int64.float_of_bits @@ Bytes.get_int64_be (Bytes.of_string bytes) 0
      in
      Seq.cons (Double f) (k seq)

    let rec signals seq () =
      match initial_byte seq with
      | Some ((major_type, ai), seq) -> (
          match major_type with
          | 0 -> unsigned_int signals ai seq ()
          | 1 -> negative_int signals ai seq ()
          | 2 -> bytestring signals ai seq ()
          | 3 -> textstring signals ai seq ()
          | 4 -> array signals ai seq ()
          | 5 -> map signals ai seq ()
          | 6 -> tag signals ai seq ()
          | 7 -> (
              match ai with
              | 20 -> Seq.Cons (Bool false, signals seq)
              | 21 -> Seq.Cons (Bool true, signals seq)
              | 22 -> Seq.Cons (Null, signals seq)
              | 23 -> Seq.Cons (Undefined, signals seq)
              | 24 ->
                  let z, seq = read_z_argument 24 seq in
                  Seq.Cons (Simple (Z.to_int z), signals seq)
              | 25 -> failwith "IEEE 754 Half-Precision Float not supported"
              | 26 -> float signals seq ()
              | 27 -> double signals seq ()
              | 31 -> Seq.Cons (Break, signals seq)
              | n when n < 20 -> Seq.Cons (Simple n, signals seq)
              | _ ->
                  not_well_formed
                    "encountered reserved additional information for major \
                     type 7")
          | _ -> failwith "invalid major type")
      | None -> Seq.Nil
  end

  let read seq = Reader.signals seq
end

let write item = Signal.of_item item |> Signal.write

let read chars =
  let item, _seq = Signal.read chars |> Signal.to_item in
  item