jsv

defmodule JSV.Helpers.Traverse do
  @moduledoc """
  Helpers to work with nested generic data structures.
  """

  @type postwalk_cont :: (map, term -> {map, term})
  @type postwalk_item :: {:key, term} | {:val, term} | {:struct, struct, postwalk_cont}
  @type prewalk_item :: {:key, term} | {:val, term} | {:struct, struct}

  @doc """
  Updates a data structure in depth-first, post-order traversal.

  Operates like `postwalk/3` but without an accumulator. Handling continuations
  for structs require to handle the accumulator, whose value MUST be `nil`.
  """
  @spec postwalk(data, (postwalk_item() -> data)) :: data when data: term
  def postwalk(data, fun) when is_function(fun, 1) do
    {value, _} = postwalk(data, nil, fn value, nil -> {fun.(value), nil} end)
    value
  end

  @doc """
  Updates a JSON-compatible data structure in depth-first, post-order traversal
  while carrying an accumulator.

  The callback must return a `{new_value, new_acc}` tuple.

  Nested data structures are given to the callback before their wrappers, and
  when the wrappers are called, their children are already updated.

  JSON-compatible only means that there are restrictions on map keys and struct
  values:

  * The callback function will be called for any key but will not traverse the
    keys. For instance, with data such as `%{{x, y} => "some city"}`, the tuple
    used as key will be passed as-is but the callback will not be called for
    individual tuple elements.
  * Structs will be passed as `{:struct, value, continuation}`. The struct keys
    and values will **NOT** have been traversed yet. To operate on the struct
    keys you MUST call it manually. To respect the post-order of traversal, it
    SHOULD be called before further transformation of the struct:

        Traverse.postwalk(%MyStruct, [], fn
          {:struct, my_struct, cont}, acc ->
            {map, acc} = cont.(Map.from_struct(my_struct), acc)
            {struct!(MyStruct, do_something_with_map(map)), acc}
          {:val, ...} -> ...
        end)

    The continuation only accepts raw maps.

  * General data is accepted: tuples, pid, refs, etc. *
  """
  @spec postwalk(data, acc, (postwalk_item, acc -> {data, acc})) :: {data, acc} when data: term, acc: term
  def postwalk(data, acc, fun) when is_function(fun, 2) do
    {data, acc} = postwalk_subs(data, acc, fun)
    postwalk_parent(data, acc, fun)
  end

  defp postwalk_parent(struct, acc, fun) when is_struct(struct) do
    cont = fn
      map, acc when is_map(map) and not is_struct(map) ->
        postwalk_map_pairs(map, acc, fun)

      other, _ ->
        raise ArgumentError, "continuation function only accepts raw maps, got: #{inspect(other)}"
    end

    fun.({:struct, struct, cont}, acc)
  end

  defp postwalk_parent(data, acc, fun) do
    fun.({:val, data}, acc)
  end

  defp postwalk_subs(struct, acc, _fun) when is_struct(struct) do
    # Iteration on the struct is made in the continuation function
    {struct, acc}
  end

  defp postwalk_subs(map, acc, fun) when is_map(map) do
    postwalk_map_pairs(map, acc, fun)
  end

  defp postwalk_subs(list, acc, fun) when is_list(list) do
    Enum.map_reduce(list, acc, fn item, acc -> {_v, _acc} = postwalk(item, acc, fun) end)
  end

  # small optimization for 2-tuples
  defp postwalk_subs({a, b}, acc, fun) do
    {a, acc} = postwalk(a, acc, fun)
    {b, acc} = postwalk(b, acc, fun)
    {{a, b}, acc}
  end

  defp postwalk_subs(tuple, acc, fun) when is_tuple(tuple) do
    {elems, acc} =
      tuple
      |> Tuple.to_list()
      |> Enum.map_reduce(acc, fn item, acc -> {_v, _acc} = postwalk(item, acc, fun) end)

    {List.to_tuple(elems), acc}
  end

  defp postwalk_subs(data, acc, _fun) do
    {data, acc}
  end

  defp postwalk_map_pairs(map, acc, fun) do
    {pairs, acc} =
      Enum.map_reduce(map, acc, fn {k, v}, acc ->
        {v, acc} = postwalk(v, acc, fun)
        {k, acc} = fun.({:key, k}, acc)
        {{k, v}, acc}
      end)

    {Map.new(pairs), acc}
  end

  @doc """
  Updates a data structure in depth-first, pre-order traversal.

  Operates like `prewalk/3` but without an accumulator.
  """

  @spec prewalk(data, (prewalk_item() -> data)) :: data when data: term
  def prewalk(data, fun) when is_function(fun, 1) do
    {value, _} = prewalk(data, nil, fn value, nil -> {fun.(value), nil} end)
    value
  end

  @doc """
  Updates a JSON-compatible data structure in depth-first, pre-order traversal
  while carrying an accumulator.

  The callback must return a `{new_value, new_acc}` tuple.


  Nested data structures are given iterated after the parent data has been given
  to the function. So it is possible to accept a container (map, list, tuple)
  and return another one from the callback before the children are iterated.

  JSON-compatible only means that there are restrictions on map keys and struct
  values:

  * The callback function will be called for any key but will not traverse the
    keys. For instance, with data such as `%{{x, y} => "some city"}`, the tuple
    used as key will be passed as-is but the callback will not be called for
    individual tuple elements.
  * Structs are passed as a `{:struct, struct}` tuple.
  * General data is accepted: tuples, pid, refs, etc. *
  """
  @spec prewalk(data, acc, (prewalk_item, acc -> {data, acc})) :: {data, acc} when data: term, acc: term
  def prewalk(data, acc, fun) when is_function(fun, 2) do
    {new_value, acc} = prewalk_parent(data, acc, fun)
    prewalk_subs(new_value, acc, fun)
  end

  defp prewalk_parent(struct, acc, fun) when is_struct(struct) do
    fun.({:struct, struct}, acc)
  end

  defp prewalk_parent(val, acc, fun) do
    fun.({:val, val}, acc)
  end

  # TODO(doc) structs keys are not traversed
  defp prewalk_subs(%mod{} = struct, acc, fun) do
    # we need to preserve the struct, that means keep its keys.
    {pairs, acc} =
      struct
      |> Map.from_struct()
      |> Enum.map_reduce(acc, fn {k, v}, acc ->
        {v, acc} = prewalk(v, acc, fun)
        {{k, v}, acc}
      end)

    {struct(mod, pairs), acc}
  end

  defp prewalk_subs(map, acc, fun) when is_map(map) do
    prewalk_map_pairs(map, acc, fun)
  end

  defp prewalk_subs(list, acc, fun) when is_list(list) do
    Enum.map_reduce(list, acc, fn item, acc -> prewalk(item, acc, fun) end)
  end

  defp prewalk_subs({a, b}, acc, fun) do
    {a, acc} = prewalk(a, acc, fun)
    {b, acc} = prewalk(b, acc, fun)
    {{a, b}, acc}
  end

  defp prewalk_subs(tuple, acc, fun) when is_tuple(tuple) do
    {elems, acc} =
      tuple
      |> Tuple.to_list()
      |> Enum.map_reduce(acc, fn item, acc -> {_v, _acc} = prewalk(item, acc, fun) end)

    {List.to_tuple(elems), acc}
  end

  defp prewalk_subs(scalar, acc, _fun) do
    {scalar, acc}
  end

  defp prewalk_map_pairs(map, acc, fun) do
    {pairs, acc} =
      Enum.map_reduce(map, acc, fn {k, v}, acc ->
        {v, acc} = prewalk(v, acc, fun)
        {k, acc} = fun.({:key, k}, acc)
        {{k, v}, acc}
      end)

    {Map.new(pairs), acc}
  end
end