bin_codex

defmodule Codex do
  @moduledoc """
  Provides functions to create composable and bidirectional serializers.

  # Introduction
  The `Codex` module provides a number of predefined codecs and combinators that you can use to build new codes.
  ```
  iex> import Codex
  ...> first_codec = [byte(), byte(), byte()] |> sequence()
  ...> <<0x10, 0xFF, 0xAB>> |> decode(first_codec)
  {:ok, [16, 255, 171], <<>>}
  ```
  """

  @derive {Inspect, except: [:encode, :decode]}
  defstruct encode: nil, decode: nil

  @type type :: any
  @type type2 :: any
  @type remaining_bits :: bitstring

  @type encode_result :: {:ok, bitstring} | {:error, String.t()}
  @type decode_result(a) :: {:ok, a, remaining_bits} | {:error, String.t(), remaining_bits}

  @type encoder(a) :: (a -> encode_result)
  @type decoder(a) :: (bitstring -> decode_result(a))

  @type codec(a) :: %Codex{encode: encoder(a), decode: decoder(a)}

  @type boolean_codec :: codec(boolean)
  @type list_codec(a) :: codec(list(a))

  @doc """
  Creates a new codec.

  It is recommended to use the other, well tested, functions in the module to build a codec instead of creating your own unless you absolutely have to.

  When creating your own codec it *MUST* be idempotent.

  ## Examples
  ```
  iex> Codex.create(fn _ -> {:ok, <<>>} end, fn
  ...>   <<>> -> {:ok, false, <<>>}
  ...>   bits -> {:ok, true, bits}
  ...> end)
  #Codex<...>
  ```
  """
  @spec create(encoder(type), decoder(type)) :: %Codex{}
  def create(encode, decode) when is_function(encode, 1) and is_function(decode, 1),
    do: %Codex{encode: safe_encode(encode), decode: safe_decode(decode)}

  @doc """
  Encodes a value to its binary form using the supplied codec.

  ## Examples
  ```
  iex> import Codex
  ...> 198 |> encode(byte())
  {:ok, <<198>>}
  ```
  """
  @spec encode(type, codec(type)) :: encode_result
  def encode(value, %Codex{encode: encode}) when is_function(encode, 1), do: encode.(value)
  def encode(_value, %Codex{encode: _encode}), do: nil

  @doc """
  Decodes a value from its binary form using the supplied codec.

  The third element of the tuple is any remaining unparsed bits.

  ## Examples
  ```
  iex> import Codex
  iex> <<198>> |> decode(byte())
  {:ok, 198, <<>>}
  ```
  """
  @spec decode(bitstring, codec(type)) :: decode_result(type)
  def decode(bits, %Codex{decode: decode}) when is_function(decode, 1), do: decode.(bits)
  def decode(_bits, %Codex{decode: _decode}), do: nil

  @doc """
  Combines two codecs into a codec that produces a 2-element tuple of each value.

  ## Examples
  ```
  iex> import Codex
  ...> {198, 2} |> encode(combine(byte(), byte()))
  {:ok, <<198, 2>>}

  ...> <<198, 2>> |> decode(combine(byte(), byte()))
  {:ok, {198, 2}, <<>>}
  ```
  """
  @spec combine(codec(type), codec(type2)) :: codec({type, type2})
  def combine(codec1, codec2),
    do: create(&combine_encoder(codec1, codec2, &1), &combine_decoder(codec1, codec2, &1))

  @doc """
  A codec that always encodes to and decodes from the same value.

  It fails if it doesn't see the expected value that it always encodes/decodes.

  ## Examples
  ```
  iex> import Codex
  ...> codec = constant(10, <<200>>)
  ...> <<200>> |> decode(codec)
  {:ok, 10, <<>>}

  ...> <<234>> |> decode(codec)
  {:error, "<<234>> did not equal <<200>> in constant", <<234>>}

  ...> 10 |> encode(codec)
  {:ok, <<200>>}

  ...> 22 |> encode(codec)
  {:error, "22 did not equal 10 in constant"}
  ```
  """
  @spec constant(type, bitstring) :: codec(type)
  def constant(value, bits),
    do: create(&constant_encoder(value, bits, &1), &constant_decoder(value, bits, &1))

  @doc """
  A codec that always decodes to a certain value.

  It never consumes bits while decoding so it can't fail. It can fail on encoding if the value to encode doesn't match.

  ## Examples
  ```
  iex> import Codex
  ...> codec = value(10)
  ...> <<>> |> decode(codec)
  {:ok, 10, <<>>}

  ...> <<200>> |> decode(codec)
  {:ok, 10, <<200>>}

  ...> 10 |> encode(codec)
  {:ok, <<>>}

  ...> 22 |> encode(codec)
  {:error, "22 did not equal 10 in constant"}
  ```
  """
  @spec value(type) :: codec(type)
  def value(value), do: constant(value, <<>>)

  @spec empty() :: codec(any)
  def empty(), do: value([])

  @spec nothing() :: codec(any)
  def nothing(), do: value(nil)

  @doc """
  Uses the first codec, falling back to the second if it fails.

  ## Examples
  ```
  iex> import Codex
  ...> optional_byte = byte() |> fallback(nothing())
  ...> <<8>> |> Codex.decode(optional_byte)
  {:ok, 8, <<>>}

  ...> <<8::4>> |> Codex.decode(optional_byte)
  {:ok, nil, <<8::4>>}
  ```
  """
  @spec fallback(codec(type), codec(type2)) :: codec(type | type2)
  def fallback(codec1, codec2),
    do: create(&fallback_encoder(codec1, codec2, &1), &fallback_decoder(codec1, codec2, &1))

  @doc """
  Creates a codec that always fails with the supplied error message.

  This is not useful on its own but can be useful when building other codecs.

  ## Examples
  ```
  iex> defmodule Example do
  ...>   import Codex
  ...>   def choose([]), do: fail("None of the choices worked")
  ...>   def choose([codec | rest]), do: fallback(codec, choice(rest))
  ...> end
  ...> <<1>> |> Codex.decode(Example.choose([Codex.byte(), Codex.bits(4)]))
  {:ok, 1, <<>>}

  ...> <<1::2>> |> Codex.decode(Example.choose([Codex.bits(4), Codex.byte()]))
  {:error, "None of the choices worked", <<1::2>>}
  ```
  """
  @spec fail(String.t()) :: codec(any)
  def fail(error), do: fail(error, error)

  @doc """
  Creates a codec that always fails with the supplied error messages.

  Same as `fail/1` but you can supply a different error message for encoding and decoding.
  """
  @spec fail(String.t(), String.t()) :: codec(any)
  def fail(encode_error, decode_error),
    do: create(fn _ -> {:error, encode_error} end, fn bits -> {:error, decode_error, bits} end)

  @spec choice([codec(type)]) :: list_codec(type)
  def choice([]), do: fail("None of the choices worked")
  def choice([codec | rest]), do: codec |> fallback(choice(rest))

  @spec optional(codec(type)) :: codec(type | nil)
  def optional(codec), do: fallback(codec, nothing())

  @spec peek(codec(type)) :: codec(type)
  def peek(codec) do
    decode = fn bits ->
      with {:ok, a, _remaining} <- codec.decode.(bits) do
        {:ok, a, bits}
      end
    end

    create(fn _ -> {:ok, <<>>} end, decode)
  end

  @spec recover(codec(any)) :: boolean_codec()
  def recover(codec), do: create(&recover_encoder(codec, &1), &recover_decoder(codec, &1))

  @spec lookahead(codec(any)) :: boolean_codec()
  def lookahead(codec), do: peek(recover(codec))

  @spec convert(codec(type), (type -> type2), (type2 -> type)) :: codec(type2)
  def convert(codec, convert_to, convert_from),
    do: create(&convert_encoder(codec, convert_from, &1), &convert_decoder(codec, convert_to, &1))

  @spec not_(boolean_codec()) :: boolean_codec()
  def not_(boolean_codec), do: boolean_codec |> convert(&!/1, &!/1)

  @spec then(codec(type), (type -> codec(type2)), (type2 -> type)) :: codec(type2)
  def then(codec, f, g),
    do: create(&then_encoder(codec, f, g, &1), &then_decoder(codec, f, &1))

  @spec ensure(codec(type), boolean_codec(), String.t()) :: codec(type)
  def ensure(codec, boolean_codec, error) do
    codec
    |> combine(boolean_codec)
    |> then(
      fn
        {_a, false} -> fail(error)
        {a, true} -> value(a)
      end,
      fn a -> {a, true} end
    )
  end

  @spec refute(codec(type), boolean_codec(), String.t()) :: codec(type)
  def refute(codec, boolean_codec, error), do: ensure(codec, not_(boolean_codec), error)

  @spec bits_remaining() :: boolean_codec()
  def bits_remaining() do
    create(fn _ -> {:ok, <<>>} end, fn
      <<>> -> {:ok, false, <<>>}
      bits -> {:ok, true, bits}
    end)
  end

  @spec done(codec(type)) :: codec(type)
  def done(codec), do: codec |> refute(bits_remaining(), "There was more to parse")

  @spec reverse(list_codec(type)) :: list_codec(type)
  def reverse(list_codec), do: list_codec |> convert(&Enum.reverse/1, &Enum.reverse/1)

  @spec cons(codec(type), list_codec(type)) :: list_codec(type)
  def cons(codec, list_codec) do
    codec
    |> combine(list_codec)
    |> convert(
      fn {head, rest} -> [head | rest] end,
      fn [head | rest] -> {head, rest} end
    )
  end

  @spec append(list_codec(type), codec(type)) :: list_codec(type)
  def append(list_codec, codec), do: cons(codec, reverse(list_codec)) |> reverse()

  @spec sequence([codec(type)]) :: list_codec(type)
  def sequence([]), do: empty()
  def sequence([codec | rest]), do: codec |> cons(sequence(rest))

  @spec take_while(boolean_codec, codec(type)) :: list_codec(type)
  def take_while(boolean_codec, codec) do
    boolean_codec
    |> then(
      fn
        true -> codec |> cons(take_while(boolean_codec, codec))
        false -> empty()
      end,
      fn
        [] -> false
        _ -> true
      end
    )
  end

  @spec take_until(codec(type), boolean_codec) :: list_codec(type)
  def take_until(codec, boolean_codec), do: take_while(not_(boolean_codec), codec)

  @spec list(codec(type)) :: list_codec(type)
  def list(codec), do: take_while(bits_remaining(), codec)

  @spec list_of(non_neg_integer, codec(type)) :: list_codec(type)
  def list_of(count, codec) when is_integer(count) and count >= 0,
    do: List.duplicate(codec, count) |> sequence()

  def list_of(count, _codec) when is_integer(count), do: raise("list_of count must be >= 0")
  def list_of(_count, _codec), do: raise("list_of count must be >= 0")

  @spec non_empty_list(codec(type)) :: codec(nonempty_list(type))
  def non_empty_list(codec), do: codec |> cons(list(codec))

  @spec length_prefixed(codec(non_neg_integer), codec(type)) :: list_codec(type)
  def length_prefixed(length_codec, codec),
    do: length_codec |> then(&list_of(&1, codec), &length/1)

  @spec join(list_codec(binary), pos_integer) :: codec(binary)
  def join(codec, group_size \\ 1), do: codec |> convert(&Enum.join/1, &split_by(&1, group_size))

  @spec bit() :: codec(0 | 1)
  def bit(), do: bits(1)

  @spec bits(non_neg_integer) :: codec(non_neg_integer)
  def bits(count), do: create(&bits_encoder(count, &1), &bits_decoder(count, &1))

  def pad(codec, bits) do
    use Bitwise
    codec |> convert(&(&1 <<< bits), &(&1 >>> bits))
  end

  @spec bool_bit() :: boolean_codec()
  def bool_bit() do
    bit()
    |> convert(
      fn
        1 -> true
        0 -> false
      end,
      fn
        true -> 1
        false -> 0
      end
    )
  end

  @spec byte() :: codec(non_neg_integer)
  def byte(), do: bytes(1)

  @spec bytes(non_neg_integer) :: codec(non_neg_integer)
  def bytes(count), do: bits(count * 8)

  defp safe_encode(encode) do
    fn a ->
      try do
        encode.(a)
      rescue
        _ -> {:error, "Failed to encode #{inspect(a)}"}
      end
    end
  end

  defp safe_decode(decode) do
    fn bits ->
      try do
        decode.(bits)
      rescue
        _ -> {:error, "Failed to decode #{inspect(bits)}"}
      end
    end
  end

  defp split_by(binary, group_size),
    do: binary |> String.codepoints() |> Enum.chunk_every(group_size) |> Enum.map(&Enum.join/1)

  defp bits_encoder(count, _n) when count <= 0, do: raise("count must be >= 0, was #{count}")

  defp bits_encoder(_count, n) when is_integer(n) and n < 0,
    do: {:error, "Cannot encode negative number #{n}"}

  defp bits_encoder(count, n) when is_integer(n) do
    if(n < Integer.pow(2, count),
      do: {:ok, <<n::size(count)>>},
      else: {:error, "#{n} cannot be encoded in #{count} bits"}
    )
  end

  defp bits_encoder(_count, other), do: {:error, "'#{inspect(other)}' is not a positive integer"}

  defp bits_decoder(count, bits) do
    case bits do
      <<n::size(count), rest::bits>> -> {:ok, n, rest}
      bits -> {:error, "Could not decode #{count} bits from #{inspect(bits)}", bits}
    end
  end

  defp combine_encoder(codec1, codec2, {a, b}) do
    with {:ok, a_bits} <- codec1.encode.(a),
         {:ok, b_bits} <- codec2.encode.(b) do
      {:ok, <<a_bits::bits, b_bits::bits>>}
    end
  end

  defp combine_decoder(codec1, codec2, bits) do
    with {:ok, a, bits} <- codec1.decode.(bits),
         {:ok, b, bits} <- codec2.decode.(bits) do
      {:ok, {a, b}, bits}
    end
  end

  defp constant_encoder(value, bits, a) do
    case a do
      ^value -> {:ok, bits}
      other -> {:error, "#{inspect(other)} did not equal #{inspect(value)} in constant"}
    end
  end

  defp constant_decoder(value, bits, b) do
    size = bit_size(bits)

    case b do
      <<b::size(size), rest::bits>> when <<b::size(size)>> == bits -> {:ok, value, rest}
      other -> {:error, "#{inspect(other)} did not equal #{inspect(bits)} in constant", other}
    end
  end

  defp recover_encoder(codec, a) do
    case codec.encode.(a) do
      {:ok, a} -> {:ok, a}
      {:error, _} -> {:ok, <<>>}
    end
  end

  defp recover_decoder(codec, bits) do
    case codec.decode.(bits) do
      {:ok, _, bits} -> {:ok, true, bits}
      {:error, _, _} -> {:ok, false, bits}
    end
  end

  defp convert_encoder(codec, convert_from, a), do: a |> convert_from.() |> codec.encode.()

  defp convert_decoder(codec, convert_to, a) do
    with {:ok, a, bits} <- codec.decode.(a) do
      {:ok, convert_to.(a), bits}
    end
  end

  defp then_encoder(codec, f, g, a) do
    prefix = g.(a)

    with {:ok, a_bits} <- codec.encode.(prefix),
         {:ok, b_bits} <- f.(prefix).encode.(a) do
      {:ok, <<a_bits::bits, b_bits::bits>>}
    end
  end

  defp then_decoder(codec, f, bits) do
    with {:ok, a, rest} <- codec.decode.(bits) do
      f.(a).decode.(rest)
    end
  end

  defp fallback_encoder(codec1, codec2, a) do
    with {:error, _} <- codec1.encode.(a) do
      codec2.encode.(a)
    end
  end

  defp fallback_decoder(codec1, codec2, bits) do
    with {:error, _, _} <- codec1.decode.(bits) do
      codec2.decode.(bits)
    end
  end
end