Ecto Polymorphism

## ecto_polymorphism.ex
defmodule Ecto.Polymorphic do

  defmacro __using__(_) do
    quote do
      require Ecto.Schema
      import Ecto.Schema, except: [belongs_to: 2, belongs_to: 3]
      import unquote(__MODULE__)
    end
  end

  defmacro belongs_to(name, types, opts \\ []) do
    if check_polymorphic_options!(opts[:polymorphic], types) do
      {:module, module, _binary, _fns} = Ecto.Polymorphic.Type.generate_ecto_type(__CALLER__.module, name, types)
      quote do
        unquote(__MODULE__).__belongs_to__(__MODULE__, unquote(name), unquote(Keyword.put(opts, :types, types)))
        field unquote(:"#{name}_type"), unquote(module)
      end
    else
      quote do
        Ecto.Schema.belongs_to(unquote(name), unquote(types), unquote(opts))
      end
    end
  end

  def check_polymorphic_options!(nil, _types), do: nil
  def check_polymorphic_options!(false, _types), do: nil
  def check_polymorphic_options!(true, types) when not(is_nil(types) or length(types) == 0), do: true
  def check_polymorhpic_options!(_, _types) do
    raise """
    Polymorphic relationships require knowing all the possible types at compile time. Pass them in as
    a keyword list mapping the expected database value to the Ecto Schema. Example:

        belongs_to :address, ["ShippingAddress": MyApp.ShippingAddress, "BillingAddress": MyApp.BillingAddress], polymophic: true

    Or if they're the same as the database value, just pass a list of Ecto Schemas:

        belongs_to :address, [MyApp.ShippingAddress, MyApp.BillingAddress], polymophic: true
    """
  end

  ### COPIED FROM ECTO w/ small modifications :( ###

  @valid_belongs_to_options [:foreign_key, :references, :define_field, :type, :types,
                             :on_replace, :defaults, :primary_key, :polymorphic]
  @doc false
  # def __belongs_to__(mod, name, queryable, opts) do
  def __belongs_to__(mod, name, opts) do
    check_options!(opts, @valid_belongs_to_options, "belongs_to/3")

    opts = Keyword.put_new(opts, :foreign_key, :"#{name}_id")
    foreign_key_type = opts[:type] || Module.get_attribute(mod, :foreign_key_type)

    if name == Keyword.get(opts, :foreign_key) do
      raise ArgumentError, "foreign_key #{inspect name} must be distinct from corresponding association name"
    end

    if Keyword.get(opts, :define_field, true) do
      Ecto.Schema.__field__(mod, opts[:foreign_key], foreign_key_type, opts)
    end

    struct =
      # association(mod, :one, name, Ecto.Association.BelongsTo, [queryable: queryable] ++ opts)
      association(mod, :one, name, Ecto.Polymorphic.Association.BelongsTo, opts)
    Module.put_attribute(mod, :changeset_fields, {name, {:assoc, struct}})
  end

  defp check_options!(opts, valid, fun_arity) do
    case Enum.find(opts, fn {k, _} -> not k in valid end) do
      {k, _} ->
        raise ArgumentError, "invalid option #{inspect k} for #{fun_arity}"
      nil ->
        :ok
    end
  end

  defp association(mod, cardinality, name, association, opts) do
    not_loaded  = %Ecto.Association.NotLoaded{__owner__: mod,
                    __field__: name, __cardinality__: cardinality}
    put_struct_field(mod, name, not_loaded)
    opts = [cardinality: cardinality] ++ opts
    struct = association.struct(mod, name, opts)
    Module.put_attribute(mod, :ecto_assocs, {name, struct})

    struct
  end

  defp put_struct_field(mod, name, assoc) do
    fields = Module.get_attribute(mod, :struct_fields)

    if List.keyfind(fields, name, 0) do
      raise ArgumentError, "field/association #{inspect name} is already set on schema"
    end

    Module.put_attribute(mod, :struct_fields, {name, assoc})
  end

  ### END -- COPIED FROM ECTO :( ###

  defmodule Association.BelongsTo do
    import Ecto.Query
    @behaviour Ecto.Association
    @on_replace_opts [:raise, :mark_as_invalid, :delete, :nilify, :update]
    defstruct [:field, :owner, :owner_key, :related_key, :type_field, :on_cast, :key_field,
               :on_replace, defaults: [], cardinality: :one, relationship: :parent, unique: true,
               assoc_query_receives_structs: true,
               # unused but required by ecto
               queryable: nil
              ]

    @doc false
    def struct(module, name, opts) do
      on_replace = Keyword.get(opts, :on_replace, :raise)

      unless on_replace in @on_replace_opts do
        raise ArgumentError, "invalid `:on_replace` option for #{inspect name}. " <>
          "The only valid options are: " <>
          Enum.map_join(@on_replace_opts, ", ", &"`#{inspect &1}`")
      end

      %__MODULE__{
        field: name,
        owner: module,
        owner_key: Keyword.fetch!(opts, :foreign_key),
        key_field: :"#{name}_id",
        type_field: :"#{name}_type",
        on_replace: on_replace,
        defaults: opts[:defaults] || [],
        # Set something queryable to make ecto happy
        queryable: default_queryable(opts[:types])
      }
    end

    @doc false
    defp default_queryable([{_db_value, module} | _]), do: module
    defp default_queryable([module | _]), do: module

    @doc false
    defdelegate build(refl, struct, attributes), to: Ecto.Association.BelongsTo

    @doc false
    def joins_query(_) do
      raise """
      #{__MODULE__} Join Error: Polymorphic associations cannot be joined with! Convert to a concrete table.
      Perhaps you meant to use preload instead?
      """
    end

    def assoc_query(_, _, _) do
      raise """
      #{__MODULE__} Association Error: Polymorphic associations cannot return an association query! Convert to a concrete table.
      Perhaps you meant to use preload instead?
      """
    end

    def preload(_refl, _repo, _query, []), do: []

    def preload(%{type_field: type_field, key_field: key_field, owner: mod} = refl, repo, base_query, [%mod{} | _] = structs, opts) do
      structs
      |> Enum.map(&(Map.fetch!(&1, type_field)))
      |> Enum.uniq()
      |> Enum.flat_map(fn type ->
        values = Enum.filter(structs, &(Map.fetch!(&1, type_field) == type)) |> Enum.map(&(Map.fetch!(&1, key_field))) |> Enum.uniq()
        [related_key] = type.__schema__(:primary_key)

        query = from x in type, where: field(x, ^related_key) in ^values
        query = %{query | select: base_query.select, prefix: base_query.prefix}

        Ecto.Repo.Preloader.normalize_query(query, refl, {0, related_key})
        |> case do
          list when is_list(list) -> list
          query -> repo.all(query, opts)
        end
      end)
    end

    def preload_info(%{type_field: _type_field, owner: _mod} = refl) do
      {:assoc, refl, nil}
    end

    defdelegate on_repo_change(data, changeset, meta, opts), to: Ecto.Association.BelongsTo

    @doc false
    def after_compile_validation(_assoc, _env) do
      # TODO: Check stuff here
      :ok
    end

    ## Relation callbacks
    @behaviour Ecto.Changeset.Relation

    defdelegate build(assoc), to: Ecto.Association.BelongsTo
  end

  defmodule Type do
    def generate_ecto_type(module, name, mapping) do
      module = Module.concat(module, :"#{name |> to_string() |> String.capitalize()}PolymorphicType")
      quoted = quote bind_quoted: [module: module, mapping: mapping] do
        @behaviour Ecto.Type

        def type(), do: :string

        for {db_value, schema} <- mapping do
          def cast(unquote(to_string(db_value))), do: {:ok, unquote(schema)}
        end
        def cast(_), do: :error

        def load(type), do: cast(type)

        for {db_value, schema} <- mapping do
          def dump(unquote(schema)), do: {:ok, unquote(to_string(db_value))}
        end
        def dump(_), do: :error
      end
      Module.create(module, quoted, Macro.Env.location(__ENV__))
    end
  end
end