Source file ppx_sexp_conv_grammar.ml

open! Base
open! Ppxlib
open Ast_builder.Default

let copy =
  object
    inherit Ast_traverse.map
    method! location loc = { loc with loc_ghost = true }
    method! attributes _ = []
  end
;;

let unsupported ~loc string =
  Location.raise_errorf ~loc "sexp_grammar: %s are unsupported" string
;;

let ewith_tag ~loc ~key ~value grammar =
  [%expr { key = [%e key]; value = [%e value]; grammar = [%e grammar] }]
;;

let eno_tag ~loc grammar = [%expr No_tag [%e grammar]]
let etag ~loc with_tag = [%expr Tag [%e with_tag]]
let etagged ~loc with_tag = [%expr Tagged [%e with_tag]]

let tag_of_doc_comment ~loc comment =
  ( [%expr Ppx_sexp_conv_lib.Sexp_grammar.doc_comment_tag]
  , [%expr Atom [%e estring ~loc comment]] )
;;

module Tags = struct
  type t =
    { defined_using_tags : expression option
    ; defined_using_tag : (expression * expression) list
    }

  let get x ~tags ~tag =
    { defined_using_tags = Attribute.get tags x
    ; defined_using_tag = Attribute.get tag x |> Option.value ~default:[]
    }
  ;;
end

let rec with_tag_assoc_list grammar ~loc ~tags_expr ~wrap_tag ~wrap_tags =
  match tags_expr with
  | [%expr []] -> grammar
  | [%expr ([%e? key], [%e? value]) :: [%e? tags_expr]] ->
    wrap_tag
      ~loc
      (ewith_tag
         ~loc
         ~key
         ~value
         (with_tag_assoc_list grammar ~loc ~tags_expr ~wrap_tag ~wrap_tags))
  | _ -> wrap_tags grammar ~loc ~tags_expr
;;

let with_tags grammar ~wrap_tag ~wrap_tags ~loc ~(tags : Tags.t) ~comments =
  let tags_from_comments = List.map comments ~f:(tag_of_doc_comment ~loc) in
  let init =
    match tags.defined_using_tags with
    | None -> grammar
    | Some tags_expr -> with_tag_assoc_list grammar ~loc ~tags_expr ~wrap_tag ~wrap_tags
  in
  List.fold_right
    (List.concat [ tags_from_comments; tags.defined_using_tag ])
    ~init
    ~f:(fun (key, value) grammar -> wrap_tag ~loc (ewith_tag ~loc ~key ~value grammar))
;;

let with_tags_as_list grammar ~core_type ~loc ~tags ~comments =
  let wrap_tags grammar ~loc ~tags_expr =
    [%expr
      Sexplib0.Sexp_conv.sexp_grammar_with_tag_list
        ([%e grammar] : [%t core_type] Sexplib0.Sexp_grammar.with_tag_list)
        ~tags:[%e tags_expr]]
  in
  with_tags (eno_tag ~loc grammar) ~wrap_tag:etag ~wrap_tags ~loc ~tags ~comments
;;

let with_tags_as_grammar grammar ~loc ~tags ~comments =
  let wrap_tags grammar ~loc ~tags_expr =
    [%expr Sexplib0.Sexp_conv.sexp_grammar_with_tags [%e grammar] ~tags:[%e tags_expr]]
  in
  with_tags grammar ~wrap_tag:etagged ~wrap_tags ~loc ~tags ~comments
;;

let grammar_name name = name ^ "_sexp_grammar"
let tyvar_grammar_name name = grammar_name ("_'" ^ name)
let estr { loc; txt } = estring ~loc txt

