(**************************************************************************)
(*                                                                        *)
(*                                 OCaml                                  *)
(*                                                                        *)
(*                               Leo White                                *)
(*                                                                        *)
(*   Copyright 1996 Institut National de Recherche en Informatique et     *)
(*     en Automatique.                                                    *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

open Location

(* Docstrings *)

(* A docstring is "attached" if it has been inserted in the AST. This
   is used for generating unexpected docstring warnings. *)
type ds_attached =
  | Unattached (* Not yet attached anything.*)
  | Info (* Attached to a field or constructor. *)
  | Docs (* Attached to an item or as floating text. *)

(* A docstring is "associated" with an item if there are no blank lines between
   them. This is used for generating docstring ambiguity warnings. *)
type ds_associated =
  | Zero (* Not associated with an item *)
  | One (* Associated with one item *)
  | Many (* Associated with multiple items (ambiguity) *)

type docstring = {
  ds_body: string;
  ds_loc: Location.t;
  mutable ds_attached: ds_attached;
  mutable ds_associated: ds_associated;
}

(* List of docstrings *)

let docstrings : docstring list ref = ref []

(* Warn for unused and ambiguous docstrings *)

let warn_bad_docstrings () =
  if Warnings.is_active (Warnings.Bad_docstring true) then
    List.iter
      (fun ds ->
        match ds.ds_attached with
        | Info -> ()
        | Unattached -> prerr_warning ds.ds_loc (Warnings.Bad_docstring true)
        | Docs -> (
          match ds.ds_associated with
          | Zero | One -> ()
          | Many -> prerr_warning ds.ds_loc (Warnings.Bad_docstring false)))
      (List.rev !docstrings)

(* Docstring constructors and destructors *)

let docstring body loc =
  let ds =
    {
      ds_body = body;
      ds_loc = loc;
      ds_attached = Unattached;
      ds_associated = Zero;
    }
  in
  ds

let register ds = docstrings := ds :: !docstrings

let docstring_body ds = ds.ds_body

let docstring_loc ds = ds.ds_loc

(* Docstrings attached to items *)

type docs = {docs_pre: docstring option; docs_post: docstring option}

let empty_docs = {docs_pre = None; docs_post = None}

let doc_loc = {txt = "ocaml.doc"; loc = Location.none}

let docs_attr ds =
  let open Parsetree in
  let exp =
    {
      pexp_desc = Pexp_constant (Pconst_string (ds.ds_body, None));
      pexp_loc = ds.ds_loc;
      pexp_attributes = [];
    }
  in
  let item = {pstr_desc = Pstr_eval (exp, []); pstr_loc = exp.pexp_loc} in
  (doc_loc, PStr [item])

let add_docs_attrs docs attrs =
  let attrs =
    match docs.docs_pre with
    | None | Some {ds_body = ""; _} -> attrs
    | Some ds -> docs_attr ds :: attrs
  in
  let attrs =
    match docs.docs_post with
    | None | Some {ds_body = ""; _} -> attrs
    | Some ds -> attrs @ [docs_attr ds]
  in
  attrs

(* Docstrings attached to constructors or fields *)

type info = docstring option

let empty_info = None

let info_attr = docs_attr

let add_info_attrs info attrs =
  match info with
  | None | Some {ds_body = ""; _} -> attrs
  | Some ds -> attrs @ [info_attr ds]

(* Docstrings not attached to a specific item *)

type text = docstring list

let empty_text = []
let empty_text_lazy = lazy []

let text_loc = {txt = "ocaml.text"; loc = Location.none}

let text_attr ds =
  let open Parsetree in
  let exp =
    {
      pexp_desc = Pexp_constant (Pconst_string (ds.ds_body, None));
      pexp_loc = ds.ds_loc;
      pexp_attributes = [];
    }
  in
  let item = {pstr_desc = Pstr_eval (exp, []); pstr_loc = exp.pexp_loc} in
  (text_loc, PStr [item])

let add_text_attrs dsl attrs =
  let fdsl =
    Ext_list.filter dsl (function
      | {ds_body = ""} -> false
      | _ -> true)
  in
  List.map text_attr fdsl @ attrs

(* Find the first non-info docstring in a list, attach it and return it *)
let get_docstring ~info dsl =
  let rec loop = function
    | [] -> None
    | {ds_attached = Info; _} :: rest -> loop rest
    | ds :: _ ->
      ds.ds_attached <- (if info then Info else Docs);
      Some ds
  in
  loop dsl

(* Find all the non-info docstrings in a list, attach them and return them *)
let get_docstrings dsl =
  let rec loop acc = function
    | [] -> List.rev acc
    | {ds_attached = Info; _} :: rest -> loop acc rest
    | ds :: rest ->
      ds.ds_attached <- Docs;
      loop (ds :: acc) rest
  in
  loop [] dsl

(* "Associate" all the docstrings in a list *)
let associate_docstrings dsl =
  List.iter
    (fun ds ->
      match ds.ds_associated with
      | Zero -> ds.ds_associated <- One
      | One | Many -> ds.ds_associated <- Many)
    dsl

(* Map from positions to pre docstrings *)

let pre_table : (Lexing.position, docstring list) Hashtbl.t = Hashtbl.create 50

let set_pre_docstrings pos dsl = if dsl <> [] then Hashtbl.add pre_table pos dsl

let get_pre_docs pos =
  try
    let dsl = Hashtbl.find pre_table pos in
    associate_docstrings dsl;
    get_docstring ~info:false dsl
  with Not_found -> None

let mark_pre_docs pos =
  try
    let dsl = Hashtbl.find pre_table pos in
    associate_docstrings dsl
  with Not_found -> ()

(* Map from positions to post docstrings *)

let post_table : (Lexing.position, docstring list) Hashtbl.t = Hashtbl.create 50

let set_post_docstrings pos dsl =
  if dsl <> [] then Hashtbl.add post_table pos dsl

let get_post_docs pos =
  try
    let dsl = Hashtbl.find post_table pos in
    associate_docstrings dsl;
    get_docstring ~info:false dsl
  with Not_found -> None

let mark_post_docs pos =
  try
    let dsl = Hashtbl.find post_table pos in
    associate_docstrings dsl
  with Not_found -> ()

let get_info pos =
  try
    let dsl = Hashtbl.find post_table pos in
    get_docstring ~info:true dsl
  with Not_found -> None

(* Map from positions to floating docstrings *)

let floating_table : (Lexing.position, docstring list) Hashtbl.t =
  Hashtbl.create 50

let set_floating_docstrings pos dsl =
  if dsl <> [] then Hashtbl.add floating_table pos dsl

let get_text pos =
  try
    let dsl = Hashtbl.find floating_table pos in
    get_docstrings dsl
  with Not_found -> []

(* Maps from positions to extra docstrings *)

let pre_extra_table : (Lexing.position, docstring list) Hashtbl.t =
  Hashtbl.create 50

let set_pre_extra_docstrings pos dsl =
  if dsl <> [] then Hashtbl.add pre_extra_table pos dsl

let get_pre_extra_text pos =
  try
    let dsl = Hashtbl.find pre_extra_table pos in
    get_docstrings dsl
  with Not_found -> []

let post_extra_table : (Lexing.position, docstring list) Hashtbl.t =
  Hashtbl.create 50

let set_post_extra_docstrings pos dsl =
  if dsl <> [] then Hashtbl.add post_extra_table pos dsl

let get_post_extra_text pos =
  try
    let dsl = Hashtbl.find post_extra_table pos in
    get_docstrings dsl
  with Not_found -> []

(* Docstrings from parser actions *)

let symbol_docs () =
  {
    docs_pre = get_pre_docs (Parsing.symbol_start_pos ());
    docs_post = get_post_docs (Parsing.symbol_end_pos ());
  }

let symbol_docs_lazy () =
  let p1 = Parsing.symbol_start_pos () in
  let p2 = Parsing.symbol_end_pos () in
  lazy {docs_pre = get_pre_docs p1; docs_post = get_post_docs p2}

let rhs_docs pos1 pos2 =
  {
    docs_pre = get_pre_docs (Parsing.rhs_start_pos pos1);
    docs_post = get_post_docs (Parsing.rhs_end_pos pos2);
  }

let rhs_docs_lazy pos1 pos2 =
  let p1 = Parsing.rhs_start_pos pos1 in
  let p2 = Parsing.rhs_end_pos pos2 in
  lazy {docs_pre = get_pre_docs p1; docs_post = get_post_docs p2}

let mark_symbol_docs () =
  mark_pre_docs (Parsing.symbol_start_pos ());
  mark_post_docs (Parsing.symbol_end_pos ())

let mark_rhs_docs pos1 pos2 =
  mark_pre_docs (Parsing.rhs_start_pos pos1);
  mark_post_docs (Parsing.rhs_end_pos pos2)

let symbol_info () = get_info (Parsing.symbol_end_pos ())

let rhs_info pos = get_info (Parsing.rhs_end_pos pos)

let symbol_text () = get_text (Parsing.symbol_start_pos ())

let symbol_text_lazy () =
  let pos = Parsing.symbol_start_pos () in
  lazy (get_text pos)

let rhs_text pos = get_text (Parsing.rhs_start_pos pos)

let rhs_text_lazy pos =
  let pos = Parsing.rhs_start_pos pos in
  lazy (get_text pos)

let symbol_pre_extra_text () = get_pre_extra_text (Parsing.symbol_start_pos ())

let symbol_post_extra_text () = get_post_extra_text (Parsing.symbol_end_pos ())

let rhs_pre_extra_text pos = get_pre_extra_text (Parsing.rhs_start_pos pos)

let rhs_post_extra_text pos = get_post_extra_text (Parsing.rhs_end_pos pos)

(* (Re)Initialise all comment state *)

let init () =
  docstrings := [];
  Hashtbl.reset pre_table;
  Hashtbl.reset post_table;
  Hashtbl.reset floating_table;
  Hashtbl.reset pre_extra_table;
  Hashtbl.reset post_extra_table
