Source file signature_help.ml

open Std

let { Logger.log } = Logger.for_section "signature-help"

type parameter_info =
  { label : Asttypes.arg_label;
    param_start : int;
    param_end : int;
    argument : Typedtree.apply_arg
  }

type application_signature =
  { function_name : string option;
    function_position : Msource.position;
    signature : string;
    parameters : parameter_info list;
    active_param : int option
  }

(* extract a properly parenthesized identifier from (expression_desc (Texp_ident
   (Longident))) *)
let extract_ident (exp_desc : Typedtree.expression_desc) =
  let rec longident ppf : Longident.t -> unit = function
    | Lident s -> Format.fprintf ppf "%s" (Misc_utils.parenthesize_name s)
    | Ldot (p, s) ->
      Format.fprintf ppf "%a.%s" longident p.txt
        (Misc_utils.parenthesize_name s.txt)
    | Lapply (p1, p2) ->
      Format.fprintf ppf "%a(%a)" longident p1.txt longident p2.txt
  in
  match exp_desc with
  | Texp_ident (_, { txt = li; _ }, _) ->
    let ppf, to_string = Format.to_string () in
    longident ppf li;
    Some (to_string ())
  | _ -> None

(* Type variables shared across arguments should all be printed with the same
   name. [Printtyp.type_scheme] ensure that a name is unique within a given
   type, but not across different invocations. [reset] followed by calls to
   [mark_loops] and [type_sch] provide that *)
let pp_type env ppf ty =
  let module Printtyp = Type_utils.Printtyp in
  Printtyp.wrap_printing_env env ~verbosity:(Lvl 0) (fun () ->
      Printtyp.shared_type_scheme ppf ty)

let rec type_is_arrow ty =
  match Types.get_desc ty with
  | Tarrow _ -> true
  | Tpoly (ty, _) -> type_is_arrow ty
  | _ -> false

(* surround function types in parentheses *)
let pp_parameter_type env ppf ty =
  if type_is_arrow ty then Format.fprintf ppf "(%a)" (pp_type env) ty
  else pp_type env ppf ty

(* print parameter labels and types *)
let pp_parameter env label ppf ty =
  match (label : Asttypes.arg_label) with
  | Nolabel -> pp_parameter_type env ppf ty
  | Labelled l -> Format.fprintf ppf "%s:%a" l (pp_parameter_type env) ty
  | Optional l ->
    (* unwrap option for optional labels the same way as
       [Raw_compat.labels_of_application] *)
    let unwrap_option ty =
      (* Since OCaml 5.5 the argument type of a [Tarrow] is wrapped in a mono
         [Tpoly] node, look through it before unwrapping the option. *)
      let ty =
        match Types.get_desc ty with
        | Types.Tpoly (ty, _) -> ty
        | _ -> ty
      in
      match Types.get_desc ty with
      | Types.Tconstr (path, [ ty ], _) when Path.same path Predef.path_option
        -> ty
      | _ -> ty
    in
    Format.fprintf ppf "?%s:%a" l (pp_parameter_type env) (unwrap_option ty)

(* record buffer offsets to be able to underline parameter types *)
let print_parameter_offset ~arg:argument ppf buffer env label ty =
  let param_start = Buffer.length buffer in
  Format.fprintf ppf "%a%!" (pp_parameter env label) ty;
  let param_end = Buffer.length buffer in
  Format.pp_print_string ppf " -> ";
  Format.pp_print_flush ppf ();
  { label; param_start; param_end; argument }

(* This function preprocesses the signature and associate already assigned
   arguments to the corresponding parameter. (They should always be in the correct
   order in the typedtree, even if they are not in order in the source file.) *)
