corsa

defmodule Corsa.Spec do
  @moduledoc """

  SpecArgs contract, also known as type requirements, must be met before executing a function. In
  Corsa, they are checked during runtime and just prior to the annotated function's execution.

  ## Example


  iex> defmodule #{__MODULE__}.Example do
  ...>   use Corsa.Type
  ...>   use Corsa.Assert
  ...>   use Corsa.Spec
  ...>   spec f(arg1 :: integer(), arg2 :: boolean()), :ok
  ...>   def f(1, false), do: :error
  ...>   def f(_arg1, _arg2), do: :ok
  ...> end
  iex> #{__MODULE__}.Example.f(1, true)
  iex> #{__MODULE__}.Example.f(1, false)
  ** (Corsa.SpecResultViolationError) @spec does not hold in call '#{inspect(__MODULE__)}.Example.f(1, true)' with result ':ok'

  """

  import Corsa.Utils

  @doc false
  @spec __using__([]) :: Macro.t()
  defmacro __using__([]) do
    context = __CALLER__.module
    Module.register_attribute(context, :corsa_specs, accumulate: true)
    Module.register_attribute(context, :corsa_specs_args, accumulate: true)
    Module.register_attribute(context, :corsa_specs_result, accumulate: true)

    quote do
      require unquote(__MODULE__)
      @before_compile unquote(__MODULE__)
    end
  end

  @doc """

  ## Errors

  iex> defmodule #{__MODULE__}.ExampleError do
  ...>   use Corsa.Type
  ...>   use Corsa.Assert
  ...>   use Corsa.Spec
  ...>   spec f(arg1 :: integer(), arg1 :: boolean()), boolean()
  ...>   def f(_arg1, _arg2), do: :ok
  ...> end
  ** (Corsa.SpecError) arguments in @spec should contain different names

  """
  defmacro spec({name, _, args}, result_type) do
    context = __CALLER__.module
    stacktrace = Macro.Env.stacktrace(__CALLER__)

    args =
      Enum.reduce(args, {[], 1}, fn
        {:"::", _, [arg, type]}, {args, n} -> {[{arg, type} | args], n + 1}
        type, {args, n} -> {[{Macro.var(:"arg#{n}", context), type} | args], n + 1}
      end)
      |> elem(0)
      |> Enum.reverse()

    {args_names, _args_types} = Enum.unzip(args)

    unless args_names == Enum.uniq(args_names) do
      "arguments in @spec should contain different names"
      |> then(&reraise(Corsa.SpecError, &1, stacktrace))
    end

    arity = length(args_names)

    Module.put_attribute(context, :corsa_specs, {name, arity})
    Module.put_attribute(context, :corsa_specs_args, {{name, arity}, args})
    Module.put_attribute(context, :corsa_specs_result, {{name, arity}, result_type})

    []
  end

  @doc false
  @spec __before_compile__(Macro.Env.t()) :: Macro.t()
  defmacro __before_compile__(env) do
    module = env.module

    specs = Module.get_attribute(module, :corsa_specs) |> Enum.uniq()
    specs_args = Module.get_attribute(module, :corsa_specs_args)
    specs_result = Module.get_attribute(module, :corsa_specs_result)

    types =
      Enum.zip_with(specs_args, specs_result, fn {f, args}, {f, result} -> {f, args, result} end)
      |> Enum.group_by(fn {k, _, _} -> k end, fn {_, a, r} -> {a, r} end)

    defs =
      Module.definitions_in(module, :def) |> Enum.map(fn {name, arity} -> {:def, name, arity} end)

    defps =
      Module.definitions_in(module, :defp) |> Enum.map(fn {name, arity} -> {:defp, name, arity} end)

    defs = defs ++ defps

    for {name, arity} <- specs,
        f = Enum.find(defs, &match?({_, ^name, ^arity}, &1)),
        types = Map.get(types, {name, arity}),
        {_, _, location, _} = Module.get_definition(module, {name, arity}) do
      file = __CALLER__.file
      Module.make_overridable(module, [{name, arity}])

      [
        spec_args(module, name, types),
        spec_result(module, name, types),
        new_def(f, module, location ++ [file: file])
      ]
    end
  end

  defp spec_args(context, name, types) do
    args_names = hd(types) |> elem(0) |> Enum.unzip() |> elem(0)

    clause_check = fn args ->
      Enum.reduce(args, true, fn
        {n, t}, true -> quote do: conforms?(unquote(n), unquote(t))
        {n, t}, acc -> quote do: conforms?(unquote(n), unquote(t)) && unquote(acc)
      end)
    end

    args_check =
      Enum.reduce(types, :fst, fn
        {args, _}, :fst -> quote do: unquote(clause_check.(args))
        {args, _}, acc -> quote do: unquote(clause_check.(args)) || unquote(acc)
      end)

    quote context: context do
      Kernel.defp unquote(:"#{name}_spec_args")(unquote_splicing(args_names)) do
        unquote(args_check)
      end
    end
  end

  defp spec_result(context, name, types) do
    args_names = hd(types) |> elem(0) |> Enum.unzip() |> elem(0)
    result = Macro.unique_var(:result, context)

    clause_check = fn args, result_type ->
      args_check =
        Enum.reduce(args, true, fn
          {n, t}, true -> quote do: conforms?(unquote(n), unquote(t))
          {n, t}, acc -> quote do: conforms?(unquote(n), unquote(t)) && unquote(acc)
        end)

      quote do: unquote(args_check) && conforms?(unquote(result), unquote(result_type))
    end

    args_check =
      Enum.reduce(types, :fst, fn
        {args, r_t}, :fst -> quote do: unquote(clause_check.(args, r_t))
        {args, r_t}, acc -> quote do: unquote(clause_check.(args, r_t)) || unquote(acc)
      end)

    quote context: context do
      Kernel.defp unquote(:"#{name}_spec_result")(unquote_splicing(args_names ++ [result])) do
        unquote(args_check)
      end
    end
  end

  defp new_def({def_t, name, arity}, context, location) do
    line = Keyword.get(location, :line)
    args = Macro.generate_arguments(arity, nil)

    before =
      quote context: context, line: line do
              Corsa.Assert.assert(
          unquote(:"#{name}_spec_args")(unquote_splicing(args)),
          Corsa.SpecArgViolationError,
          Corsa.SpecError,
          call: call
        )
      end

    result =
      quote context: context, line: line do
        Corsa.Assert.assert(
          unquote(:"#{name}_spec_result")(unquote_splicing(args), result),
          Corsa.SpecResultViolationError,
          Corsa.SpecError,
          call: call,
          result: result
        )
      end

    new_def(def_t, name, args, before, [], result, context, line)
  end
end