fast_ngram

defmodule FastNgram do
  @moduledoc """
  A fast and unicode aware letter & word N-gram library written in Elixir.
  """

  @doc """
  Returns a list of letter N-grams from the given `string`.

  ## Example
      iex> FastNgram.letter_ngrams("¥ · € · $", 3)
      ["¥ ·", " · ", "· €", " € ", "€ ·", " · ", "· $"]
      iex> FastNgram.letter_ngrams("", 2)
      []
      iex> FastNgram.letter_ngrams("abcd", 1)
      ["a", "b", "c", "d"]
      iex> FastNgram.letter_ngrams("abcde", 2)
      ["ab", "bc", "cd", "de"]
  """
  @spec letter_ngrams(String.t(), non_neg_integer) :: list
  def letter_ngrams(string, 1) do
    String.graphemes(string)
  end

  def letter_ngrams(string, n) when is_integer(n) and n > 1 do
    graphemes = String.graphemes(string)

    if length(graphemes) < n do
      []
    else
      sizes = Enum.map(graphemes, &byte_size/1)
      # Initial length of the first N-gram
      {initial_sizes, rest_sizes} = Enum.split(sizes, n)
      len = Enum.sum(initial_sizes)

      do_letter_ngrams(string, 0, len, rest_sizes, sizes)
    end
  end

  defp do_letter_ngrams(string, offset, len, [next_s | rest_sizes], [this_s | sizes]) do
    ngram = binary_part(string, offset, len)
    # New len = old_len - this_s + next_s
    [ngram | do_letter_ngrams(string, offset + this_s, len - this_s + next_s, rest_sizes, sizes)]
  end

  defp do_letter_ngrams(string, offset, len, [], _) do
    [binary_part(string, offset, len)]
  end

  @doc """
  Returns a list of word N-grams from the given `string`.

  ## Example
      iex> FastNgram.word_ngrams("the bus came to a halt", 2)
      ["the bus", "bus came", "came to", "to a", "a halt"]
      iex> FastNgram.word_ngrams("the bus came to a halt", 3)
      ["the bus came", "bus came to", "came to a", "to a halt"]
      iex> FastNgram.word_ngrams("", 2)
      []
      iex> FastNgram.word_ngrams("some words here", 1)
      ["some", "words", "here"]
  """
  @spec word_ngrams(String.t(), non_neg_integer) :: list
  def word_ngrams(string, 1) do
    string |> String.split()
  end

  def word_ngrams(string, n) when is_integer(n) and n > 0 do
    words = string |> String.split()

    do_word_ngrams(n, length(words), words)
  end

  defp do_word_ngrams(n, len, words) when len >= n do
    [Enum.take(words, n) |> Enum.join(" ") | do_word_ngrams(n, len - 1, tl(words))]
  end

  defp do_word_ngrams(_, _, _) do
    []
  end
end