let grammar_type ~loc core_type =
  [%type: [%t copy#core_type core_type] Sexplib0.Sexp_grammar.t]
;;

let abstract_grammar ~ctxt ~loc id =
  let module_name =
    ctxt |> Expansion_context.Deriver.code_path |> Code_path.fully_qualified_path
  in
  [%expr Any [%e estr { id with txt = String.concat ~sep:"." [ module_name; id.txt ] }]]
;;

let arrow_grammar ~loc = [%expr Sexplib0.Sexp_conv.fun_sexp_grammar.untyped]
let opaque_grammar ~loc = [%expr Sexplib0.Sexp_conv.opaque_sexp_grammar.untyped]
let any_grammar ~loc name = [%expr Any [%e estring ~loc name]]
let list_grammar ~loc expr = [%expr List [%e expr]]
let many_grammar ~loc expr = [%expr Many [%e expr]]
let fields_grammar ~loc expr = [%expr Fields [%e expr]]
let tyvar_grammar ~loc expr = [%expr Tyvar [%e expr]]
let recursive_grammar ~loc name args = [%expr Recursive ([%e name], [%e args])]

let tycon_grammar ~loc tycon_name params defns =
  [%expr Tycon ([%e tycon_name], [%e params], [%e defns])]
;;

let defns_type ~loc = [%type: Sexplib0.Sexp_grammar.defn Stdlib.List.t Stdlib.Lazy.t]

let untyped_grammar ~loc expr =
  match expr with
  | [%expr { untyped = [%e? untyped] }] -> untyped
  | _ -> [%expr [%e expr].untyped]
;;

let typed_grammar ~loc expr =
  match expr with
  | [%expr [%e? typed].untyped] -> typed
  | _ -> [%expr { untyped = [%e expr] }]
;;

let annotated_grammar ~loc expr core_type =
  pexp_constraint ~loc expr (grammar_type ~loc core_type)
;;

let defn_expr ~loc ~tycon ~tyvars ~grammar =
  [%expr { tycon = [%e tycon]; tyvars = [%e tyvars]; grammar = [%e grammar] }]
;;

let union_grammar ~loc exprs =
  match exprs with
  | [] -> [%expr Union []]
  | [ expr ] -> expr
  | _ -> [%expr Union [%e elist ~loc exprs]]
;;

let tuple_grammar ~loc exprs =
  List.fold_right exprs ~init:[%expr Empty] ~f:(fun expr rest ->
    [%expr Cons ([%e expr], [%e rest])])
;;

let atom_clause ~loc = [%expr Atom_clause]
let list_clause ~loc args = [%expr List_clause { args = [%e args] }]

module Variant_clause_type = struct
  type t =
    { name : label loc
    ; comments : string list
    ; tags : Tags.t
    ; clause_kind : expression
    }

  let to_grammar_expr { name; comments; tags; clause_kind } ~loc =
    [%expr { name = [%e estr name]; clause_kind = [%e clause_kind] }]
    |> with_tags_as_list
         ~loc:name.loc
         ~comments
         ~tags
         ~core_type:[%type: Sexplib0.Sexp_grammar.clause]
  ;;
end

let variant_grammars ~loc ~case_sensitivity ~clauses =
  match List.is_empty clauses with
  | true -> []
  | false ->
    let clause_exprs = List.map clauses ~f:(Variant_clause_type.to_grammar_expr ~loc) in
    let grammar =
      [%expr
        Variant
          { case_sensitivity = [%e case_sensitivity]
          ; clauses = [%e elist ~loc clause_exprs]
          }]
    in
    [ grammar ]
;;

(* Wrap [expr] in [fun a b ... ->] for type parameters. *)
let td_params_fun td expr =
  let loc = td.ptype_loc in
  let params =
    List.map td.ptype_params ~f:(fun param ->
      let { loc; txt } = get_type_param_name param in
      pvar ~loc (tyvar_grammar_name txt))
  in
  eabstract ~loc params expr
;;

module Row_field_type = struct
  type t =
    | Inherit of core_type
    | Tag_no_arg of string loc
    | Tag_with_arg of string loc * core_type

  let of_row_field ~loc row_field =
    match row_field with
    | Rinherit core_type -> Inherit core_type
    | Rtag (name, possibly_no_arg, possible_type_args) ->
      (match possibly_no_arg, possible_type_args with
       | true, [] -> Tag_no_arg name
       | false, [ core_type ] -> Tag_with_arg (name, core_type)
       | false, [] -> unsupported ~loc "empty polymorphic variant types"
       | true, _ :: _ | false, _ :: _ :: _ -> unsupported ~loc "intersection types")
  ;;
end

let attr_doc_comments attributes ~tags_of_doc_comments =
  match tags_of_doc_comments with
  | false -> []
  | true ->
    let doc_pattern = Ast_pattern.(pstr (pstr_eval (estring __) nil ^:: nil)) in
    List.filter_map attributes ~f:(fun attribute ->
      match attribute.attr_name.txt with
      | "ocaml.doc" | "doc" ->
        Ast_pattern.parse
          doc_pattern
          attribute.attr_loc
          attribute.attr_payload
          ~on_error:(fun () -> None)
          (fun doc -> Some doc)
      | _ -> None)
;;

let grammar_of_type_tags core_type grammar ~tags_of_doc_comments =
  let tags = Tags.get core_type ~tags:Attrs.tags_type ~tag:Attrs.tag_type in
  let loc = core_type.ptyp_loc in
  let comments = attr_doc_comments ~tags_of_doc_comments core_type.ptyp_attributes in
  with_tags_as_grammar grammar ~loc ~tags ~comments
;;

let grammar_of_field_tags field grammar ~tags_of_doc_comments =
  let tags = Tags.get field ~tags:Attrs.tags_ld ~tag:Attrs.tag_ld in
  let loc = field.pld_loc in
  let comments = attr_doc_comments ~tags_of_doc_comments field.pld_attributes in
  with_tags_as_list
    grammar
    ~loc
    ~tags
    ~comments
    ~core_type:[%type: Sexplib0.Sexp_grammar.field]
;;

let rec grammar_of_type core_type ~rec_flag ~tags_of_doc_comments =
  let loc = core_type.ptyp_loc in
  let grammar =
    let from_attribute =
      match
        ( Attribute.get Attrs.grammar_custom core_type
        , Attribute.get Attrs.grammar_any core_type )
      with
      | Some _, Some _ ->
        Some
          [%expr
            [%ocaml.warning
              "[@sexp_grammar.custom] and [@sexp_grammar.any] are mutually exclusive"]]
      | Some expr, None ->
        Some (untyped_grammar ~loc (annotated_grammar ~loc expr core_type))
      | None, Some maybe_name ->
        Some (any_grammar ~loc (Option.value maybe_name ~default:"ANY"))
      | None, None ->
        (* only check [[@sexp.opaque]] if neither other attribute is present, so that it
           only counts as using the attribute when we actually base the grammar on it *)
        (match Attribute.get Attrs.opaque core_type with
         | Some () -> Some (opaque_grammar ~loc)
         | None -> None)
    in
    match from_attribute with
    | Some expr -> expr
    | None ->
      (match Ppxlib_jane.Jane_syntax.Core_type.of_ast core_type with
       | Some (Jtyp_tuple ltps, _attrs) ->
         grammar_of_labeled_tuple ~loc ~rec_flag ~tags_of_doc_comments ltps
       | Some (Jtyp_layout _, _) | None ->
         (match core_type.ptyp_desc with
          | Ptyp_any -> any_grammar ~loc "_"
          | Ptyp_var name ->
            (match rec_flag with
             | Recursive ->
               (* For recursive grammars, [grammar_of_type] for any type variables is called
                  inside a [defn]. The variables should therefore be resolved as [Tyvar]
                  grammars. *)
               tyvar_grammar ~loc (estring ~loc name)
             | Nonrecursive ->
               (* Outside recursive [defn]s, type variables are passed in as function
                  arguments. *)
               unapplied_type_constr_conv
                 ~loc
                 ~f:tyvar_grammar_name
                 (Located.lident ~loc name)
               |> untyped_grammar ~loc)
          | Ptyp_arrow _ -> arrow_grammar ~loc
          | Ptyp_tuple list ->
            List.map ~f:(grammar_of_type ~rec_flag ~tags_of_doc_comments) list
            |> tuple_grammar ~loc
            |> list_grammar ~loc
          | Ptyp_constr (id, args) ->
            List.map args ~f:(fun core_type ->
              let loc = core_type.ptyp_loc in
              grammar_of_type ~rec_flag ~tags_of_doc_comments core_type
              |> typed_grammar ~loc)
            |> type_constr_conv ~loc ~f:grammar_name id
            |> untyped_grammar ~loc
          | Ptyp_object _ -> unsupported ~loc "object types"
          | Ptyp_class _ -> unsupported ~loc "class types"
          | Ptyp_alias _ -> unsupported ~loc "type aliases"
          | Ptyp_variant (rows, closed_flag, (_ : string list option)) ->
            (match closed_flag with
             | Open -> unsupported ~loc "open polymorphic variant types"
             | Closed ->
               grammar_of_polymorphic_variant ~loc ~rec_flag ~tags_of_doc_comments rows)
          | Ptyp_poly _ -> unsupported ~loc "explicitly polymorphic types"
          | Ptyp_package _ -> unsupported ~loc "first-class module types"
          | Ptyp_extension _ -> unsupported ~loc "unexpanded ppx extensions"))
  in
  grammar_of_type_tags core_type grammar ~tags_of_doc_comments

and grammar_of_labeled_tuple ~loc ~rec_flag ~tags_of_doc_comments alist =
  assert (Labeled_tuple.is_valid alist);
  let fields =
    List.concat_map alist ~f:(fun (lbl, typ) ->
      let lbl = Labeled_tuple.atom_of_label lbl in
      let field = grammar_of_type ~rec_flag ~tags_of_doc_comments typ in
      let clauses : Variant_clause_type.t list =
        (* Labeled tuples are encoded as a list of singleton variants, where the
           constructor name is used for the label. *)
        [ { name = { txt = lbl; loc }
          ; comments = []
          ; tags =
              { defined_using_tags = None; defined_using_tag = [] }
              (* We can use empty comments and tags because it's not possible to attach an
             attribute to a labeled tuple field. *)
          ; clause_kind = list_clause ~loc [%expr Cons ([%e field], Empty)]
          }
        ]
      in
      let case_sensitivity = [%expr Case_sensitive] in
      variant_grammars ~loc ~case_sensitivity ~clauses)
  in
  list_grammar ~loc (tuple_grammar ~loc fields)

and grammar_of_polymorphic_variant ~loc ~rec_flag ~tags_of_doc_comments rows =
  let inherits, clauses =
    List.partition_map rows ~f:(fun row : (_, Variant_clause_type.t) Either.t ->
      let tags = Tags.get row ~tags:Attrs.tags_poly ~tag:Attrs.tag_poly in
      let comments = attr_doc_comments ~tags_of_doc_comments row.prf_attributes in
      match Attribute.get Attrs.list_poly row with
      | Some () ->
        (match Row_field_type.of_row_field ~loc row.prf_desc with
         | Tag_with_arg (name, [%type: [%t? ty] list]) ->
           let clause_kind =
             grammar_of_type ~rec_flag ~tags_of_doc_comments ty
             |> many_grammar ~loc
             |> list_clause ~loc
           in
           Second { name; comments; tags; clause_kind }
         | _ -> Attrs.invalid_attribute ~loc Attrs.list_poly "_ list")
      | None ->
        (match Row_field_type.of_row_field ~loc row.prf_desc with
         | Inherit core_type ->
           First
             (grammar_of_type ~rec_flag ~tags_of_doc_comments core_type
              |> with_tags_as_grammar ~loc ~tags ~comments)
         | Tag_no_arg name ->
           Second { name; comments; tags; clause_kind = atom_clause ~loc }
         | Tag_with_arg (name, core_type) ->
           let clause_kind =
             [ grammar_of_type ~rec_flag ~tags_of_doc_comments core_type ]
             |> tuple_grammar ~loc
             |> list_clause ~loc
           in
           Second { name; comments; tags; clause_kind }))
  in
  variant_grammars ~loc ~case_sensitivity:[%expr Case_sensitive] ~clauses
  |> List.append inherits
  |> union_grammar ~loc
;;

let record_expr ~loc ~rec_flag ~tags_of_doc_comments ~extra_attr syntax fields =
  let fields =
    List.map fields ~f:(fun field ->
      let loc = field.pld_loc in
      let field_kind = Record_field_attrs.Of_sexp.create ~loc field in
      let required =
        match field_kind with
        | Specific Required -> true
        | Specific (Default _)
        | Sexp_bool | Sexp_option _ | Sexp_array _ | Sexp_list _ | Omit_nil -> false
      in
      let args =
        match field_kind with
        | Specific Required | Specific (Default _) | Omit_nil ->
          [%expr
            Cons
              ([%e grammar_of_type ~tags_of_doc_comments ~rec_flag field.pld_type], Empty)]
        | Sexp_bool -> [%expr Empty]
        | Sexp_option ty ->
          [%expr Cons ([%e grammar_of_type ~tags_of_doc_comments ~rec_flag ty], Empty)]
        | Sexp_list ty | Sexp_array ty ->
          [%expr
            Cons
              (List (Many [%e grammar_of_type ~tags_of_doc_comments ~rec_flag ty]), Empty)]
      in
      [%expr
        { name = [%e estr field.pld_name]
        ; required = [%e ebool ~loc required]
        ; args = [%e args]
        }]
      |> grammar_of_field_tags field ~tags_of_doc_comments)
  in
  let allow_extra_fields =
    match Attribute.get extra_attr syntax with
    | Some () -> true
    | None -> false
  in
  [%expr
    { allow_extra_fields = [%e ebool ~loc allow_extra_fields]
    ; fields = [%e elist ~loc fields]
    }]
;;

let grammar_of_variant ~loc ~rec_flag ~tags_of_doc_comments clause_decls =
  let clauses =
    List.map clause_decls ~f:(fun clause : Variant_clause_type.t ->
      let loc = clause.pcd_loc in
      let tags = Tags.get clause ~tags:Attrs.tags_cd ~tag:Attrs.tag_cd in
      let comments = attr_doc_comments ~tags_of_doc_comments clause.pcd_attributes in
      match Attribute.get Attrs.list_variant clause with
      | Some () ->
        (match clause.pcd_args with
         | Pcstr_tuple [ [%type: [%t? ty] list] ] ->
           let args =
             many_grammar ~loc (grammar_of_type ty ~rec_flag ~tags_of_doc_comments)
           in
           { name = clause.pcd_name; comments; tags; clause_kind = list_clause ~loc args }
         | _ -> Attrs.invalid_attribute ~loc Attrs.list_variant "_ list")
      | None ->
        (match clause.pcd_args with
         | Pcstr_tuple [] ->
           { name = clause.pcd_name; comments; tags; clause_kind = atom_clause ~loc }
         | Pcstr_tuple (_ :: _ as args) ->
           let args =
             tuple_grammar
               ~loc
               (List.map args ~f:(grammar_of_type ~rec_flag ~tags_of_doc_comments))
           in
           { name = clause.pcd_name; comments; tags; clause_kind = list_clause ~loc args }
         | Pcstr_record fields ->
           let args =
             record_expr
               ~loc
               ~rec_flag
               ~tags_of_doc_comments
               ~extra_attr:Attrs.allow_extra_fields_cd
               clause
               fields
             |> fields_grammar ~loc
           in
           { name = clause.pcd_name; comments; tags; clause_kind = list_clause ~loc args }))
  in
  variant_grammars
    ~loc
    ~case_sensitivity:[%expr Case_sensitive_except_first_character]
    ~clauses
  |> union_grammar ~loc
;;

let grammar_of_td ~ctxt ~rec_flag ~tags_of_doc_comments td =
  let loc = td.ptype_loc in
  match td.ptype_kind with
  | Ptype_open -> unsupported ~loc "open types"
  | Ptype_record fields ->
    record_expr
      ~loc
      ~rec_flag
      ~tags_of_doc_comments
      ~extra_attr:Attrs.allow_extra_fields_td
      td
      fields
    |> fields_grammar ~loc
    |> list_grammar ~loc
  | Ptype_variant clauses ->
    grammar_of_variant ~loc ~rec_flag ~tags_of_doc_comments clauses
  | Ptype_abstract ->
    (match td.ptype_manifest with
     | None -> abstract_grammar ~ctxt ~loc td.ptype_name
     | Some core_type -> grammar_of_type ~rec_flag ~tags_of_doc_comments core_type)
;;

let pattern_of_td td =
  let { loc; txt } = td.ptype_name in
  ppat_constraint
    ~loc
    (pvar ~loc (grammar_name txt))
    (ptyp_poly
       ~loc
       (List.map td.ptype_params ~f:get_type_param_name)
       (combinator_type_of_type_declaration td ~f:grammar_type))
;;

(* Any grammar expression that is purely a constant does no work, and does not need to be
   wrapped in [Lazy]. *)
let rec is_preallocated_constant expr =
  match expr.pexp_desc with
  | Pexp_constraint (expr, _) | Pexp_coerce (expr, _, _) | Pexp_open (_, expr) ->
    is_preallocated_constant expr
  | Pexp_constant _ -> true
  | Pexp_tuple args -> List.for_all ~f:is_preallocated_constant args
  | Pexp_variant (_, maybe_arg) | Pexp_construct (_, maybe_arg) ->
    Option.for_all ~f:is_preallocated_constant maybe_arg
  | Pexp_record (fields, maybe_template) ->
    List.for_all fields ~f:(fun (_, expr) -> is_preallocated_constant expr)
    && Option.for_all ~f:is_preallocated_constant maybe_template
  | _ -> false
;;

(* Any grammar expression that just refers to a previously defined grammar also does not
   need to be wrapped in [Lazy]. Accessing the previous grammar is work, but building the
   closure for a lazy value is at least as much work anyway. *)
let rec is_variable_access expr =
  match expr.pexp_desc with
  | Pexp_constraint (expr, _) | Pexp_coerce (expr, _, _) | Pexp_open (_, expr) ->
    is_variable_access expr
  | Pexp_ident _ -> true
  | Pexp_field (expr, _) -> is_variable_access expr
  | _ -> false
;;

let grammar_needs_lazy_wrapper expr =
  not (is_preallocated_constant expr || is_variable_access expr)
;;

let lazy_grammar ~loc td expr =
  if List.is_empty td.ptype_params
     (* polymorphic types generate functions, so the body does not need a [lazy] wrapper *)
     && grammar_needs_lazy_wrapper expr
  then [%expr Lazy (lazy [%e expr])]
  else expr
;;

let force_expr ~loc expr = [%expr Stdlib.Lazy.force [%e expr]]

(* Definitions of grammars that do not refer to each other. *)
let nonrecursive_grammars ~ctxt ~loc ~tags_of_doc_comments td_lists =
  List.concat_map td_lists ~f:(fun tds ->
    List.map tds ~f:(fun td ->
      let td = name_type_params_in_td td in
      let loc = td.ptype_loc in
      let pat = pattern_of_td td in
      let expr =
        grammar_of_td ~ctxt ~rec_flag:Nonrecursive ~tags_of_doc_comments td
        |> lazy_grammar td ~loc
        |> typed_grammar ~loc
        |> td_params_fun td
      in
      value_binding ~loc ~pat ~expr)
    |> pstr_value_list ~loc Nonrecursive)
;;

(* Type constructor grammars used to "tie the knot" for (mutally) recursive grammars. *)
let recursive_grammar_tycons tds =
  List.map tds ~f:(fun td ->
    let td = name_type_params_in_td td in
    let loc = td.ptype_loc in
    let pat = pattern_of_td td in
    let expr =
      recursive_grammar
        ~loc
        (estr td.ptype_name)
        (List.map td.ptype_params ~f:(fun param ->
           let { loc; txt } = get_type_param_name param in
           tyvar_grammar_name txt |> evar ~loc |> untyped_grammar ~loc)
         |> elist ~loc)
      |> typed_grammar ~loc
      |> td_params_fun td
    in
    value_binding ~loc ~pat ~expr)
;;

(* Recursive grammar definitions, based on the type constructors from above. *)
let recursive_grammar_defns ~ctxt ~loc ~tags_of_doc_comments tds =
  List.map tds ~f:(fun td ->
    let td = name_type_params_in_td td in
    let loc = td.ptype_loc in
    let tycon = estr td.ptype_name in
    let tyvars =
      List.map td.ptype_params ~f:(fun param -> estr (get_type_param_name param))
      |> elist ~loc
    in
    let grammar = grammar_of_td ~ctxt ~rec_flag:Recursive ~tags_of_doc_comments td in
    defn_expr ~loc ~tycon ~tyvars ~grammar)
  |> elist ~loc
;;

(* Grammar expression using [Recursive] and a shared definition of grammar definitions.
   The shared definitions are wrapped in [lazy] to avoid toplevel side effects. *)
let recursive_grammar_expr ~defns_name td =
  let td = name_type_params_in_td td in
  let loc = td.ptype_loc in
  let pat = pattern_of_td td in
  let expr =
    let tyvars =
      List.map td.ptype_params ~f:(fun param ->
        let { loc; txt } = get_type_param_name param in
        tyvar_grammar_name txt |> evar ~loc |> untyped_grammar ~loc)
      |> elist ~loc
    in
    tycon_grammar
      ~loc
      (estr td.ptype_name)
      tyvars
      (evar ~loc defns_name |> force_expr ~loc)
    |> lazy_grammar td ~loc
    |> typed_grammar ~loc
    |> td_params_fun td
  in
  value_binding ~loc ~pat ~expr
;;

(* Puts together recursive grammar definitions from the parts implemented above. *)
let recursive_grammars ~ctxt ~loc ~tags_of_doc_comments tds =
  match List.is_empty tds with
  | true -> []
  | false ->
    let defns_name = gen_symbol ~prefix:"grammars" () in
    let defns_item =
      let expr =
        recursive_grammar_defns ~ctxt ~loc ~tags_of_doc_comments tds
        |> pexp_let ~loc Nonrecursive (recursive_grammar_tycons tds)
        |> pexp_lazy ~loc
      in
      let pat = ppat_constraint ~loc (pvar ~loc defns_name) (defns_type ~loc) in
      pstr_value ~loc Nonrecursive [ value_binding ~loc ~pat ~expr ]
    in
    let grammars_item =
      List.map tds ~f:(recursive_grammar_expr ~defns_name) |> pstr_value ~loc Nonrecursive
    in
    [%str
      include struct
        open struct
          [%%i defns_item]
        end

        [%%i grammars_item]
      end]
;;

let partition_recursive_and_nonrecursive ~rec_flag tds =
  match (rec_flag : rec_flag) with
  | Nonrecursive -> [], [ tds ]
  | Recursive ->
    (* Pulling out non-recursive references repeatedly means we only "tie the knot" for
       variables that actually need it, and we don't have to manually [ignore] the added
       bindings in case they are unused. *)
    let rec loop tds ~acc =
      let obj =
        object
          inherit type_is_recursive Recursive tds
          method recursion td = {<type_names = [ td.ptype_name.txt ]>}#go ()
        end
      in
      let recursive, nonrecursive =
        List.partition_tf tds ~f:(fun td ->
          match obj#recursion td with
          | Recursive -> true
          | Nonrecursive -> false)
      in
      if List.is_empty recursive || List.is_empty nonrecursive
      then recursive, nonrecursive :: acc
      else loop recursive ~acc:(nonrecursive :: acc)
    in
    loop tds ~acc:[]
;;

let str_type_decl ~ctxt (rec_flag, tds) tags_of_doc_comments =
  let loc = Expansion_context.Deriver.derived_item_loc ctxt in
  let recursive, nonrecursive = partition_recursive_and_nonrecursive ~rec_flag tds in
  [ recursive_grammars ~ctxt ~loc ~tags_of_doc_comments recursive
  ; nonrecursive_grammars ~ctxt ~loc ~tags_of_doc_comments nonrecursive
  ]
  |> List.concat
;;

let sig_type_decl ~ctxt:_ (_rec_flag, tds) =
  List.map tds ~f:(fun td ->
    let loc = td.ptype_loc in
    value_description
      ~loc
      ~name:(Loc.map td.ptype_name ~f:grammar_name)
      ~type_:(combinator_type_of_type_declaration td ~f:grammar_type)
      ~prim:[]
    |> psig_value ~loc)
;;

let extension_loc ~ctxt =
  let loc = Expansion_context.Extension.extension_point_loc ctxt in
  { loc with loc_ghost = true }
;;

let core_type ~tags_of_doc_comments ~ctxt core_type =
  let loc = extension_loc ~ctxt in
  pexp_constraint
    ~loc
    (core_type
     |> grammar_of_type ~rec_flag:Nonrecursive ~tags_of_doc_comments
     |> typed_grammar ~loc)
    (core_type |> grammar_type ~loc)
  |> Merlin_helpers.hide_expression
;;

let type_extension ~ctxt core_type =
  assert_no_attributes_in#core_type core_type;
  let loc = extension_loc ~ctxt in
  core_type |> grammar_type ~loc
;;