Source file ptyp_of_type.ml

open Std
open Typedtree
open Types

let var_of_id id = Location.mknoloc @@ Ident.name id

type signature_elt =
  | Item of Types.signature_item
  | Type of Asttypes.rec_flag * Parsetree.type_declaration list

let rec module_type =
  let open Ast_helper in
  function
  | Mty_for_hole -> failwith "Holes are not allowed in module types"
  | Mty_signature signature_items -> Mty.signature @@ signature signature_items
  | Mty_ident path ->
    Ast_helper.Mty.ident (Location.mknoloc (Untypeast.lident_of_path path))
  | Mty_alias path ->
    Ast_helper.Mty.alias (Location.mknoloc (Untypeast.lident_of_path path))
  | Mty_functor (param, type_out) ->
    let param =
      match param with
      | Unit -> Parsetree.Unit
      | Named (id, type_in) ->
        Parsetree.Named
          (Location.mknoloc (Option.map ~f:Ident.name id), module_type type_in)
    in
    let out = module_type type_out in
    Mty.functor_ param out

and core_type type_expr =
  let open Ast_helper in
  match Types.get_desc type_expr with
  | Tvar None | Tunivar None -> Typ.any ()
  | Tvar (Some s) | Tunivar (Some s) -> Typ.var s
  | Tarrow (label, type_expr, type_expr_out, _commutable) ->
    Typ.arrow label (core_type type_expr) (core_type type_expr_out)
  | Ttuple type_exprs -> Typ.tuple @@ List.map ~f:core_type type_exprs
  | Tconstr (path, type_exprs, _abbrev) ->
    let loc = Untypeast.lident_of_path path |> Location.mknoloc in
    Typ.constr loc @@ List.map ~f:core_type type_exprs
  | Tobject (type_expr, _class_) ->
    let rec aux acc type_expr =
      match get_desc type_expr with
      | Tnil -> (acc, Asttypes.Closed)
      | Tvar None | Tunivar None -> (acc, Asttypes.Open)
      | Tfield ("*dummy method*", _, _, fields) -> aux acc fields
      | Tfield (name, _, type_expr, fields) ->
        let open Ast_helper in
        let core_type = core_type type_expr in
        let core_type = Of.tag (Location.mknoloc name) core_type in

        aux (core_type :: acc) fields
      | _ ->
        failwith
        @@ Format.asprintf "Unexpected type constructor in fields list: %a"
             Printtyp.type_expr type_expr
    in
    let fields, closed = aux [] type_expr in
    Typ.object_ fields closed
  | Tfield _ -> failwith "Found object field outside of object."
  | Tnil -> Typ.object_ [] Closed
  | Tlink type_expr | Tsubst (type_expr, _) -> core_type type_expr
  | Tvariant row ->
    let row_fields = row_fields row in
    let row_closed = row_closed row in
    let field (label, row_field) =
      let label = Location.mknoloc label in
      match row_field_repr row_field with
      | Rpresent None | Reither (true, _, _) -> Rf.tag label true []
      | Rpresent (Some type_expr) ->
        let core_type = core_type type_expr in
        Rf.tag label false [ core_type ]
      | Reither (false, type_exprs, _) ->
        Rf.tag label false @@ List.map ~f:core_type type_exprs
      | Rabsent -> assert false
    in
    let closed = if row_closed then Asttypes.Closed else Asttypes.Open in
    let fields = List.map ~f:field row_fields in
    (* TODO NOT ALWAYS NONE *)
    Typ.variant fields closed None
  | Tpoly (type_expr, type_exprs) ->
    let names =
      List.map
        ~f:(fun v ->
          match get_desc v with
          | Tunivar (Some name) | Tvar (Some name) -> mknoloc name
          | _ -> failwith "poly: not a var")
        type_exprs
    in
    Typ.poly names @@ core_type type_expr
  | Tpackage (path, lids_type_exprs) ->
    let loc = mknoloc (Untypeast.lident_of_path path) in
    let args =
      List.map lids_type_exprs ~f:(fun (id, t) -> (mknoloc id, core_type t))
    in
    Typ.package loc args

and modtype_declaration id { mtd_type; mtd_attributes; _ } =
  Ast_helper.Mtd.mk ~attrs:mtd_attributes
    ?typ:(Option.map ~f:module_type mtd_type)
    (var_of_id id)

and module_declaration id { md_type; md_attributes; _ } =
  let name = Location.mknoloc (Some (Ident.name id)) in
  Ast_helper.Md.mk ~attrs:md_attributes name @@ module_type md_type

and extension_constructor id { ext_args; ext_ret_type; ext_attributes; _ } =
  Ast_helper.Te.decl ~attrs:ext_attributes
    ~args:(constructor_arguments ext_args)
    ?res:(Option.map ~f:core_type ext_ret_type)
    (var_of_id id)

and value_description id { val_type; val_kind = _; val_loc; val_attributes; _ }
    =
  let type_ = core_type val_type in
  { Parsetree.pval_name = var_of_id id;
    pval_type = type_;
    pval_prim = [];
    pval_attributes = val_attributes;
    pval_loc = val_loc
  }

