open! Import
open Common

type (_, _) equality = Eq : ('a, 'a) equality | Ne : (_, _) equality

module Context = struct
  type 'a t =
    | Class_expr : class_expr t
    | Class_field : class_field t
    | Class_type : class_type t
    | Class_type_field : class_type_field t
    | Core_type : core_type t
    | Expression : expression t
    | Module_expr : module_expr t
    | Module_type : module_type t
    | Pattern : pattern t
    | Signature_item : signature_item t
    | Structure_item : structure_item t
    | Ppx_import : type_declaration t

  type packed = T : _ t -> packed

  let class_expr = Class_expr
  let class_field = Class_field
  let class_type = Class_type
  let class_type_field = Class_type_field
  let core_type = Core_type
  let expression = Expression
  let module_expr = Module_expr
  let module_type = Module_type
  let pattern = Pattern
  let signature_item = Signature_item
  let structure_item = Structure_item

  let desc : type a. a t -> string = function
    | Class_expr -> "class expression"
    | Class_field -> "class field"
    | Class_type -> "class type"
    | Class_type_field -> "class type field"
    | Core_type -> "core type"
    | Expression -> "expression"
    | Module_expr -> "module expression"
    | Module_type -> "module type"
    | Pattern -> "pattern"
    | Signature_item -> "signature item"
    | Structure_item -> "structure item"
    | Ppx_import -> "type declaration"

  let eq : type a b. a t -> b t -> (a, b) equality =
   fun a b ->
    match (a, b) with
    | Class_expr, Class_expr -> Eq
    | Class_field, Class_field -> Eq
    | Class_type, Class_type -> Eq
    | Class_type_field, Class_type_field -> Eq
    | Core_type, Core_type -> Eq
    | Expression, Expression -> Eq
    | Module_expr, Module_expr -> Eq
    | Module_type, Module_type -> Eq
    | Pattern, Pattern -> Eq
    | Signature_item, Signature_item -> Eq
    | Structure_item, Structure_item -> Eq
    | Ppx_import, Ppx_import -> Eq
    | _ ->
        assert (Poly.( <> ) (T a) (T b));
        Ne

  let get_ppx_import_extension type_decl =
    match type_decl with
    | { ptype_manifest = Some { ptyp_desc = Ptyp_extension (name, _); _ }; _ }
      ->
        let virtual_payload =
          Ast_builder.Default.pstr_type ~loc:type_decl.ptype_loc Recursive
            [ type_decl ]
        in
        let attr = [] in
        Some ((name, PStr [ virtual_payload ]), attr)
    | _ -> None

  let get_extension : type a. a t -> a -> (extension * attributes) option =
   fun t x ->
    match (t, x) with
    | Class_expr, { pcl_desc = Pcl_extension e; pcl_attributes = a; _ } ->
        Some (e, a)
    | Class_field, { pcf_desc = Pcf_extension e; pcf_attributes = a; _ } ->
        Some (e, a)
    | Class_type, { pcty_desc = Pcty_extension e; pcty_attributes = a; _ } ->
        Some (e, a)
    | Class_type_field, { pctf_desc = Pctf_extension e; pctf_attributes = a; _ }
      ->
        Some (e, a)
    | Core_type, { ptyp_desc = Ptyp_extension e; ptyp_attributes = a; _ } ->
        Some (e, a)
    | Expression, { pexp_desc = Pexp_extension e; pexp_attributes = a; _ } ->
        Some (e, a)
    | Module_expr, { pmod_desc = Pmod_extension e; pmod_attributes = a; _ } ->
        Some (e, a)
    | Module_type, { pmty_desc = Pmty_extension e; pmty_attributes = a; _ } ->
        Some (e, a)
    | Pattern, { ppat_desc = Ppat_extension e; ppat_attributes = a; _ } ->
        Some (e, a)
    | Signature_item, { psig_desc = Psig_extension (e, a); _ } -> Some (e, a)
    | Structure_item, { pstr_desc = Pstr_extension (e, a); _ } -> Some (e, a)
    | Ppx_import, type_decl -> get_ppx_import_extension type_decl
    | _ -> None

  let merge_attributes : type a. a t -> a -> attributes -> a =
   fun t x attrs ->
    match t with
    | Class_expr -> { x with pcl_attributes = x.pcl_attributes @ attrs }
    | Class_field -> { x with pcf_attributes = x.pcf_attributes @ attrs }
    | Class_type -> { x with pcty_attributes = x.pcty_attributes @ attrs }
    | Class_type_field -> { x with pctf_attributes = x.pctf_attributes @ attrs }
    | Core_type -> { x with ptyp_attributes = x.ptyp_attributes @ attrs }
    | Expression -> { x with pexp_attributes = x.pexp_attributes @ attrs }
    | Module_expr -> { x with pmod_attributes = x.pmod_attributes @ attrs }
    | Module_type -> { x with pmty_attributes = x.pmty_attributes @ attrs }
    | Pattern -> { x with ppat_attributes = x.ppat_attributes @ attrs }
    | Signature_item ->
        assert_no_attributes attrs;
        x
    | Structure_item ->
        assert_no_attributes attrs;
        x
    | Ppx_import ->
        assert_no_attributes attrs;
        x
