CljElixir/lib/clj_elixir/malli.ex

defmodule CljElixir.Malli do
  @moduledoc """
  Converts Malli-style schema data to Elixir typespec AST.

  Takes plain Elixir terms (atoms, lists, maps) representing Malli schemas
  and produces Elixir AST nodes suitable for `@spec` and `@type` attributes.

  ## Public API

    * `spec_ast/2` - Generate `@spec` AST nodes from a function schema
    * `type_ast/2,3` - Generate `@type` AST nodes from a type schema
    * `schema_to_typespec/2` - Convert a schema to its typespec AST (the type part only)
  """

  # Atoms that need quoted syntax in source but are valid at runtime
  @arrow :"=>"
  @optional_marker :"?"

  # ── Public API ──────────────────────────────────────────────────────

  @doc """
  Generates a list of `@spec` AST nodes for the given function name and schema.

  `schema` is either `[:=> ...]` for a single-arity function or
  `[:function ...]` for a multi-arity function.

  Returns a list because `:function` schemas and optional params
  can produce multiple `@spec` entries.
  """
  @spec spec_ast(atom(), list(), keyword()) :: list()
  def spec_ast(fun_name, schema, opts \\ [])

  def spec_ast(fun_name, [:function | clauses], opts) do
    Enum.flat_map(clauses, fn clause -> spec_ast(fun_name, clause, opts) end)
  end

  def spec_ast(fun_name, [@arrow, [:cat | param_schemas], return_schema], opts) do
    ret_ast = schema_to_typespec(return_schema, opts)
    param_groups = expand_optional_params(param_schemas)

    Enum.map(param_groups, fn params ->
      param_asts = Enum.map(params, &schema_to_typespec(&1, opts))
      wrap_spec(fun_name, param_asts, ret_ast)
    end)
  end

  @doc """
  Generates a `@type` AST node for the given type name and schema.

  Accepts an optional `opts` keyword list with `:known_types` for cross-references.
  For schemas with a `:registry` key, generates multiple types from the registry.
  """
  @spec type_ast(atom(), list(), keyword()) :: tuple() | list()
  def type_ast(type_name, schema, opts \\ [])

  def type_ast(type_name, [:schema, %{registry: registry}, ref_schema], opts) do
    type_ast_registry(type_name, [:schema, %{registry: registry}, ref_schema], registry, opts)
  end

  def type_ast(type_name, schema, opts) do
    type_body = schema_to_typespec(schema, opts)
    wrap_type(type_name, type_body)
  end

  @doc """
  Generates a list of `@type` AST nodes, one for each entry in the registry.

  `registry_types` is a map of `{name_atom => schema}` pairs or a list of
  `{name_atom, schema}` tuples.
  """
  def type_ast_registry(_type_name, [:schema, %{registry: _}, _ref], registry_types, opts) when is_map(registry_types) do
    Enum.map(registry_types, fn {name, schema} ->
      clean_name = clean_registry_name(name)
      body = schema_to_typespec(schema, Keyword.put(opts, :registry, registry_types))
      wrap_type(clean_name, body)
    end)
  end

  def type_ast_registry(_type_name, [:schema, %{registry: _}, _ref], registry_types, opts) when is_list(registry_types) do
    Enum.map(registry_types, fn {name, schema} ->
      clean_name = clean_registry_name(name)
      body = schema_to_typespec(schema, Keyword.put(opts, :registry, Map.new(registry_types)))
      wrap_type(clean_name, body)
    end)
  end

  # ── schema_to_typespec ──────────────────────────────────────────────

  @doc """
  Converts a schema to its typespec AST representation (the type part,
  not the `@type` wrapper).

  ## Options

    * `:known_types` - map of `%{"User" => :user, ...}` for cross-schema references
    * `:registry` - map of registry types for resolving `:ref` references
  """
  @spec schema_to_typespec(term(), keyword()) :: term()
  def schema_to_typespec(schema, opts \\ [])

  # ── Primitives ──────────────────────────────────────────────────────

  def schema_to_typespec(:string, _opts), do: string_t_ast()
  def schema_to_typespec(:int, _opts), do: {:integer, [], []}
  def schema_to_typespec(:integer, _opts), do: {:integer, [], []}
  def schema_to_typespec(:float, _opts), do: {:float, [], []}
  def schema_to_typespec(:number, _opts), do: {:number, [], []}
  def schema_to_typespec(:boolean, _opts), do: {:boolean, [], []}
  def schema_to_typespec(:atom, _opts), do: {:atom, [], []}
  def schema_to_typespec(:keyword, _opts), do: {:atom, [], []}
  def schema_to_typespec(:any, _opts), do: {:any, [], []}
  def schema_to_typespec(:nil, _opts), do: nil
  def schema_to_typespec(:pid, _opts), do: {:pid, [], []}
  def schema_to_typespec(:port, _opts), do: {:port, [], []}
  def schema_to_typespec(:reference, _opts), do: {:reference, [], []}
  def schema_to_typespec(:"pos-int", _opts), do: {:pos_integer, [], []}
  def schema_to_typespec(:"neg-int", _opts), do: {:neg_integer, [], []}
  def schema_to_typespec(:"nat-int", _opts), do: {:non_neg_integer, [], []}

  # ── Schema references (string keys) ────────────────────────────────

  def schema_to_typespec(name, opts) when is_binary(name) do
    known = Keyword.get(opts, :known_types, %{})

    case Map.fetch(known, name) do
      {:ok, type_name} -> {type_name, [], []}
      :error -> {:any, [], []}
    end
  end

  # ── Literal values (atoms that aren't schema keywords) ──────────────

  def schema_to_typespec(atom, _opts) when is_atom(atom), do: atom
  def schema_to_typespec(int, _opts) when is_integer(int), do: int

  # ── Compound and container types (list schemas) ─────────────────────

  def schema_to_typespec([head | _rest] = schema, opts) do
    convert_list_schema(head, schema, opts)
  end

  # ── Fallback ───────────────────────────────────────────────────────

  def schema_to_typespec(_, _opts), do: {:any, [], []}

  # ── List schema dispatch ────────────────────────────────────────────

  defp convert_list_schema(:or, [_ | types], opts) do
    type_asts = Enum.map(types, &schema_to_typespec(&1, opts))
    right_assoc_union(type_asts)
  end

  defp convert_list_schema(:and, [_ | schemas], opts) do
    resolve_and_type(schemas, opts)
  end

  defp convert_list_schema(:maybe, [:maybe, schema], opts) do
    inner = schema_to_typespec(schema, opts)
    {:|, [], [inner, nil]}
  end

  defp convert_list_schema(:enum, [_ | values], _opts) do
    right_assoc_union(values)
  end

  defp convert_list_schema(:=, [:=, value], _opts), do: value

  defp convert_list_schema(:map, [_ | field_specs], opts) do
    fields =
      Enum.map(field_specs, fn
        [name, {:optional, true}, schema] ->
          {name, schema_to_typespec(schema, opts)}

        [name, schema] ->
          {name, schema_to_typespec(schema, opts)}
      end)

    {:%{}, [], fields}
  end

  defp convert_list_schema(:"map-of", [_, key_schema, val_schema], opts) do
    key_ast = schema_to_typespec(key_schema, opts)
    val_ast = schema_to_typespec(val_schema, opts)
    {:%{}, [], [{{:optional, [], [key_ast]}, val_ast}]}
  end

  defp convert_list_schema(:list, [:list, elem_schema], opts) do
    [schema_to_typespec(elem_schema, opts)]
  end

  defp convert_list_schema(:vector, _schema, _opts) do
    persistent_vector_t_ast()
  end

  defp convert_list_schema(:set, _schema, _opts) do
    mapset_t_ast()
  end

  defp convert_list_schema(:tuple, [_ | elem_schemas], opts) do
    elems = Enum.map(elem_schemas, &schema_to_typespec(&1, opts))
    {:{}, [], elems}
  end

  defp convert_list_schema(:ref, [:ref, name], opts) do
    clean = clean_registry_name(name)
    registry = Keyword.get(opts, :registry, %{})

    if Map.has_key?(registry, name) or Map.has_key?(registry, clean) do
      {clean, [], []}
    else
      known = Keyword.get(opts, :known_types, %{})

      case Map.fetch(known, name) do
        {:ok, type_name} -> {type_name, [], []}
        :error -> {clean, [], []}
      end
    end
  end

  defp convert_list_schema(:schema, [:schema, %{registry: registry}, ref_schema], opts) do
    merged_opts = Keyword.put(opts, :registry, registry)
    schema_to_typespec(ref_schema, merged_opts)
  end

  defp convert_list_schema(:>, _, _opts), do: {:any, [], []}
  defp convert_list_schema(:>=, _, _opts), do: {:any, [], []}
  defp convert_list_schema(:<, _, _opts), do: {:any, [], []}
  defp convert_list_schema(:<=, _, _opts), do: {:any, [], []}

  defp convert_list_schema(head, schema, opts) when head == @arrow do
    [@arrow, [:cat | params], ret] = schema
    param_asts = Enum.map(params, &schema_to_typespec(&1, opts))
    ret_ast = schema_to_typespec(ret, opts)
    [{:->, [], [param_asts, ret_ast]}]
  end

  defp convert_list_schema(_, _, _opts), do: {:any, [], []}

  # ── Private helpers ─────────────────────────────────────────────────

  defp string_t_ast do
    {{:., [], [{:__aliases__, [alias: false], [:String]}, :t]}, [], []}
  end

  defp mapset_t_ast do
    {{:., [], [{:__aliases__, [alias: false], [:MapSet]}, :t]}, [], []}
  end

  defp persistent_vector_t_ast do
    {{:., [], [{:__aliases__, [alias: false], [:CljElixir, :PersistentVector]}, :t]}, [], []}
  end

  defp right_assoc_union([single]), do: single

  defp right_assoc_union([first | rest]) do
    {:|, [], [first, right_assoc_union(rest)]}
  end

  defp wrap_spec(fun_name, param_asts, ret_ast) do
    {:@, [], [
      {:spec, [], [
        {:"::", [], [
          {fun_name, [], param_asts},
          ret_ast
        ]}
      ]}
    ]}
  end

  defp wrap_type(type_name, body_ast) do
    {:@, [], [
      {:type, [], [
        {:"::", [], [
          {type_name, [], []},
          body_ast
        ]}
      ]}
    ]}
  end

  defp clean_registry_name(name) when is_atom(name) do
    name_str = Atom.to_string(name)

    cleaned =
      name_str
      |> String.replace(~r/^(Elixir\.CljElixir\.|::)/, "")

    String.to_atom(cleaned)
  end

  defp clean_registry_name(name), do: name

  # Expand optional params into all param combinations.
  # E.g., [:string, [:"?", :string], [:"?", :int]] produces:
  #   [[:string], [:string, :string], [:string, :string, :int]]
  defp expand_optional_params(param_schemas) do
    {required, optionals} = split_required_optional(param_schemas)

    for n <- 0..length(optionals) do
      required ++ Enum.take(optionals, n)
    end
  end

  defp split_required_optional(params) do
    split_required_optional(params, [])
  end

  defp split_required_optional([[@optional_marker, schema] | rest], req_acc) do
    optionals = [schema | extract_optionals(rest)]
    {Enum.reverse(req_acc), optionals}
  end

  defp split_required_optional([param | rest], req_acc) do
    split_required_optional(rest, [param | req_acc])
  end

  defp split_required_optional([], req_acc) do
    {Enum.reverse(req_acc), []}
  end

  defp extract_optionals([[@optional_marker, schema] | rest]) do
    [schema | extract_optionals(rest)]
  end

  defp extract_optionals([_ | rest]), do: extract_optionals(rest)
  defp extract_optionals([]), do: []

  # Resolve :and types — extract most specific expressible type.
  # Special cases: [:and :int [:> 0]] -> pos_integer()
  #                [:and :int [:>= 0]] -> non_neg_integer()
  defp resolve_and_type(schemas, opts) do
    base_types = Enum.filter(schemas, &recognized_schema?/1)
    constraints = Enum.filter(schemas, &constraint?/1)

    case {base_types, constraints} do
      {[:int], [[:>, 0]]} -> {:pos_integer, [], []}
      {[:integer], [[:>, 0]]} -> {:pos_integer, [], []}
      {[:int], [[:>=, 0]]} -> {:non_neg_integer, [], []}
      {[:integer], [[:>=, 0]]} -> {:non_neg_integer, [], []}
      {[base | _], _} -> schema_to_typespec(base, opts)
      {[], _} -> {:any, [], []}
    end
  end

  @primitive_types [
    :string, :int, :integer, :float, :number, :boolean, :atom, :keyword,
    :any, :nil, :pid, :port, :reference, :"pos-int", :"neg-int", :"nat-int"
  ]

  @compound_heads [:or, :and, :maybe, :enum, :=, :map, :"map-of",
                   :list, :vector, :set, :tuple, :ref, :schema]

  defp recognized_schema?(schema) when is_atom(schema) do
    schema in @primitive_types
  end

  defp recognized_schema?([head | _]) when is_atom(head) do
    head in @compound_heads or head == @arrow
  end

  defp recognized_schema?(_), do: false

  defp constraint?([:>, _]), do: true
  defp constraint?([:>=, _]), do: true
  defp constraint?([:<, _]), do: true
  defp constraint?([:<=, _]), do: true
  defp constraint?(_), do: false
end