firebird

defmodule Firebird.Compiler.Inliner do
  @moduledoc """
  Function inlining optimization for the Firebird compiler.

  Replaces calls to small, non-recursive functions with their body
  expressions. This eliminates function call overhead in the WASM output.

  ## When to inline

  A function is a candidate for inlining when:
  - It has a single clause (no pattern matching)
  - It is not recursive
  - Its body is "small" (fewer than a configurable number of IR nodes)
  - It's called from within the same module

  ## Example

  Before inlining:
  ```elixir
  def double(n), do: n * 2
  def quadruple(n), do: double(double(n))
  ```

  After inlining `double`:
  ```elixir
  def quadruple(n), do: (n * 2) * 2
  ```
  """

  alias Firebird.Compiler.IR
  alias Firebird.Compiler.TCO

  @default_max_size 10

  @doc """
  Run inlining optimization on a module IR.

  ## Options

  - `:max_size` - Maximum IR node count for inlining (default: #{@default_max_size})
  """
  @spec inline(IR.Module.t(), keyword()) :: {:ok, IR.Module.t()}
  def inline(%IR.Module{} = module, opts \\ []) do
    max_size = Keyword.get(opts, :max_size, @default_max_size)

    # Find inlinable functions
    inline_map = build_inline_map(module.functions, max_size)

    if Enum.empty?(inline_map) do
      {:ok, module}
    else
      # Apply inlining to all function bodies
      inlined_functions =
        Enum.map(module.functions, fn func ->
          %{func | body: inline_calls(func.body, inline_map, func.name)}
        end)

      {:ok, %{module | functions: inlined_functions}}
    end
  end

  @doc """
  Build a map of function name → {params, body} for inlinable functions.
  """
  @spec build_inline_map([IR.Function.t()], non_neg_integer()) :: map()
  def build_inline_map(functions, max_size) do
    functions
    |> Enum.filter(&inlinable?(&1, max_size))
    |> Map.new(fn func -> {func.name, {func.params, func.body}} end)
  end

  @doc """
  Check if a function is a good inlining candidate.
  """
  @spec inlinable?(IR.Function.t(), non_neg_integer()) :: boolean()
  def inlinable?(%IR.Function{} = func, max_size) do
    not TCO.has_self_calls?(func.body, func.name) and
      ir_size(func.body) <= max_size
  end

  @doc """
  Count the number of IR nodes in an expression.
  """
  @spec ir_size(term()) :: non_neg_integer()
  def ir_size({:literal, _}), do: 1
  def ir_size({:var, _}), do: 1
  def ir_size({:binop, _, left, right}), do: 1 + ir_size(left) + ir_size(right)
  def ir_size({:unaryop, _, expr}), do: 1 + ir_size(expr)
  def ir_size({:call, _, args}), do: 1 + Enum.sum(Enum.map(args, &ir_size/1))
  def ir_size({:if, c, t, e}), do: 1 + ir_size(c) + ir_size(t) + ir_size(e)

  def ir_size({:case, s, clauses}),
    do: 1 + ir_size(s) + Enum.sum(Enum.map(clauses, fn {_, _, b} -> ir_size(b) end))

  def ir_size({:let, _, v}), do: 1 + ir_size(v)
  def ir_size({:block, exprs}), do: Enum.sum(Enum.map(exprs, &ir_size/1))
  def ir_size({:tail_loop, _, body}), do: 1 + ir_size(body)
  def ir_size({:tail_call, _, args}), do: 1 + Enum.sum(Enum.map(args, &ir_size/1))
  def ir_size(_), do: 1

  # Inline function calls in an expression
  defp inline_calls({:call, name, args}, inline_map, current_func) do
    # Don't inline the current function (prevent infinite recursion)
    args = Enum.map(args, &inline_calls(&1, inline_map, current_func))

    if name != current_func and Map.has_key?(inline_map, name) do
      {params, body} = Map.get(inline_map, name)

      if length(args) == length(params) do
        # Substitute parameters with arguments
        var_map = Enum.zip(params, args) |> Map.new()
        substitute_inline(body, var_map)
      else
        {:call, name, args}
      end
    else
      {:call, name, args}
    end
  end

  defp inline_calls({:binop, op, left, right}, inline_map, cf) do
    {:binop, op, inline_calls(left, inline_map, cf), inline_calls(right, inline_map, cf)}
  end

  defp inline_calls({:unaryop, op, expr}, inline_map, cf) do
    {:unaryop, op, inline_calls(expr, inline_map, cf)}
  end

  defp inline_calls({:if, c, t, e}, inline_map, cf) do
    {:if, inline_calls(c, inline_map, cf), inline_calls(t, inline_map, cf),
     inline_calls(e, inline_map, cf)}
  end

  defp inline_calls({:case, s, clauses}, inline_map, cf) do
    {:case, inline_calls(s, inline_map, cf),
     Enum.map(clauses, fn {p, g, b} -> {p, g, inline_calls(b, inline_map, cf)} end)}
  end

  defp inline_calls({:let, name, value}, inline_map, cf) do
    {:let, name, inline_calls(value, inline_map, cf)}
  end

  defp inline_calls({:block, exprs}, inline_map, cf) do
    {:block, Enum.map(exprs, &inline_calls(&1, inline_map, cf))}
  end

  defp inline_calls(other, _inline_map, _cf), do: other

  # Substitute parameters with argument expressions
  defp substitute_inline({:var, name}, var_map) do
    Map.get(var_map, name, {:var, name})
  end

  defp substitute_inline({:binop, op, left, right}, var_map) do
    {:binop, op, substitute_inline(left, var_map), substitute_inline(right, var_map)}
  end

  defp substitute_inline({:unaryop, op, expr}, var_map) do
    {:unaryop, op, substitute_inline(expr, var_map)}
  end

  defp substitute_inline({:call, name, args}, var_map) do
    {:call, name, Enum.map(args, &substitute_inline(&1, var_map))}
  end

  defp substitute_inline({:if, c, t, e}, var_map) do
    {:if, substitute_inline(c, var_map), substitute_inline(t, var_map),
     substitute_inline(e, var_map)}
  end

  defp substitute_inline({:let, name, value}, var_map) do
    {:let, name, substitute_inline(value, var_map)}
  end

  defp substitute_inline({:block, exprs}, var_map) do
    {:block, Enum.map(exprs, &substitute_inline(&1, var_map))}
  end

  defp substitute_inline(other, _var_map), do: other
end