Source file pb_codegen_decode_binary.ml

module Ot = Pb_codegen_ocaml_type
module F = Pb_codegen_formatting

let sp = Pb_codegen_util.sp

let runtime_function_for_basic_type bt pk =
  match pk, bt with
  | Ot.Pk_varint false, Ot.Bt_int -> "Pbrt.Decoder.int_as_varint"
  | Ot.Pk_varint true, Ot.Bt_int -> "Pbrt.Decoder.int_as_zigzag"
  | Ot.Pk_varint false, Ot.Bt_int32 -> "Pbrt.Decoder.int32_as_varint"
  | Ot.Pk_varint true, Ot.Bt_int32 -> "Pbrt.Decoder.int32_as_zigzag"
  | Ot.Pk_varint false, Ot.Bt_uint32 -> "Pbrt.Decoder.uint32_as_varint"
  | Ot.Pk_varint true, Ot.Bt_uint32 -> "Pbrt.Decoder.uint32_as_zigzag"
  | Ot.Pk_varint false, Ot.Bt_int64 -> "Pbrt.Decoder.int64_as_varint"
  | Ot.Pk_varint true, Ot.Bt_int64 -> "Pbrt.Decoder.int64_as_zigzag"
  | Ot.Pk_varint false, Ot.Bt_uint64 -> "Pbrt.Decoder.uint64_as_varint"
  | Ot.Pk_varint true, Ot.Bt_uint64 -> "Pbrt.Decoder.uint64_as_zigzag"
  | Ot.Pk_bits32, Ot.Bt_int32 -> "Pbrt.Decoder.int32_as_bits32"
  | Ot.Pk_bits64, Ot.Bt_int64 -> "Pbrt.Decoder.int64_as_bits64"
  | Ot.Pk_bits32, Ot.Bt_uint32 -> "Pbrt.Decoder.uint32_as_bits32"
  | Ot.Pk_bits64, Ot.Bt_uint64 -> "Pbrt.Decoder.uint64_as_bits64"
  | Ot.Pk_varint false, Ot.Bt_bool -> "Pbrt.Decoder.bool"
  | Ot.Pk_bits32, Ot.Bt_float -> "Pbrt.Decoder.float_as_bits32"
  | Ot.Pk_bits64, Ot.Bt_float -> "Pbrt.Decoder.float_as_bits64"
  | Ot.Pk_bits32, Ot.Bt_int -> "Pbrt.Decoder.int_as_bits32"
  | Ot.Pk_bits64, Ot.Bt_int -> "Pbrt.Decoder.int_as_bits64"
  | Ot.Pk_bytes, Ot.Bt_string -> "Pbrt.Decoder.string"
  | Ot.Pk_bytes, Ot.Bt_bytes -> "Pbrt.Decoder.bytes"
  | _ -> failwith "Invalid decoding/OCaml type combination"

let runtime_function_for_wrapper_type { Ot.wt_type; wt_pk } =
  match wt_type, wt_pk with
  | Ot.Bt_float, Ot.Pk_bits64 -> "Pbrt.Decoder.wrapper_double_value"
  | Ot.Bt_float, Ot.Pk_bits32 -> "Pbrt.Decoder.wrapper_float_value"
  | Ot.Bt_int64, Ot.Pk_varint _ -> "Pbrt.Decoder.wrapper_int64_value"
  | Ot.Bt_int32, Ot.Pk_varint _ -> "Pbrt.Decoder.wrapper_int32_value"
  | Ot.Bt_bool, Ot.Pk_varint _ -> "Pbrt.Decoder.wrapper_bool_value"
  | Ot.Bt_string, Ot.Pk_bytes -> "Pbrt.Decoder.wrapper_string_value"
  | Ot.Bt_bytes, Ot.Pk_bytes -> "Pbrt.Decoder.wrapper_bytes_value"
  | _ -> assert false

let decode_field_expression field_type pk : string =
  match field_type with
  | Ot.Ft_user_defined_type t ->
    let f_name =
      let function_prefix = "decode_pb" in
      Pb_codegen_util.function_name_of_user_defined ~function_prefix t
    in
    (match t.Ot.udt_type with
    | `Message -> f_name ^ " (Pbrt.Decoder.nested d)"
    | `Enum -> f_name ^ " d")
  | Ot.Ft_unit -> "Pbrt.Decoder.empty_nested d"
  | Ot.Ft_basic_type bt -> runtime_function_for_basic_type bt pk ^ " d"
  | Ot.Ft_wrapper_type wt -> runtime_function_for_wrapper_type wt ^ " d"