let separate_function_signature ~args (e : Typedtree.expression) =
  Out_type.reset ();
  let buffer = Buffer.create 16 in
  let ppf = Format.formatter_of_buffer buffer in
  let rec separate ?(parameters = []) args ty =
    let desc =
      match args with
      | _ :: _ -> Types.get_desc (Ctype.expand_head e.exp_env ty)
      (* expand the type only if there are remaining arguments *)
      | [] -> Types.get_desc ty
    in
    match (args, desc) with
    | (_l, arg) :: args, Tarrow (label, ty1, ty2, _) ->
      let parameter =
        print_parameter_offset ~arg ppf buffer e.exp_env label ty1
      in
      separate args ty2 ~parameters:(parameter :: parameters)
    | [], Tarrow (label, ty1, ty2, _) ->
      let parameter =
        print_parameter_offset ~arg:(Omitted ()) ppf buffer e.exp_env label ty1
      in
      separate args ty2 ~parameters:(parameter :: parameters)
    (* end of function type, print remaining type without recording offsets *)
    | _ ->
      Format.fprintf ppf "%a%!" (pp_type e.exp_env) ty;
      { function_name = extract_ident e.exp_desc;
        function_position = `Offset e.exp_loc.loc_end.pos_cnum;
        signature = Buffer.contents buffer;
        parameters = List.rev parameters;
        active_param = None
      }
  in
  let expanded_ty = Ctype.expand_head e.exp_env e.exp_type in
  separate args expanded_ty

let active_parameter_by_arg ~arg params =
  let find_by_arg = function
    | { argument = Arg a; _ } when a == arg -> true
    | _ -> false
  in
  try Some (List.index params ~f:find_by_arg) with Not_found -> None

let first_unassigned_argument params =
  let positional = function
    | { argument = Omitted (); label = Asttypes.Nolabel; _ } -> true
    | _ -> false
  in
  let labelled = function
    | { argument = Omitted (); label = Asttypes.Labelled _ | Optional _; _ } ->
      true
    | _ -> false
  in
  try Some (List.index params ~f:positional)
  with Not_found -> (
    try Some (List.index params ~f:labelled) with Not_found -> None)

let active_parameter_by_prefix ~prefix params =
  let common = function
    | Asttypes.Nolabel -> Some 0
    | l
      when String.is_prefixed ~by:"~" prefix
           || String.is_prefixed ~by:"?" prefix ->
      Some (String.common_prefix_len (Btype.prefixed_label_name l) prefix)
    | _ -> None
  in
  let is_omitted = function
    | Typedtree.Omitted _ -> true
    | _ -> false
  in
  let rec find_by_prefix ?(i = 0) ?longest_len ?longest_i = function
    | [] -> longest_i
    | p :: ps when is_omitted p.argument -> (
      (* The search is performed only on the arguments not already given in the parameters. *)
      match (common p.label, longest_len) with
      | Some common_len, Some longest_len when common_len > longest_len ->
        find_by_prefix ps ~i:(succ i) ~longest_len:common_len ~longest_i:i
      | Some common_len, None ->
        find_by_prefix ps ~i:(succ i) ~longest_len:common_len ~longest_i:i
      | _ -> find_by_prefix ps ~i:(succ i) ?longest_len ?longest_i)
    | _ :: ps -> find_by_prefix ps ~i:(succ i) ?longest_len ?longest_i
  in
  find_by_prefix params

let is_arrow env t =
  match Types.get_desc (Ctype.expand_head env t) with
  | Tarrow _ -> true
  | _ -> false

let application_signature ~prefix ~cursor node =
  match node with
  | (_, Browse_raw.Expression arg)
    :: ( _,
         Expression
           { exp_desc = Texp_apply (({ exp_type; exp_env; _ } as e), args); _ }
       )
    :: _
    when is_arrow exp_env exp_type ->
    log ~title:"application_signature" "Last arg:\n%a" Logger.fmt (fun fmt ->
        Printtyped.expression fmt arg);
    let result = separate_function_signature e ~args in
    let active_param =
      if prefix = "" && Lexing.compare_pos cursor arg.exp_loc.loc_end > 0 then begin
        (* If the cursor is placed after the last arg it means that a whitespace
           was inserted and we want to underline the next argument. *)
        log ~title:"application_signature"
          "Current cursor position is after the last argument";
        first_unassigned_argument result.parameters
      end
      else
        (* If not, we identify the argument which is being written *)
        let active_param = active_parameter_by_arg ~arg result.parameters in
        match active_param with
        | Some _ as ap -> ap
        | None -> active_parameter_by_prefix ~prefix result.parameters
    in
    Some { result with active_param }
  | (_, Expression ({ exp_type; exp_env; _ } as e)) :: _
    when is_arrow exp_env exp_type ->
    (* provide signature information directly after an unapplied function-type
       value *)
    let result = separate_function_signature e ~args:[] in
    let active_param = active_parameter_by_prefix ~prefix result.parameters in
    Some { result with active_param }
  | _ ->
    let rec find_smallest_arrow_before_pos (e : Typedtree.expression) pos
        (acc : Typedtree.expression) =
      if Lexing.compare_pos e.exp_loc.loc_start pos > 0 then
        match acc.exp_desc with
        | Texp_let (_, vlist, _) ->
          let v =
            List.find_opt
              ~f:(fun (value_binding : Typedtree.value_binding) ->
                is_arrow value_binding.vb_expr.exp_env
                  value_binding.vb_expr.exp_type)
              vlist
          in
          Option.map ~f:(fun (v : Typedtree.value_binding) -> v.vb_expr) v
        | _ -> None
      else
        match e.exp_desc with
        | Texp_let (_, _, next) -> find_smallest_arrow_before_pos next pos e
        | _ -> None
    in
    let expressions =
      List.filter_map
        ~f:(fun n ->
          match n with
          | _, Browse_raw.Expression e -> Some e
          | _ -> None)
        node
    in
    if expressions = [] then None
    else
      let last_node = List.hd (List.rev expressions) in
      let smallest_frag_opt =
        find_smallest_arrow_before_pos last_node cursor last_node
      in
      Option.map
        ~f:(fun frag ->
          let result = separate_function_signature frag ~args:[] in
          let active_param =
            active_parameter_by_prefix ~prefix result.parameters
          in
          { result with active_param })
        smallest_frag_opt

let prefix_of_position ~short_path source position =
  match Msource.text source with
  | "" -> ""
  | text ->
    let from =
      let (`Offset index) = Msource.get_offset source position in
      min (String.length text - 1) (index - 1)
    in
    let pos =
      let should_terminate = ref false in
      let has_seen_dot = ref false in
      let is_prefix_char c =
        if !should_terminate then false
        else
          match c with
          | 'a' .. 'z'
          | 'A' .. 'Z'
          | '0' .. '9'
          | '\''
          | '_'
          (* Infix function characters *)
          | '$'
          | '&'
          | '*'
          | '+'
          | '-'
          | '/'
          | '='
          | '>'
          | '@'
          | '^'
          | '!'
          | '?'
          | '%'
          | '<'
          | ':'
          | '~'
          | '#' -> true
          | '`' ->
            if !has_seen_dot then false
            else (
              should_terminate := true;
              true)
          | '.' ->
            has_seen_dot := true;
            not short_path
          | _ -> false
      in
      String.rfindi text ~from ~f:(fun c -> not (is_prefix_char c))
    in
    let pos =
      match pos with
      | None -> 0
      | Some pos -> pos + 1
    in
    let len = from - pos + 1 in
    let reconstructed_prefix = String.sub text ~pos ~len in
    (* if we reconstructed [~f:ignore] or [?f:ignore], we should take only
       [ignore], so: *)
    log ~title:"prefix_of_position" "%S" reconstructed_prefix;
    if
      String.is_prefix reconstructed_prefix ~prefix:"~"
      || String.is_prefix reconstructed_prefix ~prefix:"?"
    then
      match String.lsplit2 reconstructed_prefix ~on:':' with
      | Some (_, s) -> s
      | None -> reconstructed_prefix
    else reconstructed_prefix