end

let registrar =
  Name.Registrar.create ~kind:"extension" ~current_file:__FILE__
    ~string_of_context:(fun (Context.T ctx) -> Some (Context.desc ctx))

module Make (Callback : sig
  type 'a t
end) =
struct
  type ('a, 'b) payload_parser =
    | Payload_parser :
        ('a, 'b, 'c) Ast_pattern.t * 'b Callback.t
        -> ('a, 'c) payload_parser

  type ('context, 'payload) t = {
    name : Name.Pattern.t;
    context : 'context Context.t;
    payload : (payload, 'payload) payload_parser;
    with_arg : bool;
  }

  let declare :
      type a.
      with_arg:bool ->
      string ->
      a Context.t ->
      (payload, 'b, 'payload) Ast_pattern.t ->
      'b Callback.t ->
      (a, 'payload) t =
   fun ~with_arg name context pattern k ->
    (* Check that there is no collisions between ppx_import and core_type
       extensions *)
    (match context with
    | Context.Ppx_import ->
        Name.Registrar.check_collisions registrar (Context.T Core_type) name
    | Context.Core_type ->
        Name.Registrar.check_collisions registrar (Context.T Ppx_import) name
    | _ -> ());
    Name.Registrar.register ~kind:`Extension registrar (Context.T context) name;
    {
      name = Name.Pattern.make name;
      context;
      payload = Payload_parser (pattern, k);
      with_arg;
    }

  let find ts (ext : extension) =
    let { txt = name; loc } = fst ext in
    let name, arg = Name.split_path name in
    match List.filter ts ~f:(fun t -> Name.Pattern.matches t.name name) with
    | [] -> None
    | _ :: _ :: _ as l ->
        Location.raise_errorf ~loc "Multiple match for extensions: %s"
          (String.concat ~sep:", "
             (List.map l ~f:(fun t -> Name.Pattern.name t.name)))
    | [ t ] ->
        if (not t.with_arg) && Option.is_some arg then
          Location.raise_errorf ~loc
            "Extension %s doesn't expect a path argument" name;
        let arg =
          Option.map arg ~f:(fun s ->
              let shift = String.length name + 1 in
              let start = loc.loc_start in
              {
                txt = Longident.parse s;
                loc =
                  {
                    loc with
                    loc_start = { start with pos_cnum = start.pos_cnum + shift };
                  };
              })
        in
        Some (t, arg)
end

module Expert = struct
  include Make (struct
    type 'a t = arg:Longident.t Loc.t option -> 'a
  end)

  let declare_with_path_arg name ctx patt f =
    declare ~with_arg:true name ctx patt f

  let declare name ctx patt f =
    declare ~with_arg:false name ctx patt (fun ~arg:_ -> f)

  let convert ts ~loc ext =
    match find ts ext with
    | None -> None
    | Some ({ payload = Payload_parser (pattern, f); _ }, arg) ->
        Some (Ast_pattern.parse pattern loc (snd ext) (f ~arg))
end

module M = Make (struct
  type 'a t =
    ctxt:Expansion_context.Extension.t -> arg:Longident.t Loc.t option -> 'a
end)

type 'a expander_result = Simple of 'a | Inline of 'a list

module For_context = struct
  type 'a t = ('a, 'a expander_result) M.t

  let convert ts ~ctxt ext =
    let loc = Expansion_context.Extension.extension_point_loc ctxt in
    match M.find ts ext with
    | None -> None
    | Some ({ payload = M.Payload_parser (pattern, f); _ }, arg) -> (
        match Ast_pattern.parse pattern loc (snd ext) (f ~ctxt ~arg) with
        | Simple x -> Some x
        | Inline _ -> failwith "Extension.convert")

  let convert_inline ts ~ctxt ext =
    let loc = Expansion_context.Extension.extension_point_loc ctxt in
    match M.find ts ext with
    | None -> None
    | Some ({ payload = M.Payload_parser (pattern, f); _ }, arg) -> (
        match Ast_pattern.parse pattern loc (snd ext) (f ~ctxt ~arg) with
        | Simple x -> Some [ x ]
        | Inline l -> Some l)
end

type t = T : _ For_context.t -> t

let check_context_for_inline : type a. func:string -> a Context.t -> unit =
 fun ~func ctx ->
  match ctx with
  | Context.Class_field -> ()
  | Context.Class_type_field -> ()
  | Context.Signature_item -> ()
  | Context.Structure_item -> ()
  | context ->
      Printf.ksprintf invalid_arg "%s: %s can't be inlined" func
        (Context.desc context)

let rec filter_by_context :
    type a. a Context.t -> t list -> a For_context.t list =
 fun context expanders ->
  match expanders with
  | [] -> []
  | T t :: rest -> (
      match Context.eq context t.context with
      | Eq -> t :: filter_by_context context rest
      | Ne -> filter_by_context context rest)

let fail ctx (name, _) =
  if
    not
      (Name.Whitelisted.is_whitelisted ~kind:`Extension name.txt
      || Name.ignore_checks name.txt)
  then
    Name.Registrar.raise_errorf registrar (Context.T ctx)
      "Extension `%s' was not translated" name

let check_unused =
  object
    inherit Ast_traverse.iter as super

    method! extension (name, _) =
      Location.raise_errorf ~loc:name.loc
        "extension not expected here, Ppxlib.Extension needs updating!"

    method! core_type_desc =
      function
      | Ptyp_extension ext -> fail Core_type ext | x -> super#core_type_desc x

    method! pattern_desc =
      function
      | Ppat_extension ext -> fail Pattern ext | x -> super#pattern_desc x

    method! expression_desc =
      function
      | Pexp_extension ext -> fail Expression ext | x -> super#expression_desc x

    method! class_type_desc =
      function
      | Pcty_extension ext -> fail Class_type ext | x -> super#class_type_desc x

    method! class_type_field_desc =
      function
      | Pctf_extension ext -> fail Class_type_field ext
      | x -> super#class_type_field_desc x

    method! class_expr_desc =
      function
      | Pcl_extension ext -> fail Class_expr ext | x -> super#class_expr_desc x

    method! class_field_desc =
      function
      | Pcf_extension ext -> fail Class_field ext
      | x -> super#class_field_desc x

    method! module_type_desc =
      function
      | Pmty_extension ext -> fail Module_type ext
      | x -> super#module_type_desc x

    method! signature_item_desc =
      function
      | Psig_extension (ext, _) -> fail Signature_item ext
      | x -> super#signature_item_desc x

    method! module_expr_desc =
      function
      | Pmod_extension ext -> fail Module_expr ext
      | x -> super#module_expr_desc x

    method! structure_item_desc =
      function
      | Pstr_extension (ext, _) -> fail Structure_item ext
      | x -> super#structure_item_desc x
  end

module V3 = struct
  type nonrec t = t

  let declare name context pattern k =
    let pattern = Ast_pattern.map_result pattern ~f:(fun x -> Simple x) in
    T
      (M.declare ~with_arg:false name context pattern (fun ~ctxt ~arg:_ ->
           k ~ctxt))

  let declare_inline name context pattern k =
    check_context_for_inline context ~func:"Extension.declare_inline";
    let pattern = Ast_pattern.map_result pattern ~f:(fun x -> Inline x) in
    T
      (M.declare ~with_arg:false name context pattern (fun ~ctxt ~arg:_ ->
           k ~ctxt))
end

let declare name context pattern f =
  V3.declare name context pattern
    (Expansion_context.Extension.with_loc_and_path f)

let declare_inline name context pattern f =
  V3.declare_inline name context pattern
    (Expansion_context.Extension.with_loc_and_path f)

let declare_with_path_arg name context pattern k =
  let k' = Expansion_context.Extension.with_loc_and_path k in
  let pattern = Ast_pattern.map_result pattern ~f:(fun x -> Simple x) in
  T (M.declare ~with_arg:true name context pattern k')

let declare_inline_with_path_arg name context pattern k =
  let k' = Expansion_context.Extension.with_loc_and_path k in
  check_context_for_inline context
    ~func:"Extension.declare_inline_with_path_arg";
  let pattern = Ast_pattern.map_result pattern ~f:(fun x -> Inline x) in
  T (M.declare ~with_arg:true name context pattern k')

let __declare_ppx_import name expand =
  (* This pattern is used to unwrap the type declaration from the payload
     assembled by [Context.get_ppx_import_extension] *)
  let pattern = Ast_pattern.(pstr (pstr_type recursive (__ ^:: nil) ^:: nil)) in
  V3.declare name Context.Ppx_import pattern expand

module V2 = struct
  type nonrec t = t

  let declare = declare
  let declare_inline = declare_inline
end
