firebird

defmodule Firebird.Compiler.ConstantPropagation do
  @moduledoc """
  Constant propagation pass for the Firebird compiler.

  Substitutes variables with their known constant values when the variable
  is bound to a literal via a `let` binding. This is especially effective
  after function inlining, where call arguments are lowered to let bindings.

  ## How it works

  1. Walk the IR expression tree, tracking a mapping of variable names to
     their known constant values (from `let` bindings).
  2. When a `{:var, name}` is encountered and `name` maps to a constant,
     replace it with the literal value.
  3. This enables subsequent constant folding to evaluate the expression
     entirely at compile time.

  ## Example

  Before constant propagation (after inlining `double(5)`):
  ```
  {:block, [
    {:let, :n, {:literal, 5}},
    {:binop, :mul, {:var, :n}, {:literal, 2}}
  ]}
  ```

  After constant propagation:
  ```
  {:block, [
    {:let, :n, {:literal, 5}},
    {:binop, :mul, {:literal, 5}, {:literal, 2}}
  ]}
  ```

  After subsequent constant folding:
  ```
  {:literal, 10}
  ```

  ## Safety

  Only propagates values that are:
  - Literal integers or floats (`{:literal, _}`)
  - Bound exactly once (single-assignment, which is always true in
    Firebird's IR since Elixir variables are single-assignment in each scope)
  - Not shadowed by a later `let` binding with the same name

  The pass does NOT remove dead `let` bindings — that is left to the
  dead code elimination pass in the optimizer.
  """

  alias Firebird.Compiler.IR

  @doc """
  Run constant propagation on a module IR.

  Returns `{:ok, optimized_module}`.
  """
  @spec propagate(IR.Module.t()) :: {:ok, IR.Module.t()}
  def propagate(%IR.Module{} = module) do
    optimized_functions =
      Enum.map(module.functions, fn func ->
        param_set = MapSet.new(func.params)
        optimized_body = propagate_expr(func.body, %{}, param_set)
        %{func | body: optimized_body}
      end)

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

  @doc """
  Run constant propagation on a single expression.

  `env` is a map of `%{variable_name => {:literal, value}}`.
  `params` is a MapSet of function parameter names (never propagated over).
  """
  @spec propagate_expr(term(), map(), MapSet.t()) :: term()
  def propagate_expr(expr, env \\ %{}, params \\ MapSet.new())

  # Variable reference: substitute if we have a known constant
  def propagate_expr({:var, name}, env, _params) do
    case Map.get(env, name) do
      {:literal, _} = lit -> lit
      _ -> {:var, name}
    end
  end

  # Let binding: if the value is a literal, record it in env for subsequent exprs
  def propagate_expr({:let, name, value}, env, params) do
    propagated_value = propagate_expr(value, env, params)
    {:let, name, propagated_value}
  end

  # Block: process sequentially, building up the constant environment
  def propagate_expr({:block, exprs}, env, params) do
    {propagated_exprs, _final_env} =
      Enum.map_reduce(exprs, env, fn expr, current_env ->
        propagated = propagate_expr(expr, current_env, params)

        # If this is a let binding to a literal, record it
        new_env =
          case propagated do
            {:let, name, {:literal, _} = lit} ->
              # Don't propagate function params (they can vary per call)
              if MapSet.member?(params, name) do
                current_env
              else
                Map.put(current_env, name, lit)
              end

            {:let, name, _} ->
              # Non-constant binding — remove any previous constant for this name
              Map.delete(current_env, name)

            _ ->
              current_env
          end

        {propagated, new_env}
      end)

    {:block, propagated_exprs}
  end

  # Binary operations
  def propagate_expr({:binop, op, left, right}, env, params) do
    {:binop, op, propagate_expr(left, env, params), propagate_expr(right, env, params)}
  end

  # Unary operations
  def propagate_expr({:unaryop, op, inner}, env, params) do
    {:unaryop, op, propagate_expr(inner, env, params)}
  end

  # Function calls
  def propagate_expr({:call, name, args}, env, params) do
    {:call, name, Enum.map(args, &propagate_expr(&1, env, params))}
  end

  # If expressions — propagate into all branches with the same env
  # (branches don't introduce new bindings visible outside)
  def propagate_expr({:if, cond_expr, then_b, else_b}, env, params) do
    {:if, propagate_expr(cond_expr, env, params), propagate_expr(then_b, env, params),
     propagate_expr(else_b, env, params)}
  end

  # Case expressions
  def propagate_expr({:case, subject, clauses}, env, params) do
    {:case, propagate_expr(subject, env, params),
     Enum.map(clauses, fn {pat, guard, body} ->
       # Pattern-bound variables shadow any constants in env
       pat_vars = extract_pattern_vars(pat)
       inner_env = Map.drop(env, pat_vars)

       propagated_guard =
         case guard do
           nil -> nil
           g -> propagate_expr(g, inner_env, params)
         end

       {pat, propagated_guard, propagate_expr(body, inner_env, params)}
     end)}
  end

  # Tail loop — propagate into body but don't propagate loop params
  # (they change on each iteration via tail_call)
  def propagate_expr({:tail_loop, loop_params, body}, env, params) do
    # Loop parameter names shadow any constants
    loop_param_names =
      Enum.map(loop_params, fn
        {name, _init} -> name
        name when is_atom(name) -> name
      end)

    inner_env = Map.drop(env, loop_param_names)
    {:tail_loop, loop_params, propagate_expr(body, inner_env, params)}
  end

  # Tail call — propagate into args
  def propagate_expr({:tail_call, name, args}, env, params) do
    {:tail_call, name, Enum.map(args, &propagate_expr(&1, env, params))}
  end

  # Literals and anything else pass through unchanged
  def propagate_expr({:literal, _} = lit, _env, _params), do: lit
  def propagate_expr(other, _env, _params), do: other

  # --- Helpers ---

  # Extract variable names bound by a pattern
  defp extract_pattern_vars({:var_pat, name}), do: [name]
  defp extract_pattern_vars({:literal_pat, _}), do: []
  defp extract_pattern_vars(:wildcard), do: []
  defp extract_pattern_vars(_), do: []
end