let pbrt_payload_kind payload_kind is_packed =
  if is_packed then
    "Bytes"
  else
    Pb_codegen_util.string_of_payload_kind ~capitalize:() payload_kind false

let gen_field_common sc encoding_number payload_kind message_name
    ?(is_packed = false) f =
  F.linep sc "| Some (%i, Pbrt.%s) -> begin" encoding_number
    (pbrt_payload_kind payload_kind is_packed);
  F.sub_scope sc f;
  F.line sc "end";
  F.linep sc "| Some (%i, pk) -> " encoding_number;
  F.linep sc "  Pbrt.Decoder.unexpected_payload \"%s\" pk"
    (sp "Message(%s), field(%i)" message_name encoding_number)

let gen_rft_nolabel sc r_name rf_label (field_type, encoding_number, pk) =
  gen_field_common sc encoding_number pk r_name (fun sc ->
      F.linep sc "v.%s <- %s;" rf_label (decode_field_expression field_type pk))

(* return the variable name used for keeping track if a required
 * field has been set during decoding. *)
let is_set_variable_name rf_label = sp "%s_is_set" rf_label

let gen_rft_required sc r_name rf_label (field_type, encoding_number, pk, _) =
  gen_field_common sc encoding_number pk r_name (fun sc ->
      F.linep sc "v.%s <- %s; %s := true;" rf_label
        (decode_field_expression field_type pk)
        (is_set_variable_name rf_label))

let gen_rft_optional sc r_name rf_label optional_field =
  let field_type, encoding_number, pk, _ = optional_field in
  gen_field_common sc encoding_number pk r_name (fun sc ->
      F.linep sc "v.%s <- Some (%s);" rf_label
        (decode_field_expression field_type pk))

let gen_rft_repeated sc r_name rf_label repeated_field =
  let rt, field_type, encoding_number, pk, is_packed = repeated_field in
  match rt, is_packed with
  | Ot.Rt_list, false ->
    gen_field_common sc encoding_number pk r_name ~is_packed (fun sc ->
        F.linep sc "v.%s <- (%s) :: v.%s;" rf_label
          (decode_field_expression field_type pk)
          rf_label)
  | Ot.Rt_repeated_field, false ->
    gen_field_common sc encoding_number pk r_name ~is_packed (fun sc ->
        F.linep sc "Pbrt.Repeated_field.add (%s) v.%s; "
          (decode_field_expression field_type pk)
          rf_label)
  | Ot.Rt_list, true ->
    gen_field_common sc encoding_number pk r_name ~is_packed (fun sc ->
        F.linep sc "v.%s <- Pbrt.Decoder.packed_fold (fun l d -> (%s)::l) [] d;"
          rf_label
          (decode_field_expression field_type pk))
  | Ot.Rt_repeated_field, true ->
    gen_field_common sc encoding_number pk r_name ~is_packed (fun sc ->
        F.line sc "Pbrt.Decoder.packed_fold (fun () d -> ";
        F.sub_scope sc (fun sc ->
            F.linep sc "Pbrt.Repeated_field.add (%s) v.%s;"
              (decode_field_expression field_type pk)
              rf_label);
        F.line sc ") () d;")

let gen_rft_associative sc r_name rf_label associative_field =
  let at, encoding_number, (key_type, key_pk), (value_type, value_pk) =
    associative_field
  in

  let decode_key_f = runtime_function_for_basic_type key_type key_pk in

  (* Because key can never be nested we can assign the decoding function
   * directly rather wrapping up in a closure like for the value
   * below
   *)
  (* TODO enhancement
   * For the value decoding function passed as an argument to
   * [Pbrt.Decoder.map_entry] it's not always the case that it would
   * require nesting. In the case it does not neeed a nested decoder
   * we can avoid creating a closure and therefore improving
   * the performance. *)
  gen_field_common sc encoding_number Ot.Pk_bytes r_name (fun sc ->
      F.line sc "let decode_value = (fun d ->";
      F.sub_scope sc (fun sc ->
          F.line sc @@ decode_field_expression value_type value_pk);
      F.line sc ") in";
      let decode_expression =
        sp "(Pbrt.Decoder.map_entry d ~decode_key:%s ~decode_value)"
          decode_key_f
      in

      match at with
      | Ot.At_list ->
        F.linep sc "v.%s <- (" rf_label;
        F.sub_scope sc (fun sc ->
            F.linep sc "%s::v.%s;" decode_expression rf_label);
        F.line sc ");"
      | Ot.At_hashtable ->
        F.linep sc "let a, b = %s in" decode_expression;
        F.linep sc "Hashtbl.add v.%s a b;" rf_label)

