A wrapper for the fastest Elixir JSON encode/decode library (:jiffy) with atomized keys

defmodule JSON do
  @moduledoc false

  @encode_opts [:use_nil]
  @decode_opts [:return_maps, :use_nil]

  alias :jiffy, as: Jiffy

  @doc """
  Encode and return tuple
  """
  @spec encode(any) :: {:ok, Strint.t()} | {:error, tuple()}
  def encode(nil) do
    {:ok, nil}
  end

  def encode(payload) do
    {:ok, Jiffy.encode(payload, @encode_opts)}
  catch
    {:error, reason} -> {:error, reason}
  end

  @doc """
  Encode and return encoded json
  """
  @spec encode!(any) :: Strint.t() | {:error, tuple()}
  def encode!(payload) do
    case encode(payload) do
      {:ok, json} -> json
      {:error, reason} -> {:error, reason}
    end
  end

  @doc """
  Decode and return tuple
  """
  @spec decode(String.t() | nil) :: {:ok, Map.t() | list()} | {:error, tuple()}
  def decode(payload) do
    {:ok, Jiffy.decode("#{payload}", @decode_opts)}
  catch
    {:error, reason} -> {:error, reason}
  end

  @doc """
  Decode and return decoded data
  """
  @spec decode!(String.t() | nil) :: Map.t() | list() | {:error, tuple()}
  def decode!(payload) do
    case decode(payload) do
      {:ok, data} -> data
      {:error, reason} -> {:error, reason}
    end
  end

  @doc """
  Convert map string keys to atom keys
  This function is an optional to atomize map keys instead of having
  string keys. 
  Warning: By default, the maximum number of atoms is 1,048,576.
           This function should be used responsibly. Personally,
           I recommend to use this function if you only can validate
           the JSON schemas.
  """
  @spec atomize_keys(any()) :: any()
  # Walk through the map
  def atomize_keys(map = %{}) do
    map
    |> Enum.map(fn {k, v} -> {to_atom(k), atomize_keys(v)} end)
    |> Enum.into(%{})
  end

  # Walk through the list
  def atomize_keys([head | rest]) do
    [atomize_keys(head) | atomize_keys(rest)]
  end

  # Other data types
  def atomize_keys(other) do
    other
  end

  defp to_atom(key) when is_atom(key) do
    key
  end

  defp to_atom(key) when is_binary(key) do
    String.to_atom(key)
  end

  defp to_atom(key) do
    key
  end
end