Source file attribute.ml

open! Import

let poly_equal a b =
  let module Poly = struct
    type t = T : _ -> t
  end in
  Stdppx.Poly.equal (Poly.T a) (Poly.T b)
;;

module Context = struct
  type 'a t =
    | Label_declaration       : label_declaration       t
    | Constructor_declaration : constructor_declaration t
    | Type_declaration        : type_declaration        t
    | Type_exception          : type_exception          t
    | Type_extension          : type_extension          t
    | Extension_constructor   : extension_constructor   t
    | Pattern                 : pattern                 t
    | Core_type               : core_type               t
    | Expression              : expression              t
    | Value_description       : value_description       t
    | Class_type              : class_type              t
    | Class_type_field        : class_type_field        t
    | Class_infos             : _ class_infos           t
    | Class_expr              : class_expr              t
    | Class_field             : class_field             t
    | Module_type             : module_type             t
    | Module_declaration      : module_declaration      t
    | Module_type_declaration : module_type_declaration t
    | Module_substitution     : module_substitution     t
    | Open_description        : open_description        t
    | Open_declaration        : open_declaration        t
    | Include_infos           : _ include_infos         t
    | Module_expr             : module_expr             t
    | Value_binding           : value_binding           t
    | Module_binding          : module_binding          t
    | Pstr_eval               : structure_item          t
    | Pstr_extension          : structure_item          t
    | Psig_extension          : signature_item          t
    | Rtag                    : row_field               t
    | Object_type_field       : object_field            t

  let label_declaration       = Label_declaration
  let constructor_declaration = Constructor_declaration
  let type_declaration        = Type_declaration
  let type_extension          = Type_extension
  let type_exception          = Type_exception
  let extension_constructor   = Extension_constructor
  let pattern                 = Pattern
  let core_type               = Core_type
  let expression              = Expression
  let value_description       = Value_description
  let class_type              = Class_type
  let class_type_field        = Class_type_field
  let class_infos             = Class_infos
  let class_expr              = Class_expr
  let class_field             = Class_field
  let module_type             = Module_type
  let module_declaration      = Module_declaration
  let module_type_declaration = Module_type_declaration
  let open_description        = Open_description
  let include_infos           = Include_infos
  let module_expr             = Module_expr
  let value_binding           = Value_binding
  let module_binding          = Module_binding
  let pstr_eval               = Pstr_eval
  let pstr_extension          = Pstr_extension
  let psig_extension          = Psig_extension
  let rtag                    = Rtag
  let object_type_field       = Object_type_field

  let get_pstr_eval st =
    match st.pstr_desc with
    | Pstr_eval (e, l) -> (e, l)
    | _ -> failwith "Attribute.Context.get_pstr_eval"

  let get_pstr_extension st =
    match st.pstr_desc with
    | Pstr_extension (e, l) -> (e, l)
    | _ -> failwith "Attribute.Context.get_pstr_extension"

  let get_psig_extension st =
    match st.psig_desc with
    | Psig_extension (e, l) -> (e, l)
    | _ -> failwith "Attribute.Context.get_psig_extension"

  let get_attributes : type a. a t -> a -> attributes = fun t x ->
    match t with
    | Label_declaration       -> x.pld_attributes
    | Constructor_declaration -> x.pcd_attributes
    | Type_declaration        -> x.ptype_attributes
    | Type_extension          -> x.ptyext_attributes
    | Type_exception          -> x.ptyexn_attributes
    | Extension_constructor   -> x.pext_attributes
    | Pattern                 -> x.ppat_attributes
    | Core_type               -> x.ptyp_attributes
    | Expression              -> x.pexp_attributes
    | Value_description       -> x.pval_attributes
    | Class_type              -> x.pcty_attributes
    | Class_type_field        -> x.pctf_attributes
    | Class_infos             -> x.pci_attributes
    | Class_expr              -> x.pcl_attributes
    | Class_field             -> x.pcf_attributes
    | Module_type             -> x.pmty_attributes
    | Module_declaration      -> x.pmd_attributes
    | Module_type_declaration -> x.pmtd_attributes
    | Module_substitution     -> x.pms_attributes
    | Open_description        -> x.popen_attributes
    | Open_declaration        -> x.popen_attributes
    | Include_infos           -> x.pincl_attributes
    | Module_expr             -> x.pmod_attributes
    | Value_binding           -> x.pvb_attributes
    | Module_binding          -> x.pmb_attributes
    | Pstr_eval               -> snd (get_pstr_eval      x)
    | Pstr_extension          -> snd (get_pstr_extension x)
    | Psig_extension          -> snd (get_psig_extension x)
    | Rtag                    -> x.prf_attributes
    | Object_type_field       -> x.pof_attributes

  let set_attributes : type a. a t -> a -> attributes -> a = fun t x attrs ->
    match t with
    | Label_declaration       -> { x with pld_attributes    = attrs }
    | Constructor_declaration -> { x with pcd_attributes    = attrs }
    | Type_declaration        -> { x with ptype_attributes  = attrs }
    | Type_extension          -> { x with ptyext_attributes = attrs }
    | Type_exception          -> { x with ptyexn_attributes = attrs }
    | Extension_constructor   -> { x with pext_attributes   = attrs }
    | Pattern                 -> { x with ppat_attributes   = attrs }
    | Core_type               -> { x with ptyp_attributes   = attrs }
    | Expression              -> { x with pexp_attributes   = attrs }
    | Value_description       -> { x with pval_attributes   = attrs }
    | Class_type              -> { x with pcty_attributes   = attrs }
    | Class_type_field        -> { x with pctf_attributes   = attrs }
    | Class_infos             -> { x with pci_attributes    = attrs }
    | Class_expr              -> { x with pcl_attributes    = attrs }
    | Class_field             -> { x with pcf_attributes    = attrs }
    | Module_type             -> { x with pmty_attributes   = attrs }
    | Module_declaration      -> { x with pmd_attributes    = attrs }
    | Module_type_declaration -> { x with pmtd_attributes   = attrs }
    | Module_substitution     -> { x with pms_attributes   = attrs }
    | Open_description        -> { x with popen_attributes  = attrs }
    | Open_declaration        -> { x with popen_attributes  = attrs }
    | Include_infos           -> { x with pincl_attributes  = attrs }
    | Module_expr             -> { x with pmod_attributes   = attrs }
    | Value_binding           -> { x with pvb_attributes    = attrs }
    | Module_binding          -> { x with pmb_attributes    = attrs }
    | Pstr_eval ->
      { x with pstr_desc = Pstr_eval      (get_pstr_eval x      |> fst, attrs) }
    | Pstr_extension ->
      { x with pstr_desc = Pstr_extension (get_pstr_extension x |> fst, attrs) }
    | Psig_extension ->
      { x with psig_desc = Psig_extension (get_psig_extension x |> fst, attrs) }
    | Rtag                    -> { x with prf_attributes    = attrs}
    | Object_type_field       -> { x with pof_attributes    = attrs}

  let desc : type a. a t -> string = function
    | Label_declaration       -> "label declaration"
    | Constructor_declaration -> "constructor declaration"
    | Type_declaration        -> "type declaration"
    | Type_extension          -> "type extension"
    | Type_exception          -> "type exception"
    | Extension_constructor   -> "extension constructor"
    | Pattern                 -> "pattern"
    | Core_type               -> "core type"
    | Expression              -> "expression"
    | Value_description       -> "value"
    | Class_type              -> "class type"
    | Class_type_field        -> "class type field"
    | Class_infos             -> "class declaration"
    | Class_expr              -> "class expression"
    | Class_field             -> "class field"
    | Module_type             -> "module type"
    | Module_declaration      -> "module declaration"
    | Module_type_declaration -> "module type declaration"
    | Module_substitution     -> "module substitution"
    | Open_description        -> "open"
    | Open_declaration        -> "open"
    | Include_infos           -> "include"
    | Module_expr             -> "module expression"
    | Value_binding           -> "value binding"
    | Module_binding          -> "module binding"
    | Pstr_eval               -> "toplevel expression"
    | Pstr_extension          -> "toplevel extension"
    | Psig_extension          -> "toplevel signature extension"
    | Rtag                    -> "polymorphic variant tag"
    | Object_type_field       -> "object type field"

