corsa

# TODO: custom error messages in contracts
# TODO: write docs
# TODO: auto-generate contract docs
defmodule Corsa do
  @moduledoc """
  Corsa is a library to write runnable code contracts in Elixir.
  """

  @doc false
  @spec __using__(logger: Logger.level()) :: Macro.t()
  defmacro __using__(opts) do
    context = __CALLER__.module
    Module.register_attribute(context, :pre, accumulate: true)
    Module.register_attribute(context, :post, accumulate: true)
    Module.register_attribute(context, :throws, accumulate: true)
    Module.register_attribute(context, :within, accumulate: true)
    Module.register_attribute(context, :decreases, accumulate: true)
    Module.register_attribute(context, :specs, accumulate: true)
    Module.register_attribute(context, :options, accumulate: true)

    Keyword.get(opts, :logger, Application.get_env(:corsa, :logger))
    |> then(&Module.put_attribute(context, :options, {:logger, &1}))

    quote context: context do
      require Corsa
      import Corsa, except: [assert: 1, assert: 2]
      @before_compile unquote(__MODULE__)
      import Kernel, except: [@: 1]
      use TypeCheck
      import TypeCheck.Macros, except: [@: 1]
      import TypeCheck, only: [conforms?: 2]
      require Logger
    end
  end

  defmacro @{:assert, _line, expr} do
    quote context: __CALLER__.module do
      Corsa.assert(unquote_splicing(expr))
    end
  end

  defmacro @{:pre, _line, [{:when, _context, _args}, [do: _body]]} do
    reraise(Corsa.PreError, "guards are not supported in @pre", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:pre, _line, [{name, _context, args}, [do: body]]} do
    arity = length(args)
    context = __CALLER__.module
    stacktrace = Macro.Env.stacktrace(__CALLER__)

    if {name, arity} in Module.get_attribute(context, :pre) do
      "@pre for function '#{name}/#{arity}' already defined"
      |> then(&reraise(Corsa.PreError, &1, stacktrace))
    end

    for {arg, _, _} <- args, arg = Atom.to_string(arg), match?("_" <> _, arg) do
      "arguments in @pre cannot be ignored with _ "
      |> then(&reraise(Corsa.PreError, &1, stacktrace))
    end

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

    Module.put_attribute(context, :pre, {name, arity})
    opts = Macro.unique_var(:opts, context)

    quote context: context do
      defp unquote(:"#{name}_pre")(unquote_splicing(args)) do
        unquote(opts) =
          [unquote_splicing(args)]
          |> to_corsa_call(__MODULE__, unquote(name))
          |> then(&[call: &1])

        Corsa.assert(unquote(body), Corsa.PreViolationError, Corsa.PreError, unquote(opts))
      end
    end
  end

  defmacro @{:pre, _line, _expr} do
    reraise(Corsa.PreError, "syntax error in @pre", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:post, _line, [{:when, _context, _args}, [do: _body]]} do
    reraise(Corsa.PostError, "guards are not supported in @post", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:post, _line, [{name, _context, args}, [do: body]]} do
    arity = length(args)
    context = __CALLER__.module
    result = Macro.var(:result, nil)
    stacktrace = Macro.Env.stacktrace(__CALLER__)

    if {name, arity} in Module.get_attribute(context, :post) do
      "@post for function '#{name}/#{arity}' already defined"
      |> then(&reraise(Corsa.PostError, &1, stacktrace))
    end

    for {arg, _, _} <- args, arg = Atom.to_string(arg), match?("_" <> _, arg) do
      "arguments in @post cannot be ignored with _ "
      |> then(&reraise(Corsa.PostError, &1, stacktrace))
    end

    for {arg, _, _} <- args, arg == :result do
      "result cannot be an argument in @post"
      |> then(&reraise(Corsa.PostError, &1, stacktrace))
    end

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

    Module.put_attribute(context, :post, {name, arity})
    opts = Macro.unique_var(:opts, context)

    quote context: context do
      defp unquote(:"#{name}_post")(unquote_splicing(args ++ [result])) do
        unquote(opts) =
          [unquote_splicing(args)]
          |> to_corsa_call(__MODULE__, unquote(name))
          |> then(&[call: &1, result: unquote(result)])

        Corsa.assert(unquote(body), Corsa.PostViolationError, Corsa.PostError, unquote(opts))
      end
    end
  end

  defmacro @{:post, _line, _expr} do
    reraise(Corsa.PostError, "syntax error in @post", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:throws, _line, [{:when, _context, _args}, [do: _body]]} do
    reraise(
      Corsa.ThrowsError,
      "guards are not supported in @throws",
      Macro.Env.stacktrace(__CALLER__)
    )
  end

  defmacro @{:throws, _line, [{name, _context, args}, [do: body]]} do
    arity = length(args)
    context = __CALLER__.module
    stacktrace = Macro.Env.stacktrace(__CALLER__)

    if {name, arity} in Module.get_attribute(context, :throws) do
      "@throws for function '#{name}/#{arity}' already defined"
      |> then(&reraise(Corsa.ThrowsError, &1, stacktrace))
    end

    for {arg, _, _} <- args, arg = Atom.to_string(arg), match?("_" <> _, arg) do
      "arguments in @throws cannot be ignored with _ "
      |> then(&reraise(Corsa.ThrowsError, &1, stacktrace))
    end

    for {arg, _, _} <- args, arg == :kind or arg == :value do
      "'#{arg}' cannot be an argument in @throws"
      |> then(&reraise(Corsa.ThrowsError, &1, stacktrace))
    end

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

    Module.put_attribute(context, :throws, {name, arity})
    opts = Macro.unique_var(:opts, context)

    quote context: context do
      defp unquote(:"#{name}_throws")(unquote_splicing(args), kind, value) do
        unquote(opts) =
          [unquote_splicing(args)]
          |> to_corsa_call(__MODULE__, unquote(name))
          |> then(&[call: &1, kind: kind, reason: value])

        Corsa.assert(unquote(body), Corsa.ThrowsViolationError, Corsa.ThrowsError, unquote(opts))
      end
    end
  end

  defmacro @{:throws, _line, _expr} do
    reraise(Corsa.ThrowsError, "syntax error in @throws", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:within, _line, [{:when, _context, _args}, [do: _body]]} do
    reraise(
      Corsa.WithinError,
      "guards are not supported in @within",
      Macro.Env.stacktrace(__CALLER__)
    )
  end

  defmacro @{:within, _line, [{name, _context, args}, [do: body]]} do
    arity = length(args)
    context = __CALLER__.module
    stacktrace = Macro.Env.stacktrace(__CALLER__)

    if {name, arity} in Module.get_attribute(context, :within) do
      "@within for function '#{name}/#{arity}' already defined"
      |> then(&reraise(Corsa.WithinError, &1, stacktrace))
    end

    for {arg, _, _} <- args, arg = Atom.to_string(arg), match?("_" <> _, arg) do
      "arguments in @within cannot be ignored with _ "
      |> then(&reraise(Corsa.WithinError, &1, stacktrace))
    end

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

    Module.put_attribute(context, :within, {name, arity})

    quote context: context do
      defp unquote(:"#{name}_within")(unquote_splicing(args)) do
        [unquote_splicing(args)]
        |> to_corsa_call(__MODULE__, unquote(name))
        |> then(
          &try do
            unquote(body)
          catch
            k, v ->
              reraise(Corsa.WithinError, %{call: &1, kind: k, reason: v}, __STACKTRACE__)
          end
        )
      end
    end
  end

  defmacro @{:within, _line, _expr} do
    reraise(Corsa.WithinError, "syntax error in @within", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:decreases, _line, [{:when, _context, _args}, [do: _body]]} do
    reraise(
      Corsa.DecreasesError,
      "guards are not supported in @decreases",
      Macro.Env.stacktrace(__CALLER__)
    )
  end

  defmacro @{:decreases, _line, [{name, _context, args}, [do: body]]} do
    arity = length(args)
    context = __CALLER__.module
    stacktrace = Macro.Env.stacktrace(__CALLER__)

    if {name, arity} in Module.get_attribute(context, :decreases) do
      "@decreases for function '#{name}/#{arity}' already defined"
      |> then(&reraise(Corsa.DecreasesError, &1, stacktrace))
    end

    for {arg, _, _} <- args, arg = Atom.to_string(arg), match?("_" <> _, arg) do
      "arguments in @decreases cannot be ignored with _ "
      |> then(&reraise(Corsa.DecreasesError, &1, stacktrace))
    end

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

    Module.put_attribute(context, :decreases, {name, arity})

    old_args = Enum.map(args, fn {name, env, ctx} -> {:"#{name}_old", env, ctx} end)
    id = Macro.unique_var(:id, context)
    opts = Macro.unique_var(:opts, context)

    quote context: context do
      defp unquote(:"#{name}_decreases_measure")(unquote_splicing(args)) do
        unquote(body)
      end

      defp unquote(:"#{name}_decreases")(unquote_splicing(args)) do
        unquote(id) = {__MODULE__, unquote(name), unquote(arity)}

        unless Corsa.Stack.empty?(unquote(id)) do
          [unquote_splicing(old_args)] = Corsa.Stack.peek(unquote(id))

          unquote(opts) = [
            old_call:
              [unquote_splicing(old_args)]
              |> Corsa.to_corsa_call(__MODULE__, unquote(name)),
            call:
              [unquote_splicing(args)]
              |> Corsa.to_corsa_call(__MODULE__, unquote(name))
          ]

          Corsa.assert(
            unquote(:"#{name}_decreases_measure")(unquote_splicing(args)) <
              unquote(:"#{name}_decreases_measure")(unquote_splicing(old_args)),
            Corsa.DecreasesViolationError,
            Corsa.DecreasesError,
            unquote(opts)
          )
        end

        Corsa.Stack.push(unquote(id), [unquote_splicing(args)])
      end
    end
  end

  defmacro @{:decreases, _line, _expr} do
    reraise(Corsa.DecreasesError, "syntax error in @decreases", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:expects, _line, expr} when is_list(expr) do
    context = __CALLER__.module
    {predicates, clauses} = parse_expects(expr, __CALLER__)
    global_vars = __CALLER__.versioned_vars |> Map.keys() |> Enum.map(&elem(&1, 0)) |> MapSet.new()

    clauses =
      for {:->, env, [pattern, block]} <- clauses do
        {_, vars} = free_vars_names(pattern)
        overlapped = MapSet.intersection(vars, global_vars)

        if overlapped != MapSet.new() do
          env = Macro.Env.stacktrace(__CALLER__)

          if Enum.count(overlapped) == 1 do
            var = overlapped |> Enum.at(0) |> Atom.to_string()
            IO.warn("Variable \"#{var}\" is overriden inside pattern in @expects contract", env)
          else
            vars = overlapped |> Enum.map_join(", ", fn v -> "\"#{Atom.to_string(v)}\"" end)
            IO.warn("Variables #{vars} are overriden inside pattern in @expects contract", env)
          end
        end

        checks =
          Enum.filter(predicates, fn {_, pred_vars} ->
            MapSet.subset?(pred_vars, MapSet.union(global_vars, vars))
          end)
          |> Enum.map(fn {pred, _} ->
            quote context: context do
              Corsa.assert(unquote(pred), Corsa.ExpectsViolationError, Corsa.ExpectsError)
            end
          end)

        block =
          quote context: context do
            unquote_splicing(checks)
            unquote(block)
          end

        {:->, env, [pattern, block]}
      end

    quote context: context do
      receive do
        unquote(clauses)
      end
    end
  end

  defmacro @{:expects, _line, _expr} do
    reraise(Corsa.ExpectsError, "syntax error in @expects", Macro.Env.stacktrace(__CALLER__))
  end

  defmacro @{:spec, _line, expr = [{:"::", _, [{name, _, args}, result_type]}]} do
    context = __CALLER__.module
    arity = length(args)
    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 == Enum.uniq(args) do
      "arguments in @spec should contain different names"
      |> then(&reraise(Corsa.SpecError, &1, stacktrace))
    end

    args_check =
      Enum.map(args, fn {arg, type} ->
        quote do: TypeCheck.conforms?(unquote(arg), unquote(type))
      end)
      |> Enum.reduce(true, fn expr, acc -> quote do: unquote(acc) and unquote(expr) end)

    result = Macro.unique_var(:result, context)

    if {name, arity} in Module.get_attribute(context, :specs) do
      overridable = [{:"#{name}_args_spec", arity}, {:"#{name}_result_spec", arity + 1}]

      quote context: context do
        @spec unquote_splicing(expr)

        defoverridable(unquote(overridable))
        @dialyzer {:no_return, unquote(:"#{name}_args_spec/#{arity}")}
        defp unquote(:"#{name}_args_spec")(unquote_splicing(args_names)) do
          super(unquote_splicing(args_names)) or unquote(args_check)
        end

        @dialyzer {:no_return, unquote(:"#{name}_result_spec/#{arity + 1}")}
        defp unquote(:"#{name}_result_spec")(unquote_splicing(args_names ++ [result])) do
          super(unquote_splicing(args_names ++ [result])) or
            (unquote(args_check) and TypeCheck.conforms?(unquote(result), unquote(result_type)))
        end
      end
    else
      Module.put_attribute(context, :specs, {name, arity})

      quote context: context do
        @spec unquote_splicing(expr)

        @dialyzer {:nowarn_function, unquote({:"#{name}_args_spec", arity})}
        defp unquote(:"#{name}_args_spec")(unquote_splicing(args_names)) do
          unquote(args_check)
        end

        @dialyzer {:nowarn_function, unquote({:"#{name}_result_spec", arity + 1})}
        defp unquote(:"#{name}_result_spec")(unquote_splicing(args_names ++ [result])) do
          unquote(args_check) and TypeCheck.conforms?(unquote(result), unquote(result_type))
        end
      end
    end
  end

  defmacro @ast do
    case ast do
      {name, _, expr} when name in ~w[type! typep! opaque! spec!]a ->
        quote do
          TypeCheck.Macros.unquote(name)(unquote_splicing(expr))
        end

      _ ->
        quote do
          Kernel.@(unquote(ast))
        end
    end
  end

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

    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)

    for f = {_def_t, name, arity} <- defs ++ defps,
        function_with_contracts?(f, module),
        {_, _, location, _} = Module.get_definition(module, {name, arity}) do
      file = __CALLER__.file
      Module.make_overridable(module, [{name, arity}])
      new_def(f, module, location ++ [file: file])
    end
  end

  defp new_def({def_t, name, arity}, context, line: line, file: _file) do
    args = Macro.generate_arguments(arity, nil)
    defs = Module.definitions_in(context)

    args_spec =
      if {:"#{name}_args_spec", arity} in defs do
        quote context: context, line: line do
          Corsa.assert(
            unquote(:"#{name}_args_spec")(unquote_splicing(args)),
            Corsa.SpecArgViolationError,
            Corsa.SpecError,
            call: call
          )
        end
      end

    precondition =
      if {:"#{name}_pre", arity} in defs do
        quote context: context, line: line do
          unquote(:"#{name}_pre")(unquote_splicing(args))
        end
      end

    within =
      if {:"#{name}_within", arity} in defs do
        quote context: context, line: line do
          within = unquote(:"#{name}_within")(unquote_splicing(args))

          timer_ref =
            if within do
              Gtimer.new_timer(
                within,
                fn pid ->
                  %{call: call}
                  |> Corsa.WithinViolationError.exception()
                  |> then(&Process.exit(pid, &1))
                end
              )
            end
        end
      end

    post_within =
      if {:"#{name}_within", arity} in defs do
        quote context: context, line: line do
          Gtimer.cancel_timer(timer_ref)
        end
      end

    throws =
      cond do
        {:"#{name}_throws", arity + 2} in defs ->
          quote context: context, line: line do
            unquote(:"#{name}_throws")(unquote_splicing(args), kind, value)
          end

        {:"#{name}_post", arity + 1} in defs ->
          quote context: context, line: line do
            reraise(
              Corsa.PostViolationError,
              "@post does not hold in call \"#{call}\" which raised an exception of kind \"#{kind}\" with the value \"#{inspect(value)}\"",
              __STACKTRACE__
            )
          end

        true ->
          nil
      end

    postcondition =
      if {:"#{name}_post", arity + 1} in defs do
        quote context: context, line: line do
          unquote(:"#{name}_post")(unquote_splicing(args), result)
        end
      end

    result_spec =
      if {:"#{name}_result_spec", arity + 1} in defs do
        quote context: context, line: line do
          Corsa.assert(
            unquote(:"#{name}_result_spec")(unquote_splicing(args), result),
            Corsa.SpecResultViolationError,
            Corsa.SpecError,
            call: call,
            result: result
          )
        end
      end

    decreases =
      if {:"#{name}_decreases", arity} in defs do
        quote context: context, line: line do
          unquote(:"#{name}_decreases")(unquote_splicing(args))
        end
      end

    quote context: context, line: line do
      unquote(def_t)(unquote(name)(unquote_splicing(args))) do
        call = to_corsa_call([unquote_splicing(args)], __MODULE__, unquote(name))
        unquote(within)

        try do
          unquote(args_spec)
          unquote(precondition)
          unquote(decreases)
          super(unquote_splicing(args))
        catch
          kind, value = %{corsa_exception: true} ->
            Corsa.Stack.flush({__MODULE__, unquote(name), unquote(arity)})
            :erlang.raise(kind, value, __STACKTRACE__)

          kind, value ->
            unquote(throws)
            unquote(post_within)
            Corsa.Stack.flush({__MODULE__, unquote(name), unquote(arity)})
            :erlang.raise(kind, value, __STACKTRACE__)
        else
          result ->
            unquote(post_within)
            unquote(result_spec)
            unquote(postcondition)
            Corsa.Stack.pop({__MODULE__, unquote(name), unquote(arity)})
            result
        end
      end
    end
  end

  ##############################################################################
  # Helper macros
  ##############################################################################

  defmacro assert(
             expr,
             violation \\ Corsa.AssertViolationError,
             error \\ Corsa.AssertError,
             opts \\ []
           ) do
    {_, vars} = free_vars(expr)
    context = __CALLER__.module
    vars = for var = {name, _, _} <- vars, do: {name, var}
    logger_level = Module.get_attribute(__CALLER__.module, :options) |> Keyword.get(:logger)

    expr_var = Macro.unique_var(:expr, context)
    vars_var = Macro.unique_var(:vars, context)
    opts_var = Macro.unique_var(:opts, context)
    violation_var = Macro.unique_var(:violation, context)
    error_var = Macro.unique_var(:error, context)
    result_var = Macro.unique_var(:result, context)

    check =
      if logger_level do
        quote context: context do
          {:current_stacktrace, [_entry1, entry2 | stacktrace]} =
            Process.info(self(), :current_stacktrace)

          {module, function, arity, [file: source, line: line]} = entry2

          Map.new(unquote(opts_var))
          |> unquote(violation).exception()
          |> then(
            &Logger.log(unquote(logger_level), "#{&1.message} [#{source}:#{line}]",
              crash_reason: {&1, entry2}
            )
          )
        end
      else
        quote context: context do
          raise(unquote(violation_var), Map.new(unquote(opts_var)))
        end
      end

    quote context: context do
      unquote(expr_var) = unquote(Macro.escape(expr))
      unquote(vars_var) = unquote(vars)
      unquote(opts_var) = unquote(opts) ++ [expr: unquote(expr_var), vars: unquote(vars_var)]
      unquote(violation_var) = unquote(violation)
      unquote(error_var) = unquote(error)

      try do
        unquote(expr)
      else
        unquote(result_var) ->
          if unquote(result_var) do
            true
          else
            unquote(check)
          end
      catch
        k, v ->
          unquote(opts_var) = unquote(opts_var) ++ [kind: k, reason: v]
          reraise unquote(error), Map.new(unquote(opts_var)), __STACKTRACE__
      end
    end
  end

  ##############################################################################
  # Helper functions
  ##############################################################################

  defp free_vars(expr) do
    {_, free_vars} =
      Macro.prewalk(expr, [], fn
        e = {_, _, nil}, acc -> {e, [e | acc]}
        e = {_, _, ctx}, acc when is_atom(ctx) -> {e, [e | acc]}
        e, acc -> {e, acc}
      end)

    {expr, MapSet.new(free_vars)}
  end

  defp free_vars_names(expr) do
    {expr, vars} = free_vars(expr)
    {expr, Enum.map(vars, &elem(&1, 0)) |> MapSet.new()}
  end

  defp parse_expects(expr, caller) do
    expr = Enum.reverse(expr) |> check_expects_syntax(caller)

    predicates =
      expr
      |> List.delete_at(0)
      |> List.delete_at(0)
      |> Enum.map(&free_vars_names/1)
      |> Enum.reverse()

    clauses = expr |> hd() |> hd() |> elem(1)
    {predicates, clauses}
  end

  defp check_expects_syntax([[do: {:__block__, _, []}] | _rest], caller) do
    reraise(Corsa.ExpectsError, "syntax error in @expects", Macro.Env.stacktrace(caller))
  end

  defp check_expects_syntax(expr = [[do: _block] | rest], caller) do
    if rest == [] || not match?({:receive, _, _}, hd(rest)) do
      reraise(Corsa.ExpectsError, "syntax error in @expects", Macro.Env.stacktrace(caller))
    end

    expr
  end

  defp check_expects_syntax(_expr, caller) do
    reraise(Corsa.ExpectsError, "syntax error in @expects", Macro.Env.stacktrace(caller))
  end

  defp function_with_contracts?({_, name, arity}, module) do
    defs = Module.definitions_in(module)

    {:"#{name}_args_spec", arity} in defs or
      {:"#{name}_pre", arity} in defs or
      {:"#{name}_within", arity} in defs or
      {:"#{name}_throws", arity + 2} in defs or
      {:"#{name}_post", arity + 1} in defs or
      {:"#{name}_result_spec", arity + 1} in defs or
      {:"#{name}_decreases", arity} in defs
  end

  def to_corsa_call(args, module, name) do
    Enum.map_join(args, ", ", &inspect(&1)) |> then(&"#{inspect(module)}.#{name}(#{&1})")
  end

  # def debug_stacktrace(label) do
  #   IO.puts(inspect(label))
  #   {_, stacktrace} = Process.info(self(), :current_stacktrace)
  #   Exception.format_stacktrace(tl(tl(stacktrace))) |> IO.puts()
  #   :ok
  # end

  # defp macro_inspect(m) do
  #   Macro.to_string(m) |> IO.puts()
  #   m
  # end
end