macro.ex

defmodule Zipper.Macro do
  @heads 2..10

  defmacro __using__(_) do
    zipper_heads = Enum.map(@heads, &build_zipper_head(&1, __CALLER__))
    zipper_bodies = Enum.map(@heads, &build_zipper_body(&1, __CALLER__))

    unzipper_heads = Enum.map(@heads, &build_unzipper_head(&1, __CALLER__))
    unzipper_bodies = Enum.map(@heads, &build_unzipper_body(&1, __CALLER__))
    unzipper_base = build_unzipper_base(__CALLER__)

    quote do
      unquote_splicing(zipper_heads)

      def unzip(e) when not is_list(e) do
        e |> Enum.to_list() |> unzip()
      end

      def unzip([]) do
        {[], []}
      end

      unquote_splicing(unzipper_heads)

      unquote_splicing(zipper_bodies)

      unquote_splicing(unzipper_bodies)
      unquote(unzipper_base)
    end
  end

  defp build_zipper_head(arg_count, %Macro.Env{} = caller) do
    list_args = prefixed_args("list", arg_count, caller.module)
    body_fn_name = zipper_body_name()

    ensure_list =
      Enum.map(list_args, fn arg_name ->
        quote do
          Enum.to_list(unquote(arg_name))
        end
      end)

    quote do
      def zip(unquote_splicing(list_args)) do
        unquote(body_fn_name)(unquote_splicing(ensure_list), [])
      end
    end
  end

  defp build_zipper_body(arg_count, %Macro.Env{} = caller) do
    fn_name = zipper_body_name()

    elems = prefixed_args("elem", arg_count, caller.module)
    rests = prefixed_args("rest", arg_count, caller.module)
    ignores = Enum.map(1..arg_count, fn _ -> Macro.var(:_, caller.module) end)
    elems_and_rest = Enum.zip(elems, rests)

    args =
      Enum.map(elems_and_rest, fn {elem, rest} ->
        quote do
          [unquote(elem) | unquote(rest)]
        end
      end)

    quote do
      defp unquote(fn_name)(unquote_splicing(args), acc) do
        unquote(fn_name)(unquote_splicing(rests), [{unquote_splicing(elems)} | acc])
      end

      defp unquote(fn_name)(unquote_splicing(ignores), acc) do
        Enum.reverse(acc)
      end
    end
  end

  defp build_zipper_base_case(arg_count, %Macro.Env{} = caller) do
    fn_name = zipper_body_name()
    ignored_args = Enum.map(1..arg_count, fn _ -> Macro.var(:_, caller.module) end)

    quote do
      defp unquote(fn_name)(unquote_splicing(ignored_args), acc) do
        Enum.reverse(acc)
      end
    end
  end

  defp zipper_body_name() do
    :do_zip
  end

  defp unzipper_body_name() do
    :do_unzip
  end

  defp build_unzipper_head(arg_count, %Macro.Env{} = caller) do
    body_fn_name = unzipper_body_name()
    tuple_match = Enum.map(1..arg_count, fn _ -> Macro.var(:_, caller.module) end)
    accumulators = Enum.map(1..arg_count, fn _ -> quote do: [] end)

    quote do
      def unzip([{unquote_splicing(tuple_match)} | _rest] = l) do
        l
        |> Enum.reverse()
        |> unquote(body_fn_name)({unquote_splicing(accumulators)})
      end
    end
  end

  defp build_unzipper_body(arg_count, %Macro.Env{} = caller) do
    body_fn_name = unzipper_body_name()
    elem_names = prefixed_args("elem", arg_count, caller.module)
    accumulator_names = prefixed_args("rest", arg_count, caller.module)

    unzippers =
      Enum.zip_with(elem_names, accumulator_names, fn elem_name, accumulator_name ->
        quote do
          [unquote(elem_name) | unquote(accumulator_name)]
        end
      end)

    quote do
      defp unquote(body_fn_name)(
             [{unquote_splicing(elem_names)} | rest],
             {unquote_splicing(accumulator_names)}
           ) do
        unquote(body_fn_name)(rest, {unquote_splicing(unzippers)})
      end
    end
  end

  defp build_unzipper_base(%Macro.Env{}) do
    body_fn_name = unzipper_body_name()

    quote do
      defp unquote(body_fn_name)([], acc) do
        acc
      end
    end
  end

  defp prefixed_args(prefix, amount, context) do
    for n <- 1..amount do
      Macro.var(:"#{prefix}_#{n}", context)
    end
  end
end