livebook

defmodule Livebook.Intellisense.Docs do
  # This module is responsible for extracting and normalizing
  # information like documentation, signatures and specs.
  #
  # Note that we only extract the docs information when requested for
  # for the current node. For remote nodes, making several requests
  # for docs (which may be necessary if there are multiple modules)
  # adds to the overall latency. Remote intellisense is primarily used
  # with remote release nodes, which have docs stripped out anyway.

  @type member_info :: %{
          kind: member_kind(),
          name: atom(),
          arity: non_neg_integer(),
          from_default: boolean(),
          documentation: documentation(),
          signatures: list(signature()),
          specs: list(spec()),
          type_spec: type_spec(),
          meta: meta()
        }

  @type member_kind :: :function | :macro | :type

  @type documentation :: {format :: String.t(), content :: String.t()} | :hidden | nil

  @type signature :: String.t()

  @type meta :: map()

  @typedoc """
  A single spec annotation in the Erlang Abstract Format.
  """
  @type spec :: term()

  @typedoc """
  A tuple containing a single type annotation in the Erlang Abstract Format,
  tagged by its kind.
  """
  @type type_spec() :: {type_kind(), term()}
  @type type_kind() :: :type | :opaque

  @type definition ::
          {:module, module()} | {:function | :type, name :: atom(), arity :: pos_integer()}

  @doc """
  Fetches documentation for the given module if available.
  """
  @spec get_module_documentation(module(), node()) :: documentation()
  def get_module_documentation(module, node)

  def get_module_documentation(_module, node) when node != node(), do: nil

  def get_module_documentation(module, _node) do
    case Code.fetch_docs(module) do
      {:docs_v1, _, _, format, %{"en" => docstring}, _, _} ->
        {format, docstring}

      {:docs_v1, _, _, _, :hidden, _, _} ->
        :hidden

      _ ->
        nil
    end
  end

  @doc """
  Fetches information about the given module members if available.

  The given `members` are used to limit the result to the relevant
  entries. Arity may be given as `:any`, in which case all entries
  matching the name are returned.

  Functions with default arguments are normalized, such that each
  arity is treated as a separate member, sourcing documentation from
  the original one.

  ## Options

    * `:kinds` - a list of member kinds to limit the lookup to. Valid
      kinds are `:function`, `:macro` and `:type`. Defaults to all
      kinds

  """
  @spec lookup_module_members(
          module(),
          list({name :: atom(), arity :: non_neg_integer() | :any}),
          node(),
          keyword()
        ) :: list(member_info())
  def lookup_module_members(module, members, node, opts \\ [])

  def lookup_module_members(_module, _members, node, _opts) when node != node(), do: []

  def lookup_module_members(module, members, _node, opts) do
    members = MapSet.new(members)
    kinds = opts[:kinds] || [:function, :macro, :type]

    specs =
      with true <- :function in kinds or :macro in kinds,
           {:ok, specs} <- Code.Typespec.fetch_specs(module) do
        Map.new(specs)
      else
        _ -> %{}
      end

    type_specs =
      with true <- :type in kinds,
           {:ok, types} <- Code.Typespec.fetch_types(module) do
        for {type_kind, {name, _defs, vars}} = type <- types,
            type_kind in [:type, :opaque],
            into: Map.new(),
            do: {{name, Enum.count(vars)}, type}
      else
        _ -> %{}
      end

    case Elixir.Code.fetch_docs(module) do
      {:docs_v1, _, _, format, _, _, docs} ->
        for {{kind, name, base_arity}, _line, signatures, doc, meta} <- docs,
            kind in kinds,
            defaults = Map.get(meta, :defaults, 0),
            arity <- (base_arity - defaults)..base_arity,
            MapSet.member?(members, {name, arity}) or MapSet.member?(members, {name, :any}),
            do: %{
              kind: kind,
              name: name,
              arity: arity,
              from_default: arity != base_arity,
              documentation: documentation(doc, format),
              signatures: signatures,
              specs: Map.get(specs, {name, base_arity}, []),
              type_spec: Map.get(type_specs, {name, base_arity}, nil),
              meta: meta
            }

      _ ->
        []
    end
  end

  defp documentation(%{"en" => docstr}, format), do: {format, docstr}
  defp documentation(:hidden, _format), do: :hidden
  defp documentation(_doc, _format), do: nil

  @doc """
  Determines a more specific module type if any.
  """
  @spec get_module_subtype(module) ::
          :protocol | :implementation | :exception | :struct | :behaviour | nil
  def get_module_subtype(module) do
    cond do
      not ensure_loaded?(module) ->
        nil

      function_exported?(module, :__protocol__, 1) ->
        :protocol

      function_exported?(module, :__impl__, 1) ->
        :implementation

      function_exported?(module, :__struct__, 0) ->
        if function_exported?(module, :exception, 1) do
          :exception
        else
          :struct
        end

      function_exported?(module, :behaviour_info, 1) ->
        :behaviour

      true ->
        nil
    end
  end

  # In case insensitive file systems, attempting to load Elixir will
  # log a warning in the terminal as it wrongly loads elixir.beam,
  # so we explicitly list it.
  defp ensure_loaded?(Elixir), do: false
  defp ensure_loaded?(module), do: Code.ensure_loaded?(module)

  @doc """
  Extracts the location about an identifier found.

  The function returns the line where the identifier is located.
  """
  @spec locate_definition(list() | binary(), definition()) :: {:ok, pos_integer()} | :error
  def locate_definition(path, identifier)

  def locate_definition(path, {:module, module}) do
    with {:ok, {:raw_abstract_v1, annotations}} <- beam_lib_chunks(path, :abstract_code) do
      {:attribute, anno, :module, ^module} =
        Enum.find(annotations, &match?({:attribute, _, :module, _}, &1))

      {:ok, :erl_anno.line(anno)}
    end
  end

  def locate_definition(path, {:function, name, arity}) do
    with {:ok, {:debug_info_v1, _, {:elixir_v1, meta, _}}} <- beam_lib_chunks(path, :debug_info),
         {_pair, _kind, kw, _body} <- keyfind(meta.definitions, {name, arity}) do
      Keyword.fetch(kw, :line)
    end
  end

  def locate_definition(path, {:type, name, arity}) do
    with {:ok, {:raw_abstract_v1, annotations}} <- beam_lib_chunks(path, :abstract_code) do
      fetch_type_line(annotations, name, arity)
    end
  end

  defp fetch_type_line(annotations, name, arity) do
    for {:attribute, anno, :type, {^name, _, vars}} <- annotations, length(vars) == arity do
      :erl_anno.line(anno)
    end
    |> case do
      [] -> :error
      lines -> {:ok, Enum.min(lines)}
    end
  end

  defp beam_lib_chunks(path, key) do
    case :beam_lib.chunks(path, [key]) do
      {:ok, {_, [{^key, value}]}} -> {:ok, value}
      _ -> :error
    end
  end

  defp keyfind(list, key) do
    List.keyfind(list, key, 0) || :error
  end
end