defmodule Recode.Context do
@moduledoc ~S'''
This module provides functions to traverse an AST with a `%Context{}`.
## Examples
The following example shows the `%Context{}` for the definition of
`MyApp.Bar.bar/1`.
iex> alias Rewrite.Source
...> alias Recode.Context
...> """
...> defmodule MyApp.Foo do
...> def foo, do: :foo
...> end
...>
...> defmodule MyApp.Bar do
...> alias MyApp.Foo
...>
...> def bar(x) do
...> {Foo.foo(), x}
...> end
...> end
...> """
...> |> Source.Ex.from_string()
...> |> Source.get(:quoted)
...> |> Zipper.zip()
...> |> Context.traverse(nil, fn
...> zipper, context, nil ->
...> case context.definition do
...> {{:def, :bar, 1}, _meta} -> {zipper, context, context}
...> _def -> {zipper, context, nil}
...> end
...> zipper, context, acc ->
...> {zipper, context, acc}
...> end)
...> |> elem(1) |> Map.put(:node, nil)
%Context{
node: nil,
aliases: [
{MyApp.Foo,
[
trailing_comments: [],
leading_comments: [],
end_of_expression: [newlines: 2, line: 6, column: 18],
line: 6,
column: 3
], nil}
],
assigns: %{},
definition:
{{:def, :bar, 1},
[
trailing_comments: [],
leading_comments: [],
do: [line: 8, column: 14],
end: [line: 10, column: 3],
line: 8,
column: 3
]},
imports: [],
module:
{MyApp.Bar,
[
trailing_comments: [],
leading_comments: [],
do: [line: 5, column: 21],
end: [line: 11, column: 1],
line: 5,
column: 1
]},
requirements: [],
usages: []
}
'''
alias Recode.Context
alias Sourceror.Zipper
defstruct module: nil,
aliases: [],
imports: [],
usages: [],
requirements: [],
definition: nil,
assigns: %{},
moduledoc: nil,
node: nil,
doc: nil,
spec: nil,
impl: nil
@type t :: %Context{
module: term() | nil,
aliases: list(),
imports: list(),
usages: list(),
requirements: list(),
definition: term(),
assigns: map(),
moduledoc: Macro.t() | nil,
node: term() | nil,
doc: {term() | nil, Macro.t()} | nil,
spec: {term() | nil, Macro.t()} | nil,
impl: {term() | nil, Macro.t()} | nil
}
@type zipper :: Zipper.t()
@doc """
Returns the current module of a context.
"""
@spec module(t()) :: module() | nil
def module(%Context{module: nil}), do: nil
def module(%Context{module: {name, _meta}}), do: name
@doc """
Expands the module alias for the given `mfa`.
"""
@spec expand_mfa(t(), mfa()) :: {:ok, mfa()} | :error
def expand_mfa(%Context{aliases: aliases}, {module, fun, arity}) do
with {:ok, alias} <- find_alias(aliases, module) do
{:ok, {alias, fun, arity}}
end
end
@doc """
Assigns the given `value` under `key` to the `context`.
"""
@spec assign(t(), atom(), term()) :: t()
def assign(%Context{assigns: assigns} = context, key, value) when is_atom(key) do
%{context | assigns: Map.put(assigns, key, value)}
end
@doc """
Returns `true` if a `spec` is available.
"""
@spec spec?(t()) :: boolean
def spec?(%Context{spec: spec}), do: not is_nil(spec)
@doc """
Returns `true` if an `impl` is available.
"""
@spec impl?(t()) :: boolean
def impl?(%Context{impl: impl}), do: not is_nil(impl)
@doc """
Returns true if `definition` satisfies the assumption.
"""
@spec definition?(t(), :public | :visible) :: boolean()
def definition?(%Context{definition: definition}, :public)
when not is_nil(definition) do
{{kind, _name, _arity}, _meta} = definition
kind in [:def, :defmacro]
end
def definition?(%Context{definition: definition} = context, :visible)
when not is_nil(definition) do
with true <- definition?(context, :public) do
not (doc?(context, false) or moduledoc?(context, false))
end
end
@doc """
Returns `true` if `@doc` is available and not set to `@doc false`.
"""
@spec doc?(t()) :: boolean()
def doc?(%Context{doc: nil}), do: false
def doc?(%Context{doc: {_to, doc}}) do
case attribute_args(doc) do
[{_name, _meta, [false]}] -> false
_else -> true
end
end
@doc """
Returns `true` if `@doc false`.
"""
@spec doc?(t(), false) :: boolean()
def doc?(%Context{doc: nil}, false), do: false
def doc?(%Context{doc: {_to, doc}}, false) do
case attribute_args(doc) do
[{_name, _meta, [false]}] -> true
_else -> false
end
end
@doc """
Returns `true` if `@moduledoc` is available and not set to `@moduledoc false`.
"""
@spec moduledoc?(t()) :: boolean()
def moduledoc?(%Context{moduledoc: moduledoc}) do
case attribute_args(moduledoc) do
nil -> false
[{_name, _meta, [false]}] -> false
_else -> true
end
end
@doc """
Returns `true` if `@moduledoc false`.
"""
@spec moduledoc?(t(), false) :: boolean()
def moduledoc?(%Context{moduledoc: nil}, false), do: false
def moduledoc?(%Context{moduledoc: moduledoc}, false) do
case attribute_args(moduledoc) do
[{_name, _meta, [false]}] -> true
_else -> false
end
end
@doc """
Merges the given `map` to the assigns of the `context`.
"""
@spec assigns(t(), map()) :: t()
def assigns(%Context{assigns: assigns} = context, map) when is_map(map) do
%{context | assigns: Map.merge(assigns, map)}
end
@doc """
Traverses the given `zipper` and applies `fun` on each node.
The `fun` gets the current `zipper` and `context` as arguments.
"""
@spec traverse(zipper(), fun) :: zipper()
when fun: (zipper(), t() -> {zipper(), t()})
def traverse(zipper, fun) when is_function(fun, 2) do
zipper
|> run_traverse(%Context{}, fun)
|> elem(0)
end
@doc """
Traverses the given `zipper` with an `acc` and applies `fun` on each node.
"""
@spec traverse(zipper(), acc, fun) :: {zipper(), acc}
when acc: term(),
fun: (zipper(), t(), acc -> {zipper(), t(), acc})
def traverse(zipper, acc, fun) when is_function(fun, 3) do
{zipper, {_context, acc}} = run_traverse(zipper, %Context{}, acc, fun)
{zipper, acc}
end
# helpers for traverse/2
defp run_traverse(zipper, context, fun) do
Zipper.traverse_while(zipper, context, fn zipper, context ->
do_traverse(zipper, context, fun)
end)
end
defp do_traverse(%Zipper{node: {:@, _meta, _args} = attribute} = zipper, context, fun) do
context = add_attribute(context, attribute)
cont(zipper, context, fun)
end
defp do_traverse(%Zipper{node: {:alias, meta, args}} = zipper, context, fun)
when not is_nil(args) do
aliases = get_aliases(args, meta, context)
context = add_aliases(context, aliases)
cont(zipper, context, fun)
end
defp do_traverse(%Zipper{node: {:import, meta, [arg, opts]}} = zipper, context, fun) do
import = get_alias(arg, context)
context = add_import(context, {import, meta, opts})
cont(zipper, context, fun)
end
defp do_traverse(%Zipper{node: {:import, meta, args}} = zipper, context, fun)
when not is_nil(args) do
imports = get_aliases(args, meta, context)
context = add_imports(context, imports)
cont(zipper, context, fun)
end
defp do_traverse(%Zipper{node: {:use, meta, [arg, opts]}} = zipper, context, fun) do
use = get_alias(arg, context)
context = add_use(context, {use, meta, opts})
cont(zipper, context, fun)
end
defp do_traverse(%Zipper{node: {:use, meta, [arg]}} = zipper, context, fun) do
use = get_alias(arg, context)
context = add_use(context, {use, meta, nil})
cont(zipper, context, fun)
end
defp do_traverse(%Zipper{node: {:require, meta, [arg, opts]}} = zipper, context, fun) do
require = get_alias(arg, context)
context = add_require(context, {require, meta, opts})
cont(zipper, context, fun)
end
defp do_traverse(%Zipper{node: {:require, meta, [arg]}} = zipper, context, fun) do
require = get_alias(arg, context)
context = add_require(context, {require, meta, nil})
cont(zipper, context, fun)
end
defp do_traverse(
%Zipper{node: {:defmodule, meta, [{:__aliases__, _, name} | _]}} = zipper,
context,
fun
) do
case module(context, :meta) == meta do
true ->
cont(zipper, context, fun)
false ->
module = module_concat(context, Module.concat(name))
do_traverse_sub(zipper, context, fun, module: {module, meta})
end
end
defp do_traverse(%Zipper{node: {:defimpl, _meta, _args}} = zipper, context, _fun) do
{:skip, zipper, context}
end
defp do_traverse(%Zipper{node: {definition, meta, args}} = zipper, context, fun)
when definition in [:def, :defp, :defmacro, :defmacrop] and not is_nil(args) do
case definition(context, :meta) == meta do
true ->
cont(zipper, context, fun)
false ->
definition = get_definition(definition, args)
context = update_attributes(context, definition)
do_traverse_sub(zipper, context, fun, definition: {definition, meta})
end
end
defp do_traverse(zipper, context, fun) do
cont(zipper, context, fun)
end
defp do_traverse_sub(zipper, context, fun, [{key, value}]) do
sub_context = Map.put(context, key, value)
{%Zipper{node: ast}, %Context{assigns: assigns}} = run_traverse(zipper, sub_context, fun)
zipper = Zipper.replace(zipper, ast)
context = Context.assigns(context, assigns)
{:skip, zipper, context}
end
# helpers for traverse/3
defp run_traverse(zipper, context, acc, fun) do
Zipper.traverse_while(zipper, {context, acc}, fn zipper, {context, acc} ->
do_traverse(zipper, context, acc, fun)
end)
end
defp do_traverse(%Zipper{node: {:@, _meta, _args} = attribute} = zipper, context, acc, fun) do
context = add_attribute(context, attribute)
cont(zipper, context, acc, fun)
end
defp do_traverse(%Zipper{node: {:alias, meta, args}} = zipper, context, acc, fun)
when not is_nil(args) do
aliases = get_aliases(args, meta, context)
context = add_aliases(context, aliases)
cont(zipper, context, acc, fun)
end
defp do_traverse(%Zipper{node: {:import, meta, [arg, opts]}} = zipper, context, acc, fun) do
import = get_alias(arg, context)
context = add_import(context, {import, meta, opts})
cont(zipper, context, acc, fun)
end
defp do_traverse(%Zipper{node: {:import, meta, args}} = zipper, context, acc, fun)
when not is_nil(args) do
imports = get_aliases(args, meta, context)
context = add_imports(context, imports)
cont(zipper, context, acc, fun)
end
defp do_traverse(%Zipper{node: {:use, meta, [arg, opts]}} = zipper, context, acc, fun) do
use = get_alias(arg, context)
context = add_use(context, {use, meta, opts})
cont(zipper, context, acc, fun)
end
defp do_traverse(%Zipper{node: {:use, meta, [arg]}} = zipper, context, acc, fun) do
use = get_alias(arg, context)
context = add_use(context, {use, meta, nil})
cont(zipper, context, acc, fun)
end
defp do_traverse(%Zipper{node: {:require, meta, [arg, opts]}} = zipper, context, acc, fun) do
require = get_alias(arg, context)
context = add_require(context, {require, meta, opts})
cont(zipper, context, acc, fun)
end
defp do_traverse(%Zipper{node: {:require, meta, [arg]}} = zipper, context, acc, fun) do
require = get_alias(arg, context)
context = add_require(context, {require, meta, nil})
cont(zipper, context, acc, fun)
end
defp do_traverse(
%Zipper{node: {:defmodule, meta, [{:__aliases__, _, name} | _]}} = zipper,
context,
acc,
fun
) do
case module(context, :meta) == meta do
true ->
cont(zipper, context, acc, fun)
false ->
module = module_concat(context, Module.concat(name))
do_traverse_sub(zipper, context, acc, fun, module: {module, meta})
end
end
defp do_traverse(%Zipper{node: {:defimpl, _meta, _args}} = zipper, context, acc, _fun) do
{:skip, zipper, {context, acc}}
end
defp do_traverse(%Zipper{node: {definition, meta, args}} = zipper, context, acc, fun)
when definition in [:def, :defp, :defmacro, :defmacrop] and length(args) == 2 do
case definition(context, :meta) == meta do
true ->
cont(zipper, context, acc, fun)
false ->
definition = get_definition(definition, args)
context = update_attributes(context, definition)
do_traverse_sub(zipper, context, acc, fun, definition: {definition, meta})
end
end
defp do_traverse(zipper, context, acc, fun) do
cont(zipper, context, acc, fun)
end
defp do_traverse_sub(zipper, context, acc, fun, [{key, value}]) do
sub_context = Map.put(context, key, value)
{%Zipper{node: ast}, {%Context{assigns: assigns}, acc}} =
run_traverse(zipper, sub_context, acc, fun)
zipper = Zipper.replace(zipper, ast)
context = Context.assigns(context, assigns)
{:skip, zipper, {context, acc}}
end
# helpers for traverse/2/3
defp add_attribute(context, {:@, _meta, [arg]} = attribute) do
case arg do
{:moduledoc, _meta, _args} ->
%{context | moduledoc: attribute}
{:doc, _meta, _args} ->
%{context | doc: {nil, attribute}}
{:spec, _meta, _args} ->
%{context | spec: {nil, attribute}}
{:impl, _meta, _args} ->
%{context | impl: {nil, attribute}}
_args ->
context
end
end
defp update_attributes(context, definition) do
context
|> update_attribute(definition, :doc)
|> update_attribute(definition, :spec)
|> update_attribute(definition, :impl)
end
defp update_attribute(context, definition, key) do
Map.update!(context, key, fn
{nil, value} -> {definition, value}
{^definition, value} -> {definition, value}
_spec -> nil
end)
end
# other helpers
defp get_definition(kind, [{:when, _meta1, [{name, _meta2, args}, _expr]}, _block]) do
{kind, name, length(args)}
end
defp get_definition(kind, [{name, _meta, args}]) do
{kind, name, length(args)}
end
defp get_definition(kind, [{name, _meta, nil}, _expr]) do
{kind, name, 0}
end
defp get_definition(kind, [{name, _meta, args}, _expr]) do
{kind, name, length(args)}
end
defp get_alias({:__aliases__, _meta, name}, _context) when is_list(name) do
Module.concat(name)
rescue
_error -> :unknown
end
defp get_alias({:__block__, _meta, [name]}, _context) when is_atom(name) do
name
end
defp get_alias({:unquote, _meta, _args} = expr, _context) do
expr
end
defp get_alias({:alias, _meta, [args]}, %Context{} = context) do
get_alias(args, context)
end
defp get_alias({:alias, meta, args}, %Context{} = context) do
get_aliases(args, meta, context)
end
defp get_alias({:require, _meta, [[arg]]}, %Context{} = context) do
get_alias(arg, context)
end
defp get_alias({:require, _meta, [arg]}, %Context{} = context) do
get_alias(arg, context)
end
defp get_alias({:__MODULE__, _meta1, _args}, %Context{} = context) do
module(context)
end
defp get_aliases([{{:., _, [alias, :{}]}, _, aliases}], meta, %Context{} = context) do
base = get_alias(alias, context)
Enum.map(aliases, fn alias ->
{Module.concat(base, get_alias(alias, context)), meta, nil}
end)
end
defp get_aliases([arg], meta, context), do: [{get_alias(arg, context), meta, nil}]
defp get_aliases([arg, opts], meta, context), do: [{get_alias(arg, context), meta, opts}]
defp cont(%Zipper{node: node} = zipper, context, fun) do
{zipper, context} = fun.(zipper, %Context{context | node: node})
{:cont, zipper, context}
end
defp cont(%Zipper{node: node} = zipper, context, acc, fun) do
{zipper, context, acc} = fun.(zipper, %Context{context | node: node}, acc)
{:cont, zipper, {context, acc}}
end
defp find_alias(aliases, module) do
Enum.find_value(aliases, :error, fn
{alias, _meta, nil} ->
case ends_with?(alias, module) do
false -> false
true -> {:ok, alias}
end
{alias, _meta, opts} ->
case get_as(opts) == module do
false -> false
true -> {:ok, alias}
end
end)
end
defp get_as([{{:__block__, _meta_block, [:as]}, {:__aliases__, _meta_aliases, as}}]) do
Module.concat(as)
end
defp get_as(_ast), do: :error
defp definition(%Context{definition: nil}, :meta), do: nil
defp definition(%Context{definition: {_definition, meta}}, :meta), do: meta
defp add_aliases(%Context{aliases: aliases} = context, list) do
%Context{context | aliases: aliases ++ list}
end
defp add_import(%Context{imports: imports} = context, import) do
%Context{context | imports: imports ++ [import]}
end
defp add_imports(%Context{imports: imports} = context, list) do
%Context{context | imports: imports ++ list}
end
defp add_use(%Context{usages: usages} = context, use) do
%Context{context | usages: usages ++ [use]}
end
defp add_require(%Context{requirements: requirements} = context, require) do
%Context{context | requirements: requirements ++ [require]}
end
defp module(%Context{module: nil}, :meta), do: nil
defp module(%Context{module: {_name, meta}}, :meta), do: meta
defp module_concat(%Context{module: nil}, module), do: module
defp module_concat(%Context{module: {module1, _meta}}, module2) do
Module.concat(module1, module2)
end
defp attribute_args(nil), do: nil
defp attribute_args({:@, _meta1, [{_name, _meta2, args}]}), do: args
defp ends_with?(value, suffix) when is_atom(value) and is_atom(suffix) do
list = Module.split(value)
suffix = Module.split(suffix)
ends_with?(list, suffix)
end
defp ends_with?(list, suffix) when is_list(list) and is_list(suffix) do
case length(list) - length(suffix) do
diff when diff < 0 -> false
diff -> Enum.drop(list, diff) == suffix
end
end
end