let gen_rft_variant sc r_name rf_label { Ot.v_constructors; _ } =
  List.iter
    (fun variant_constructor ->
      let {
        Ot.vc_constructor;
        vc_field_type;
        vc_encoding_number;
        vc_payload_kind = pk;
        vc_options = _;
      } =
        variant_constructor
      in

      gen_field_common sc vc_encoding_number pk r_name (fun sc ->
          match vc_field_type with
          | Ot.Vct_nullary ->
            F.line sc "Pbrt.Decoder.empty_nested d;";
            F.linep sc "v.%s <- %s;" rf_label vc_constructor
          | Ot.Vct_non_nullary_constructor field_type ->
            F.linep sc "v.%s <- %s (%s);" rf_label vc_constructor
              (decode_field_expression field_type pk)))
    v_constructors

let gen_record ?and_ { Ot.r_name; r_fields } sc =
  (* list fields have a special treatement when decoding since each new element
     of a repeated field is appended to the front of the list. In order
     to retreive the right order efficiently we reverse all the repeated field
     lists values when the message is done being decoded. *)
  let all_lists =
    List.fold_left
      (fun acc { Ot.rf_label; rf_field_type; _ } ->
        match rf_field_type with
        | Ot.Rft_repeated (Ot.Rt_list, _, _, _, _) -> rf_label :: acc
        | Ot.Rft_associative (Ot.At_list, _, _, _) -> rf_label :: acc
        | _ -> acc)
      [] r_fields
  in

  let all_required_rf_labels =
    List.fold_left
      (fun acc { Ot.rf_label; rf_field_type; _ } ->
        match rf_field_type with
        | Ot.Rft_required _ -> rf_label :: acc
        | _ -> acc)
      [] r_fields
  in

  let mutable_record_name = Pb_codegen_util.mutable_record_name r_name in

  F.linep sc "%s decode_pb_%s d =" (Pb_codegen_util.let_decl_of_and and_) r_name;
  F.sub_scope sc (fun sc ->
      F.linep sc "let v = default_%s () in" mutable_record_name;
      F.line sc "let continue__= ref true in";

      (* Add the is_set_<field_name> boolean variable which keeps track
       * of whether a required field is set during the decoding. *)
      List.iter
        (fun rf_label ->
          F.linep sc "let %s = ref false in" (is_set_variable_name rf_label))
        all_required_rf_labels;

      (* Decoding is done with recursively (tail - recursive). The
       * function loop iterate over all fields returned by the Protobuf
       * runtime. *)
      F.line sc "while !continue__ do";
      F.sub_scope sc (fun sc ->
          F.line sc "match Pbrt.Decoder.key d with";

          (* termination condition *)
          F.line sc "| None -> (";
          F.sub_scope sc (fun sc ->
              List.iter
                (fun field_name ->
                  F.linep sc "v.%s <- List.rev v.%s;" field_name field_name)
                all_lists);
          F.line sc "); continue__ := false";

          (* compare the decoded field with the one defined in the
           * .proto file. Unknown fields are ignored. *)
          List.iter
            (fun { Ot.rf_label; rf_field_type; _ } ->
              match rf_field_type with
              | Ot.Rft_nolabel x -> gen_rft_nolabel sc r_name rf_label x
              | Ot.Rft_required x -> gen_rft_required sc r_name rf_label x
              | Ot.Rft_optional x -> gen_rft_optional sc r_name rf_label x
              | Ot.Rft_repeated x -> gen_rft_repeated sc r_name rf_label x
              | Ot.Rft_associative x -> gen_rft_associative sc r_name rf_label x
              | Ot.Rft_variant x -> gen_rft_variant sc r_name rf_label x)
            r_fields;
          F.line sc
            ("| Some (_, payload_kind) -> " ^ "Pbrt.Decoder.skip d payload_kind"));
      F.line sc "done;";

      (* Add the check to see if all required fields are set if not
       * a Protobuf.Decoder.Failure exception is raised *)
      List.iter
        (fun rf_label ->
          F.linep sc
            "begin if not !%s then Pbrt.Decoder.missing_field \"%s\" end;"
            (is_set_variable_name rf_label)
            rf_label)
        all_required_rf_labels;

      F.line sc "({";
      F.sub_scope sc (fun sc ->
          List.iter
            (fun { Ot.rf_label; _ } ->
              F.linep sc "%s = v.%s;" rf_label rf_label)
            r_fields);
      F.linep sc "} : %s)" r_name)