(*
  let pattern : type a b c d. a t
    -> (attributes, b, c) Ast_pattern.t
    -> (a, c, d) Ast_pattern.t
    -> (a, b, d) Ast_pattern.t = function
    | Label_declaration       -> Ast_pattern.pld_attributes
    | Constructor_declaration -> Ast_pattern.pcd_attributes
    | Type_declaration        -> Ast_pattern.ptype_attributes
    | Type_extension          -> Ast_pattern.ptyext_attributes
    | Extension_constructor   -> Ast_pattern.pext_attributes
*)

  let equal : _ t -> _ t -> bool = poly_equal
end

module Floating_context = struct
  type 'a t =
    | Structure_item   : structure_item   t
    | Signature_item   : signature_item   t
    | Class_field      : class_field      t
    | Class_type_field : class_type_field t

  let structure_item   = Structure_item
  let signature_item   = Signature_item
  let class_field      = Class_field
  let class_type_field = Class_type_field

  let get_attribute_if_is_floating_node : type a. a t -> a -> attribute option
    = fun t x ->
      match t, x with
      | Structure_item   , { pstr_desc = Pstr_attribute a; _ } -> Some a
      | Signature_item   , { psig_desc = Psig_attribute a; _ } -> Some a
      | Class_field      , { pcf_desc  = Pcf_attribute  a; _ } -> Some a
      | Class_type_field , { pctf_desc = Pctf_attribute a; _ } -> Some a
      | _ -> None

  let get_attribute t x =
    match get_attribute_if_is_floating_node t x with
    | Some a -> a
    | None   -> failwith "Attribute.Floating.Context.get_attribute"

  let replace_by_dummy : type a. a t -> a -> a =
    let dummy_ext = ({ txt = ""; loc = Location.none }, PStr []) in
    fun t x ->
    match t with
    | Structure_item   -> { x with pstr_desc = Pstr_extension (dummy_ext, []) }
    | Signature_item   -> { x with psig_desc = Psig_extension (dummy_ext, []) }
    | Class_field      -> { x with pcf_desc  = Pcf_extension   dummy_ext      }
    | Class_type_field -> { x with pctf_desc = Pctf_extension  dummy_ext      }

  let desc : type a. a t -> string = function
    | Structure_item   -> "structure item"
    | Signature_item   -> "signature item"
    | Class_field      -> "class field"
    | Class_type_field -> "class type field"

  let equal : _ t -> _ t -> bool = poly_equal
