ex_ahocora.ex

ex_ahocora.ex

defmodule State do
  defstruct id: 0, next: %{}, failure: 0, outputs: []

  def update_failure(%State{} = state, new_id) do
    Map.update!(state, :failure, fn _ -> new_id end)
  end

  def add_next(%State{} = state, character, id) do
    Map.update!(state, :next, &Map.put(&1, character, id))
  end

  def add_outputs(%State{} = state, output) when is_binary(output) do
    Map.update!(state, :outputs, &[output | &1])
  end

  def add_outputs(%State{} = state, outputs) when is_list(outputs) do
    Map.update!(state, :outputs, &(&1 ++ outputs))
  end

  def next_id(%State{} = state, character) do
    state.next
    |> Map.get(character)
  end
end

defmodule ExAhocora do
  def build(terms) do
    make_goto(terms)
    |> make_failure()
  end

  def make_goto(terms, states \\ [%State{}])
  def make_goto([], states), do: states

  def make_goto([term | tail], states) do
    {cur_id, states} =
      term
      |> String.codepoints()
      |> Enum.reduce({0, states}, fn char, acc ->
        make_goto_process(char, acc)
      end)

    states = List.update_at(states, cur_id, &State.add_outputs(&1, term))

    make_goto(tail, states)
  end

  defp make_goto_process(char, {cur_id, states}) do
    cur = Enum.at(states, cur_id)

    unless Map.has_key?(cur.next, char) do
      new = %State{id: length(states)}
      states = List.update_at(states, cur.id, &State.add_next(&1, char, new.id))
      {new.id, states ++ [new]}
    else
      cur_id = Map.get(cur.next, char)
      {cur_id, states}
    end
  end

  def make_failure(states) do
    q =
      states
      |> List.first()
      |> Map.fetch!(:next)
      |> Map.values()
      |> Enum.reduce(:queue.new(), fn index, queue ->
        :queue.in(index, queue)
      end)

    make_failure_process(:queue.out(q), states)
  end

  def make_failure_process({:empty, _}, states), do: states

  def make_failure_process({{:value, i}, next_q}, states) do
    cur_state = Enum.at(states, i)

    {new_states, q} =
      cur_state
      |> Map.fetch!(:next)
      |> Map.keys()
      |> Enum.reduce({states, next_q}, fn character, {states, q} ->
        target_state_id = State.next_id(cur_state, character)
        q = :queue.in(target_state_id, q)

        f = trace_states(cur_state.failure, states, character)

        new_states =
          List.update_at(states, target_state_id, fn s ->
            s
            |> State.update_failure(f.id)
            |> State.add_outputs(f.outputs)
          end)

        {new_states, q}
      end)

    make_failure_process(:queue.out(q), new_states)
  end

  def trace_states(id, states, character) when is_integer(id) do
    cur = Enum.at(states, id)

    case State.next_id(cur, character) do
      nil ->
        case id do
          0 -> Enum.at(states, 0)
          _ -> trace_states(cur.failure, states, character)
        end

      next_id ->
        Enum.at(states, next_id)
    end
  end

  def match(states, word) do
    {_, _, matches} =
      (word <> <<0>>)
      |> String.codepoints()
      |> Enum.reduce({0, 0, []}, fn character, {cur_id, i, matches} ->
        s = Enum.at(states, cur_id)

        next_id =
          case State.next_id(s, character) do
            nil ->
              case Enum.at(states, s.failure) |> State.next_id(character) do
                nil -> 0
                next_id -> next_id
              end

            next_id ->
              next_id
          end

        matches =
          case s.outputs do
            [] ->
              matches

            outputs ->
              matches ++
                Enum.reduce(outputs, [], fn o, acc -> [{i - String.length(o), i - 1, o} | acc] end)
          end

        {next_id, i + 1, matches}
      end)

    matches
  end
end