lib/archeometer/util/code.ex

defmodule Archeometer.Util.Code do
  @moduledoc """
  This module provides several helpers functions to deal with Elixir ASTs.

  The functions have specific uses in other modules, but are general enough to
  be useful in other contexts.
  """

  @doc """
  Walks the given AST and stores all the ocurrences of the given atom as an
  operator.
  """
  def collect_nodes(ast, atoms)

  def collect_nodes(ast, atom) when is_atom(atom) do
    Macro.prewalk(ast, [], fn ast, acc ->
      case ast do
        {^atom, _, _} ->
          {ast, [ast | acc]}

        _ ->
          {ast, acc}
      end
    end)
    |> elem(1)
  end

  def collect_nodes(ast, atoms) when is_list(atoms) do
    Enum.flat_map(atoms, &collect_nodes(ast, &1))
  end

  @doc """
  Walks the AST and finds all the definitions matching the given atoms, except
  if they are inside a macro definition, as quoted definitions are not yet
  instantiated.
  """
  def collect_defs(full_ast, defatom) do
    collect_nodes(full_ast, defatom)
    |> Enum.filter(fn ast ->
      case {ast, resolve_scope(full_ast, ast)} do
        {{:defmacro, _, _}, {:defmacro, _}} -> true
        {_, {:defmacro, _}} -> false
        _ -> true
      end
    end)
  end

  @doc """
  Get the metadata of the given AST.
  """
  def get_meta({_, meta, _}, atom), do: Keyword.get(meta, atom, nil)

  def get_meta(_, _), do: nil

  @doc """
  Get declaration information from the AST of a `def`-like macro. That is
  usually the name and arguments, but it dependes on the exact construct.

  Current ones are `def, defp, defmacro, defmacrop, defmodule, defstruct`.
  """
  def get_decl({macro, _, [{:when, _, [decl | _]} | _]})
      when macro in [:def, :defp, :defmacro, :defmacrop] do
    {:ok, decl}
  end

  def get_decl({macro, _, [decl | _]})
      when macro in [:def, :defp, :defmacro, :defmacrop] do
    {:ok, decl}
  end

  def get_decl({:defmodule, _, [decl | _]}), do: {:ok, decl}

  def get_decl({:defstruct, _decl, _fields} = struct), do: {:ok, struct}

  def get_decl(_ast), do: {:error, :no_decl}

  @doc """
  Get the maximum line number present in the metadata of an AST.
  """
  def num_lines(ast) do
    {min, max} =
      ast
      |> Macro.prewalk([], fn ast, acc ->
        case get_meta(ast, :line) do
          nil -> {ast, acc}
          line -> {ast, [line | acc]}
        end
      end)
      |> elem(1)
      |> Enum.min_max()

    max - min + 1
  end

  @doc """
  Determine the scope of the `ast` inside the `full_ast`.
  """
  def resolve_scope(full_ast, ast) do
    decl_line =
      ast
      |> get_decl()
      |> case do
        {:ok, elem} -> elem
        {:error, _} -> ast
      end
      |> get_meta(:line)

    full_ast
    |> Credo.Code.Scope.name(line: decl_line)
  end

  @doc """
  Determine the name of the module of some subset of the AST.
  """
  def resolve_mod_name(full_ast, ast) do
    resolve_scope(full_ast, ast)
    |> elem(1)
    |> Credo.Code.Scope.mod_name()
  end

  @doc """
  Create a new atom by concatenaing two other existing ones.
  """
  def atom_concat(atom0, atom1) do
    atom0
    |> Atom.to_string()
    |> Kernel.<>(Atom.to_string(atom1))
    |> String.to_atom()
  end

  @doc """
  Given a module name, it will return the underscored version of the last
  part of the module name.
  """
  def snakefy({_alias, _meta, [module]}), do: snakefy(module)

  def snakefy(module) do
    module
    |> Atom.to_string()
    |> String.split(".")
    |> List.last()
    |> Macro.underscore()
    |> String.to_atom()
  end
end