defmodule Ast do
  def var(name), do: {:var, name}
  def abs(param, body), do: {:abs, param, body}
  def app(func, param), do: {:app, func, param}

  def pprint(expr) do
    case expr do
      {:var, name} -> name
      {:abs, param, body} -> "(λ #{param} . #{pprint(body)})"
      {:app, func, param} -> "#{pprint(func)} #{pprint(param)}"
    end
  end

  def free_vars(expr) do
    case expr do
      {:var, name} -> MapSet.new([name])
      {:abs, param, body} -> MapSet.delete(free_vars(body), param)
      {:app, func, param} -> MapSet.union(free_vars(func), free_vars(param))
    end
  end

  def subst(expr, old_var, new_expr) do
    case expr do
      {:var, name} ->
        if name == old_var, do: new_expr, else: expr

      {:abs, param, body} ->
        if param == old_var, do: expr, else: {:abs, param, subst(body, old_var, new_expr)}

      {:app, func, param} ->
        {:app, subst(func, old_var, new_expr), subst(param, old_var, new_expr)}
    end
  end
end