end

type packed_context =
  | On_item  : _ Context.t          -> packed_context
  | Floating : _ Floating_context.t -> packed_context

type _ payload_parser =
    Payload_parser
    : (payload, 'a, 'b) Ast_pattern.t * (name_loc:Location.t -> 'a)
    -> 'b payload_parser

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

type packed = T : (_, _) t -> packed

let name t = Name.Pattern.name t.name
let context t = t.context

let registrar =
  Name.Registrar.create
    ~kind:"attribute"
    ~current_file:__FILE__
    ~string_of_context:(function
      | On_item  t -> Some (Context         .desc t)
      | Floating t -> Some (Floating_context.desc t ^ " (floating)"))
;;

let declare_with_name_loc name context pattern k =
  Name.Registrar.register ~kind:`Attribute registrar (On_item context) name;
  { name = Name.Pattern.make name
  ; context
  ; payload = Payload_parser (pattern, k)
  }
;;

let declare name context pattern k =
  declare_with_name_loc name context pattern (fun ~name_loc:_ -> k)
;;

module Attribute_table = Caml.Hashtbl.Make(struct
    type t = string loc
    let hash : t -> int = Hashtbl.hash
    let equal : t -> t -> bool = Poly.equal
  end)

let not_seen = Attribute_table.create 128

let mark_as_seen { attr_name; _ } =
  Attribute_table.remove not_seen attr_name
;;

let mark_as_handled_manually = mark_as_seen

let explicitly_drop = object
  inherit Ast_traverse.iter
  method! attribute = mark_as_seen
end

let get_internal =
  let rec find_best_match t attributes longest_match =
    match attributes with
    | [] -> longest_match
    | { attr_name = name; _ } as attr :: rest ->
      if Name.Pattern.matches t.name name.txt then begin
        match longest_match with
        | None -> find_best_match t rest (Some attr)
        | Some { attr_name = name'; _ } ->
          let len = String.length name.txt in
          let len' = String.length name'.txt in
          if len > len' then
            find_best_match t rest (Some attr)
          else if len < len' then
            find_best_match t rest longest_match
          else
            Location.raise_errorf ~loc:name.loc "Duplicated attribute"
      end else
        find_best_match t rest longest_match
  in
  fun t attributes ->
    find_best_match t attributes None
;;

let convert ?(do_mark_as_seen = true) pattern attr =
  if do_mark_as_seen then mark_as_seen attr;
  let (Payload_parser (pattern, k)) = pattern in
  Ast_pattern.parse pattern (Common.loc_of_payload attr) attr.attr_payload
    (k ~name_loc:attr.attr_name.loc)
;;

let get t ?mark_as_seen:do_mark_as_seen x =
  let attrs = Context.get_attributes t.context x in
  match get_internal t attrs with
  | None -> None
  | Some attr -> Some (convert t.payload attr ?do_mark_as_seen)
;;

let consume t x =
  let attrs = Context.get_attributes t.context x in
  match get_internal t attrs with
  | None -> None
  | Some attr ->
    let attrs = List.filter attrs ~f:(fun attr' -> not (attr == attr')) in
    let x = Context.set_attributes t.context x attrs in
    Some (x, convert t.payload attr)
;;

let remove_seen (type a) (context : a Context.t) packeds (x : a) =
  let attrs = Context.get_attributes context x in
  let matched =
    let rec loop acc = function
      | [] -> acc
      | T t :: rest ->
        if Context.equal t.context context then
          match get_internal t attrs with
          | None      -> loop acc rest
          | Some attr ->
            let name = attr.attr_name in
            if Attribute_table.mem not_seen name then
              loop acc rest
            else
              loop (attr :: acc) rest
        else
          loop acc rest
    in
    loop [] packeds
  in
  let attrs =
    List.filter attrs ~f:(fun attr' ->
      not (List.memq ~set:matched attr'))
  in
  Context.set_attributes context x attrs
;;

let pattern t p =
  let f = Ast_pattern.to_func p in
  Ast_pattern.of_func (fun ctx loc x k ->
    match consume t x with
    | None        -> f ctx loc x (k None)
    | Some (x, v) -> f ctx loc x (k (Some v))
  )
;;

module Floating = struct
  module Context = Floating_context

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

  let name t = Name.Pattern.name t.name

  let declare name context pattern k =
    Name.Registrar.register ~kind:`Attribute registrar (Floating context) name;
    { name = Name.Pattern.make name
    ; context
    ; payload = Payload_parser (pattern, fun ~name_loc:_ -> k)
    }
  ;;

  let convert ts x =
    match ts with
    | [] -> None
    | { context; _ } :: _ ->
      assert (List.for_all ts ~f:(fun t -> Context.equal t.context context));
      let attr = Context.get_attribute context x in
      let name = attr.attr_name in
      match List.filter ts ~f:(fun t -> Name.Pattern.matches t.name name.txt) with
      | [] -> None
      | [t] -> Some (convert t.payload attr)
      | l ->
        Location.raise_errorf ~loc:name.loc
          "Multiple match for floating attributes: %s"
          (String.concat ~sep:", " (List.map l ~f:(fun t -> Name.Pattern.name t.name)))
  ;;
end

let check_attribute registrar context name =
  if not (Name.Whitelisted.is_whitelisted ~kind:`Attribute name.txt
          || Name.ignore_checks name.txt)
  && Attribute_table.mem not_seen name then
    let white_list = Name.Whitelisted.get_attribute_list () in
    Name.Registrar.raise_errorf registrar context ~white_list
      "Attribute `%s' was not used" name
;;


let check_unused = object(self)
  inherit Ast_traverse.iter as super

  method! attribute { attr_name = name; _ } =
    Location.raise_errorf ~loc:name.loc
      "attribute not expected here, Ppxlib.Attribute needs updating!"

  method private check_node : type a. a Context.t -> a -> a = fun context node ->
    let attrs = Context.get_attributes context node in
    match attrs with
    | [] -> node
    | _  ->
      List.iter attrs ~f:(fun ({ attr_name = name; attr_payload = payload; _ } as attr) ->
        self#payload payload;
        check_attribute registrar (On_item context) name;
        (* If we allow the attribute to pass through, mark it as seen *)
        mark_as_seen attr);
      Context.set_attributes context node []

  method private check_floating : type a. a Floating.Context.t -> a -> a
    = fun context node ->
      match Floating.Context.get_attribute_if_is_floating_node context node with
      | None -> node
      | Some  ({ attr_name = name; attr_payload = payload; _ } as attr) ->
        self#payload payload;
        check_attribute registrar (Floating context) name;
        mark_as_seen attr;
        Floating.Context.replace_by_dummy context node

  method! label_declaration       x = super#label_declaration       (self#check_node Label_declaration       x)
  method! constructor_declaration x = super#constructor_declaration (self#check_node Constructor_declaration x)
  method! type_declaration        x = super#type_declaration        (self#check_node Type_declaration        x)
  method! type_extension          x = super#type_extension          (self#check_node Type_extension          x)
  method! type_exception          x = super#type_exception          (self#check_node Type_exception          x)
  method! extension_constructor   x = super#extension_constructor   (self#check_node Extension_constructor   x)
  method! pattern                 x = super#pattern                 (self#check_node Pattern                 x)
  method! core_type               x = super#core_type               (self#check_node Core_type               x)
  method! expression              x = super#expression              (self#check_node Expression              x)
  method! value_description       x = super#value_description       (self#check_node Value_description       x)
  method! class_type              x = super#class_type              (self#check_node Class_type              x)
  method! class_infos f           x = super#class_infos f           (self#check_node Class_infos             x)
  method! class_expr              x = super#class_expr              (self#check_node Class_expr              x)
  method! module_type             x = super#module_type             (self#check_node Module_type             x)
  method! module_declaration      x = super#module_declaration      (self#check_node Module_declaration      x)
  method! module_type_declaration x = super#module_type_declaration (self#check_node Module_type_declaration x)
  method! open_description        x = super#open_description        (self#check_node Open_description        x)
  method! open_declaration        x = super#open_declaration        (self#check_node Open_declaration        x)
  method! include_infos f         x = super#include_infos f         (self#check_node Include_infos           x)
  method! module_expr             x = super#module_expr             (self#check_node Module_expr             x)
  method! value_binding           x = super#value_binding           (self#check_node Value_binding           x)
  method! module_binding          x = super#module_binding          (self#check_node Module_binding          x)

  method! class_field x =
    let x = self#check_node     Class_field x in
    let x = self#check_floating Class_field x in
    super#class_field x

  method! class_type_field x =
    let x = self#check_node     Class_type_field x in
    let x = self#check_floating Class_type_field x in
    super#class_type_field x

  method! row_field x =
    let x =
      match x.prf_desc with
      | Rtag _ -> self#check_node Rtag x
      | _      -> x
    in
    super#row_field x

  method! core_type_desc x =
    let x =
      match x with
      | Ptyp_object (fields, closed_flag) ->
        let fields = List.map fields ~f:(self#check_node Object_type_field) in
        Ptyp_object (fields, closed_flag)
      | _ -> x
    in
    super#core_type_desc x

  method! structure_item item =
    let item = self#check_floating Structure_item item in
    let item =
      match item.pstr_desc with
      | Pstr_eval      _ -> self#check_node Pstr_eval      item
      | Pstr_extension _ -> self#check_node Pstr_extension item
      | _                -> item
    in
    super#structure_item item

  method! signature_item item =
    let item = self#check_floating Signature_item item in
    let item =
      match item.psig_desc with
      | Psig_extension _ -> self#check_node Psig_extension item
      | _                -> item
    in
    super#signature_item item
end

let reset_checks () = Attribute_table.clear not_seen

let collect = object
  inherit Ast_traverse.iter as super

  method! attribute ({ attr_name = name; attr_payload = payload; _ } as attr) =
    let loc = Common.loc_of_attribute attr in
    super#payload payload;
    Attribute_table.add not_seen name loc
end

let check_all_seen () =
  let fail name loc =
    let txt = name.txt in
    if not (Name.ignore_checks txt) then
      Location.raise_errorf ~loc "Attribute `%s' was silently dropped" txt
  in
  Attribute_table.iter fail not_seen
;;

let remove_attributes_present_in table = object
  inherit Ast_traverse.iter as super

  method!  attribute { attr_name = name; attr_payload = payload; _ } =
    super#payload payload;
    Attribute_table.remove table name
end

let copy_of_not_seen () =
  let copy = Attribute_table.create (Attribute_table.length not_seen) in
  Attribute_table.iter (Attribute_table.add copy) not_seen;
  copy
;;

let dropped_so_far_structure st =
  let table = copy_of_not_seen () in
  (remove_attributes_present_in table)#structure st;
  Attribute_table.fold (fun name loc acc -> { txt = name.txt; loc } :: acc) table []
;;

let dropped_so_far_signature sg =
  let table = copy_of_not_seen () in
  (remove_attributes_present_in table)#signature sg;
  Attribute_table.fold (fun name loc acc -> { txt = name.txt; loc } :: acc) table []
;;