(** Extend Search to subsume SearchAbout, also adding hidden Type coercions. *)

(* Main prefilter *)

DECLARE PLUGIN "ssrsearch_plugin"

{

module CoqConstr = Constr
open CoqConstr
open Constrexpr
open Evar_kinds
open Glob_term
open Ltac_plugin
open Notation_ops
open Notation_term
open Pcoq.Prim
open Pcoq.Constr
open Pp
open Ppconstr
open Printer
open Stdarg
open Ssreflect_plugin.Ssrprinters
open Ssreflect_plugin.Ssrcommon
open Ssreflect_plugin.Ssrparser
open Termops
open Util

type raw_glob_search_about_item =
  | RGlobSearchSubPattern of constr_expr
  | RGlobSearchString of Loc.t * string * string option

let pr_search_item env sigma = function
  | RGlobSearchString (_,s,_) -> str s
  | RGlobSearchSubPattern p -> pr_constr_expr env sigma p

let wit_ssr_searchitem = add_genarg "ssr_searchitem" pr_search_item

let pr_ssr_search_item env sigma _ _ _ = pr_search_item env sigma

(* Workaround the notation API that can only print notations *)

let is_ident s = try CLexer.check_ident s; true with _ -> false

let is_ident_part s = is_ident ("H" ^ s)

let interp_search_notation ?loc tag okey =
  let err msg = CErrors.user_err ?loc ~hdr:"interp_search_notation" msg in
  let mk_pntn s for_key =
    let n = String.length s in
    let s' = Bytes.make (n + 2) ' ' in
    let rec loop i i' =
      if i >= n then s', i' - 2 else if s.[i] = ' ' then loop (i + 1) i' else
      let j = try String.index_from s (i + 1) ' ' with _ -> n in
      let m = j - i in
      if s.[i] = '\'' && i < j - 2 && s.[j - 1] = '\'' then
        (String.blit s (i + 1) s' i' (m - 2); loop (j + 1) (i' + m - 1))
      else if for_key && is_ident (String.sub s i m) then
         (Bytes.set s' i' '_'; loop (j + 1) (i' + 2))
      else (String.blit s i s' i' m; loop (j + 1) (i' + m + 1)) in
    loop 0 1 in
  let trim_ntn (pntn, m) = (InConstrEntry,Bytes.sub_string pntn 1 (max 0 m)) in
  let pr_ntn ntn = str "(" ++ Notation.pr_notation ntn ++ str ")" in
  let pr_and_list pr = function
    | [x] -> pr x
    | x :: lx -> pr_list pr_comma pr lx ++ pr_comma () ++ str "and " ++ pr x
    | [] -> mt () in
  let pr_sc sc = str (if sc = "" then "independently" else sc) in
  let pr_scs = function
    | [""] -> pr_sc ""
    | scs -> str "in " ++ pr_and_list pr_sc scs in
  let generator, pr_tag_sc =
    let ign _ = mt () in match okey with
  | Some key ->
    let sc = Notation.find_delimiters_scope ?loc key in
    let pr_sc s_in = str s_in ++ spc() ++ str sc ++ pr_comma() in
    Notation.pr_scope ign sc, pr_sc
  | None -> Notation.pr_scopes ign, ign in
  let qtag s_in = pr_tag_sc s_in ++ qstring tag ++ spc()in
  let ptag, ttag =
    let ptag, m = mk_pntn tag false in
    if m <= 0 then err (str "empty notation fragment");
    ptag, trim_ntn (ptag, m) in
  let last = ref "" and last_sc = ref "" in
  let scs = ref [] and ntns = ref [] in
  let push_sc sc = match !scs with
  | "" :: scs' ->  scs := "" :: sc :: scs'
  | scs' -> scs := sc :: scs' in
  let get s _ _ = match !last with
  | "Scope " -> last_sc := s; last := ""
  | "Lonely notation" -> last_sc := ""; last := ""
  | "\"" ->
      let pntn, m = mk_pntn s true in
      if String.string_contains ~where:(Bytes.to_string pntn) ~what:(Bytes.to_string ptag) then begin
        let ntn = trim_ntn (pntn, m) in
        match !ntns with
        | [] -> ntns := [ntn]; scs := [!last_sc]
        | ntn' :: _ when ntn' = ntn -> push_sc !last_sc
        | _ when ntn = ttag -> ntns := ntn :: !ntns; scs := [!last_sc]
        | _ :: ntns' when List.mem ntn ntns' -> ()
        | ntn' :: ntns' -> ntns := ntn' :: ntn :: ntns'
      end;
      last := ""
  | _ -> last := s in
  pp_with (Format.make_formatter get (fun _ -> ())) generator;
  let ntn = match !ntns with
  | [] ->
    err (hov 0 (qtag "in" ++ str "does not occur in any notation"))
  | ntn :: ntns' when ntn = ttag ->
    if ntns' <> [] then begin
      let pr_ntns' = pr_and_list pr_ntn ntns' in
      Feedback.msg_warning (hov 4 (qtag "In" ++ str "also occurs in " ++ pr_ntns'))
    end; ntn
  | [ntn] ->
    Feedback.msg_notice (hov 4 (qtag "In" ++ str "is part of notation " ++ pr_ntn ntn)); ntn
  | ntns' ->
    let e = str "occurs in" ++ spc() ++ pr_and_list pr_ntn ntns' in
    err (hov 4 (str "ambiguous: " ++ qtag "in" ++ e)) in
  let (nvars, body), ((_, pat), osc) = match !scs with
  | [sc] -> Notation.interp_notation ?loc ntn (None, [sc])
  | scs' ->
    try Notation.interp_notation ?loc ntn (None, []) with _ ->
    let e = pr_ntn ntn ++ spc() ++ str "is defined " ++ pr_scs scs' in
    err (hov 4 (str "ambiguous: " ++ pr_tag_sc "in" ++ e)) in
  let sc = Option.default "" osc in
  let _ =
    let m_sc =
      if osc <> None then str "In " ++ str sc ++ pr_comma() else mt() in
    let ntn_pat = trim_ntn (mk_pntn pat false) in
    let rbody = glob_constr_of_notation_constr ?loc body in
    let m_body = hov 0 (Constrextern.without_symbols prl_glob_constr rbody) in
    let m = m_sc ++ pr_ntn ntn_pat ++ spc () ++ str "denotes " ++ m_body in
    Feedback.msg_notice (hov 0 m) in
  if List.length !scs > 1 then
    let scs' = List.remove (=) sc !scs in
    let w = pr_ntn ntn ++ str " is also defined " ++ pr_scs scs' in
    Feedback.msg_warning (hov 4 w)
  else if String.string_contains ~where:(snd ntn) ~what:" .. " then
    err (pr_ntn ntn ++ str " is an n-ary notation");
  let nvars = List.filter (fun (_,(_,typ)) -> typ = NtnTypeConstr) nvars in
  let rec sub () = function
  | NVar x when List.mem_assoc x nvars -> DAst.make ?loc @@ GPatVar (FirstOrderPatVar x)
  | c ->
    glob_constr_of_notation_constr_with_binders ?loc (fun _ x t -> (), None, x, Explicit, t) sub () c in
  let _, npat = Patternops.pattern_of_glob_constr (sub () body) in
  Search.GlobSearchSubPattern (Vernacexpr.Anywhere,false,npat)

}

ARGUMENT EXTEND ssr_search_item TYPED AS ssr_searchitem
  PRINTED BY { pr_ssr_search_item env sigma }
  | [ string(s) ] -> {  RGlobSearchString (loc,s,None)  }
  | [ string(s) "%" preident(key) ] -> {  RGlobSearchString (loc,s,Some key)  }
  | [ constr_pattern(p) ] -> {  RGlobSearchSubPattern p  }
END

{

let pr_ssr_search_arg env sigma _ _ _ =
  let pr_item (b, p) = str (if b then "-" else "") ++ pr_search_item env sigma p in
  pr_list spc pr_item

}

ARGUMENT EXTEND ssr_search_arg TYPED AS (bool * ssr_searchitem) list
  PRINTED BY { pr_ssr_search_arg env sigma }
  | [ "-" ssr_search_item(p) ssr_search_arg(a) ] -> {  (false, p) :: a  }
  | [ ssr_search_item(p) ssr_search_arg(a) ] -> {  (true, p) :: a  }
  | [ ] -> {  []  }
END

{

(* Main type conclusion pattern filter *)

let rec splay_search_pattern na = function
  | Pattern.PApp (fp, args) -> splay_search_pattern (na + Array.length args) fp
  | Pattern.PLetIn (_, _, _, bp) -> splay_search_pattern na bp
  | Pattern.PRef hr -> hr, na
  | _ -> CErrors.user_err (Pp.str "no head constant in head search pattern")

let push_rels_assum l e =
  let l = List.map (fun (n,t) -> n, EConstr.Unsafe.to_constr t) l in
  push_rels_assum l e

let coerce_search_pattern_to_sort hpat =
  let env = Global.env () in
  let sigma = Evd.(from_env env) in
  let mkPApp fp n_imps args =
    let args' = Array.append (Array.make n_imps (Pattern.PMeta None)) args in
    Pattern.PApp (fp, args') in
  let hr, na = splay_search_pattern 0 hpat in
  let dc, ht =
    let hr, _ = Typeops.type_of_global_in_context env hr (* FIXME *) in
    Reductionops.splay_prod env sigma (EConstr.of_constr hr) in
  let np = List.length dc in
  if np < na then CErrors.user_err (Pp.str "too many arguments in head search pattern") else
  let hpat' = if np = na then hpat else mkPApp hpat (np - na) [||] in
  let warn () =
    Feedback.msg_warning (str "Listing only lemmas with conclusion matching " ++
      pr_constr_pattern_env env sigma hpat') in
  if EConstr.isSort sigma ht then begin warn (); true, hpat' end else
  let filter_head, coe_path =
    try
      let _, cp =
        Coercionops.lookup_path_to_sort_from (push_rels_assum dc env) sigma ht in
      warn ();
      true, cp
    with _ -> false, [] in
  let coerce hp coe_index =
    let coe_ref = coe_index.Coercionops.coe_value in
    try
      let n_imps = Option.get (Coercionops.hide_coercion coe_ref) in
      mkPApp (Pattern.PRef coe_ref) n_imps [|hp|]
    with Not_found | Option.IsNone ->
    errorstrm (str "need explicit coercion " ++ pr_global coe_ref ++ spc ()
            ++ str "to interpret head search pattern as type") in
  filter_head, List.fold_left coerce hpat' coe_path

let interp_head_pat hpat =
  let filter_head, p = coerce_search_pattern_to_sort hpat in
  let rec loop c = match CoqConstr.kind c with
  | Cast (c', _, _) -> loop c'
  | Prod (_, _, c') -> loop c'
  | LetIn (_, _, _, c') -> loop c'
  | _ ->
    let env = Global.env () in
    let sigma = Evd.from_env env in
    Constr_matching.is_matching env sigma p (EConstr.of_constr c) in
  filter_head, loop

let all_true _ = true

let rec interp_search_about args accu = match args with
| [] -> accu
| (flag, arg) :: rem ->
  fun gr kind env typ ->
    let ans = Search.search_filter arg gr kind env (Evd.from_env env) typ in
    (if flag then ans else not ans) && interp_search_about rem accu gr kind env typ

let interp_search_arg arg =
  let arg = List.map (fun (x,arg) -> x, match arg with
  | RGlobSearchString (loc,s,key) ->
      if is_ident_part s then Search.GlobSearchString s else
      interp_search_notation ~loc s key
  | RGlobSearchSubPattern p ->
    let env = Global.env () in
    let _, p = Constrintern.intern_constr_pattern env (Evd.from_env env) p in
    Search.GlobSearchSubPattern (Vernacexpr.Anywhere,false,p)) arg
  in
  let hpat, a1 = match arg with
  | (_, Search.GlobSearchSubPattern (Vernacexpr.Anywhere,false,Pattern.PMeta _)) :: a' -> all_true, a'
  | (true, Search.GlobSearchSubPattern (Vernacexpr.Anywhere,false,p)) :: a' ->
     let filter_head, p = interp_head_pat p in
     if filter_head then p, a' else all_true, arg
  | (_, (Search.GlobSearchSubPattern (Vernacexpr.(InHyp|InConcl),_,_)
        |Search.GlobSearchSubPattern (Vernacexpr.Anywhere,true,_))) :: a' -> CErrors.user_err (str "Unsupported.")
  | _ -> all_true, arg in
  let is_string =
    function (_, Search.GlobSearchString _) -> true | _ -> false in
  let a2, a3 = List.partition is_string a1 in
  interp_search_about (a2 @ a3) (fun gr kind env typ -> hpat typ)

(* Module path postfilter *)

let pr_modloc (b, m) = if b then str "-" ++ pr_qualid m else pr_qualid m

let wit_ssrmodloc = add_genarg "ssrmodloc" (fun env sigma -> pr_modloc)

let pr_ssr_modlocs _ _ _ ml =
  if ml = [] then str "" else spc () ++ str "in " ++ pr_list spc pr_modloc ml

}

ARGUMENT EXTEND ssr_modlocs TYPED AS ssrmodloc list PRINTED BY { pr_ssr_modlocs }
  | [ ] -> {  []  }
END

GRAMMAR EXTEND Gram
  GLOBAL: ssr_modlocs;
  modloc: [[ "-"; m = global -> { true, m } | m = global -> { false, m } ]];
  ssr_modlocs: [[ "in"; ml = LIST1 modloc -> { ml } ]];
END

{

let interp_modloc mr =
  let interp_mod (_, qid) =
    try Nametab.full_name_module qid with Not_found ->
    CErrors.user_err ?loc:qid.CAst.loc (str "No Module " ++ pr_qualid qid) in
  let mr_out, mr_in = List.partition fst mr in
  let interp_bmod b = function
  | [] -> fun _ _ _ _ _ -> true
  | rmods -> Search.module_filter (List.map interp_mod rmods, b) in
  let is_in = interp_bmod false mr_in and is_out = interp_bmod true mr_out in
  fun gr kind env typ -> is_in gr kind env (Evd.from_env env) typ && is_out gr kind env (Evd.from_env env) typ

(* The unified, extended vernacular "Search" command *)

let ssrdisplaysearch gr env t =
  let pr_res = pr_global gr ++ str ":" ++ spc () ++ pr_lconstr_env env Evd.empty t in
  Feedback.msg_notice (hov 2 pr_res ++ fnl ())

let deprecated_search =
  CWarnings.create
    ~name:"deprecated-ssr-search"
    ~category:"deprecated"
    (fun () -> Pp.(str"SSReflect's Search command is deprecated."))

}

VERNAC COMMAND EXTEND SsrSearchPattern CLASSIFIED AS QUERY
| [ "Search" ssr_search_arg(a) ssr_modlocs(mr) ] ->
  { deprecated_search ();
    let hpat = interp_search_arg a in
    let in_mod = interp_modloc mr in
    let post_filter gr kind env typ = in_mod gr kind env typ && hpat gr kind env typ in
    let display gr kind env typ =
      if post_filter gr kind env typ then ssrdisplaysearch gr env typ
    in
    let env = Global.env () in
    Search.generic_search env display }
END