open Ppxlib

(* Generates a [let derived_<type_name> = "ok"] or a
   [let derived_<type_name> = "uninterpreted extension in input"] if
   the type manifest is an uninterpreted extension. *)
let deriver =
  let binding ~loc type_name expr =
    let var_name = "derived_" ^ type_name in
    let pat = Ast_builder.Default.ppat_var ~loc {txt = var_name; loc} in
    let vb = Ast_builder.Default.value_binding ~loc ~pat ~expr in
    [Ast_builder.Default.pstr_value ~loc Nonrecursive [vb]]
  in
  let str_type_decl =
    Deriving.Generator.V2.make_noarg
      (fun ~ctxt (_rec_flag, type_decls) ->
         let loc = Expansion_context.Deriver.derived_item_loc ctxt in
         match type_decls with
         | { ptype_manifest = Some {ptyp_desc = Ptyp_extension _; _}
           ; ptype_name = {txt; _}; _}::_ ->
           binding ~loc txt [%expr "uninterpreted extension in input"]
         | {ptype_name = {txt; _}; _}::_ ->
           binding ~loc txt [%expr "ok"]
         | [] -> assert false)
  in
  Deriving.add ~str_type_decl "derived"

[%%expect{|
val deriver : Deriving.t = <abstr>
|}]

(* Generates a [type t = int] *)
let gen_type_decl =
  Extension.V3.declare
    "gen_type_decl"
    Extension.Context.structure_item
    Ast_pattern.(pstr nil)
    (fun ~ctxt ->
       let loc = Expansion_context.Extension.extension_point_loc ctxt in
       [%stri type t = int])
  |> Context_free.Rule.extension

let () = Driver.register_transformation ~rules:[gen_type_decl] "gen_type_decl"

[%%expect{|
val gen_type_decl : Context_free.Rule.t = <abstr>
|}]

(* You cannot attach attributes to structure item extension points *)
[%%gen_type_decl]
[@@deriving derived]

[%%expect{|
Line _, characters 3-19:
Error: Attributes not allowed here
|}]

(* Generates a [type t = int[@@deriving derived]] *)
let gen_type_decl_with_derived =
  Extension.V3.declare
    "gen_type_decl_with_derived"
    Extension.Context.structure_item
    Ast_pattern.(pstr nil)
    (fun ~ctxt ->
       let loc = Expansion_context.Extension.extension_point_loc ctxt in
       [%stri type t = int[@@deriving derived]])
  |> Context_free.Rule.extension

let () =
  Driver.register_transformation
    ~rules:[gen_type_decl_with_derived]
    "gen_type_decl_with_derived"

[%%expect{|
val gen_type_decl_with_derived : Context_free.Rule.t = <abstr>
|}]

(* Attributes rule must be applied in code generated by a structure item
   extension *)
[%%gen_type_decl_with_derived]

[%%expect{|
type t = int
val derived_t : string = "ok"
|}]

let gen_inline_type_decls_with_derived =
  Extension.V3.declare_inline
    "gen_inline_type_decls_with_derived"
    Extension.Context.structure_item
    Ast_pattern.(pstr nil)
    (fun ~ctxt ->
       let loc = Expansion_context.Extension.extension_point_loc ctxt in
       [%str
         type t = int[@@deriving derived]
         type u = float[@@deriving derived]])
  |> Context_free.Rule.extension

let () =
  Driver.register_transformation
    ~rules:[gen_inline_type_decls_with_derived]
    "gen_inline_type_decls_with_derived"

[%%expect{|
val gen_inline_type_decls_with_derived : Context_free.Rule.t = <abstr>
|}]

(* That also stands for inline extension rules *)
[%%gen_inline_type_decls_with_derived]

[%%expect{|
type t = int
val derived_t : string = "ok"
type u = float
val derived_u : string = "ok"
|}]

let id =
  Extension.V3.declare
    "id"
    Extension.Context.core_type
    Ast_pattern.(ptyp __)
    (fun ~ctxt:_ core_type -> core_type)
  |> Context_free.Rule.extension

let () = Driver.register_transformation ~rules:[id] "id"

[%%expect{|
val id : Context_free.Rule.t = <abstr>
|}]

(* Nodes with attributes are expanded before attribute-based, inline
   code generation rules are applied.
   In this below, the `[[%id: int]]` is interpreted before the deriver
   is applied, meaning it can't see this extension point in its expand
   function argument. *)
type t = [%id: int]
[@@deriving derived]

[%%expect{|
type t = int
val derived_t : string = "ok"
|}]