let gen_unit ?and_ { Ot.er_name } sc =
  F.linep sc "%s decode_pb_%s d ="
    (Pb_codegen_util.let_decl_of_and and_)
    er_name;
  F.sub_scope sc (fun sc ->
      F.line sc "match Pbrt.Decoder.key d with";
      F.line sc "| None -> ();";
      F.line sc "| Some (_, pk) -> ";
      F.linep sc "  Pbrt.Decoder.unexpected_payload \"%s\" pk"
        (sp "Unexpected fields in empty message(%s)" er_name))

let gen_variant ?and_ { Ot.v_name; v_constructors } sc =
  let process_ctor sc variant_constructor =
    let {
      Ot.vc_constructor;
      vc_field_type;
      vc_encoding_number;
      vc_payload_kind = pk;
      vc_options = _;
    } =
      variant_constructor
    in

    match vc_field_type with
    | Ot.Vct_nullary ->
      F.linep sc "| Some (%i, _) -> begin " vc_encoding_number;
      F.sub_scope sc (fun sc ->
          F.line sc "Pbrt.Decoder.empty_nested d ;";
          F.linep sc "(%s : %s)" vc_constructor v_name);
      F.line sc "end"
    | Ot.Vct_non_nullary_constructor field_type ->
      F.linep sc "| Some (%i, _) -> (%s (%s) : %s) " vc_encoding_number
        vc_constructor
        (decode_field_expression field_type pk)
        v_name
  in

  F.linep sc "%s decode_pb_%s d = "
    (Pb_codegen_util.let_decl_of_and and_)
    v_name;
  F.sub_scope sc (fun sc ->
      F.linep sc "let rec loop () = ";
      F.sub_scope sc (fun sc ->
          F.linep sc "let ret:%s = match Pbrt.Decoder.key d with" v_name;

          F.sub_scope sc (fun sc ->
              F.linep sc "| None -> Pbrt.Decoder.malformed_variant \"%s\""
                v_name;
              List.iter (fun ctor -> process_ctor sc ctor) v_constructors;
              F.line sc "| Some (n, payload_kind) -> (";
              F.line sc "  Pbrt.Decoder.skip d payload_kind; ";
              F.line sc "  loop () ";
              F.line sc ")");

          F.line sc "in";
          F.line sc "ret");
      F.line sc "in";
      F.line sc "loop ()")

let gen_const_variant ?and_ { Ot.cv_name; cv_constructors } sc =
  F.linep sc "%s decode_pb_%s d = "
    (Pb_codegen_util.let_decl_of_and and_)
    cv_name;
  F.sub_scope sc (fun sc ->
      F.line sc "match Pbrt.Decoder.int_as_varint d with";
      List.iter
        (fun { Ot.cvc_name; cvc_binary_value; _ } ->
          F.linep sc "| %i -> (%s:%s)" cvc_binary_value cvc_name cv_name)
        cv_constructors;
      F.linep sc "| _ -> Pbrt.Decoder.malformed_variant \"%s\"" cv_name)

let gen_struct ?and_ t sc =
  let { Ot.spec; _ } = t in

  let has_encoded =
    match spec with
    | Ot.Record r ->
      gen_record ?and_ r sc;
      true
    | Ot.Variant v ->
      gen_variant ?and_ v sc;
      true
    | Ot.Const_variant v ->
      gen_const_variant ?and_ v sc;
      true
    | Ot.Unit u ->
      gen_unit ?and_ u sc;
      true
  in

  has_encoded

let gen_sig ?and_ t sc =
  let _ = and_ in

  let { Ot.spec; _ } = t in

  let f type_name =
    F.linep sc "val decode_pb_%s : Pbrt.Decoder.t -> %s" type_name type_name;
    F.linep sc
      ("(** [decode_pb_%s decoder] decodes a "
     ^^ "[%s] binary value from [decoder] *)")
      type_name type_name
  in

  let has_encoded =
    match spec with
    | Ot.Record { Ot.r_name; _ } ->
      f r_name;
      true
    | Ot.Variant { Ot.v_name; _ } ->
      f v_name;
      true
    | Ot.Const_variant { Ot.cv_name; _ } ->
      f cv_name;
      true
    | Ot.Unit { Ot.er_name; _ } ->
      f er_name;
      true
  in

  has_encoded

let ocamldoc_title = "Protobuf Decoding"
let requires_mutable_records = true

let plugin : Pb_codegen_plugin.t =
  let module P = struct
    let gen_sig = gen_sig
    let gen_struct = gen_struct
    let ocamldoc_title = ocamldoc_title
    let requires_mutable_records = requires_mutable_records
  end in
  (module P)