and label_declaration { ld_id; ld_mutable; ld_type; ld_attributes; _ } =
  Ast_helper.Type.field ~attrs:ld_attributes ~mut:ld_mutable (var_of_id ld_id)
    (core_type ld_type)

and constructor_arguments = function
  | Cstr_tuple type_exprs ->
    Parsetree.Pcstr_tuple (List.map ~f:core_type type_exprs)
  | Cstr_record label_decls ->
    Parsetree.Pcstr_record (List.map ~f:label_declaration label_decls)

and constructor_declaration { cd_id; cd_args; cd_res; cd_attributes; _ } =
  Ast_helper.Type.constructor ~attrs:cd_attributes
    ~args:(constructor_arguments cd_args)
    ?res:(Option.map ~f:core_type cd_res)
  @@ var_of_id cd_id

and type_declaration id
    { type_params;
      type_variance;
      type_manifest;
      type_kind;
      type_attributes;
      type_private;
      _
    } =
  let params =
    List.map2 type_params type_variance ~f:(fun type_ variance ->
        let core_type = core_type type_ in
        let pos, neg, inj = Types.Variance.get_lower variance in
        let v =
          if pos then Asttypes.Covariant
          else if neg then Contravariant
          else NoVariance
        in
        let i = if inj then Asttypes.Injective else NoInjectivity in
        (core_type, (v, i)))
  in
  let kind =
    match type_kind with
    | Type_abstract _ -> Parsetree.Ptype_abstract
    | Type_open -> Ptype_open
    | Type_variant (constrs, _) ->
      Ptype_variant (List.map ~f:constructor_declaration constrs)
    | Type_record (labels, _repr) ->
      Ptype_record (List.map ~f:label_declaration labels)
  in
  let manifest = Option.map ~f:core_type type_manifest in
  Ast_helper.Type.mk ~attrs:type_attributes ~params ~kind ~priv:type_private
    ?manifest (var_of_id id)

and signature_item (str_item : Types.signature_item) =
  let open Ast_helper in
  match str_item with
  | Sig_value (id, vd, _visibility) ->
    let vd = value_description id vd in
    Sig.value vd
  | Sig_type (id, type_decl, rec_flag, _visibility) ->
    let rec_flag =
      match rec_flag with
      | Trec_first -> Asttypes.Recursive
      | Trec_next -> Asttypes.Recursive
      | Trec_not -> Nonrecursive
    in
    (* mutually recursive types are really handled by [signature] *)
    Sig.type_ rec_flag [ type_declaration id type_decl ]
  | Sig_modtype (id, modtype_decl, _visibility) ->
    Sig.modtype @@ modtype_declaration id modtype_decl
  | Sig_module (id, _, mod_decl, _, _) ->
    Sig.module_ @@ module_declaration id mod_decl
  | Sig_typext (id, ext_constructor, _, _) ->
    let ext =
      Te.mk
        (Location.mknoloc @@ Longident.Lident (Ident.name id))
        [ extension_constructor id ext_constructor ]
    in
    Sig.type_extension ext
  | Sig_class_type (id, _, _, _) ->
    let str =
      Format.asprintf
        "Construct does not handle class types yet. Please replace this \
         comment by [%s]'s definition."
        (Ident.name id)
    in
    Sig.text [ Docstrings.docstring str Location.none ] |> List.hd
  | Sig_class (id, _, _, _) ->
    let str =
      Format.asprintf
        "Construct does not handle classes yet. Please replace this comment by \
         [%s]'s definition."
        (Ident.name id)
    in
    Sig.text [ Docstrings.docstring str Location.none ] |> List.hd

and signature (items : Types.signature_item list) =
  List.map (group_items items) ~f:(function
    | Item item -> signature_item item
    | Type (rec_flag, type_decls) -> Ast_helper.Sig.type_ rec_flag type_decls)

and group_items (items : Types.signature_item list) =
  let rec read_type type_acc items =
    match items with
    | Sig_type (id, type_decl, Trec_next, _) :: rest ->
      let td = type_declaration id type_decl in
      read_type (td :: type_acc) rest
    | _ -> (List.rev type_acc, items)
  in
  let rec group acc items =
    match items with
    | Sig_type (id, type_decl, Trec_first, _) :: rest ->
      let type_, rest = read_type [ type_declaration id type_decl ] rest in
      group (Type (Asttypes.Recursive, type_) :: acc) rest
    | Sig_type (id, type_decl, Trec_not, _) :: rest ->
      let type_, rest = read_type [ type_declaration id type_decl ] rest in
      group (Type (Asttypes.Nonrecursive, type_) :: acc) rest
    | (Sig_class _ as item) :: _ :: _ :: rest -> group (Item item :: acc) rest
    | (Sig_class_type _ as item) :: _ :: rest -> group (Item item :: acc) rest
    | item :: rest -> group (Item item :: acc) rest
    | [] -> List.rev acc
  in
  group [] items