defmodule Retrieval do alias Retrieval.Trie alias Retrieval.PatternParser @moduledoc """ Provides an interface for creating and collecting data from the trie data structure. """ @doc """ Returns a new trie. Providing no arguments creates an empty trie. Optionally a binary or list of binaries can be passed to `new/1`. ## Examples Retrieval.new %Retrieval.Trie{...} Retrieval.new("apple") %Retrieval.Trie{...} Retrieval.new(~w/apple apply ape ample/) %Retrieval.Trie{...} """ def new, do: %Trie{} def new(binaries) when is_list(binaries) do insert(%Trie{}, binaries) end def new(binary) when is_binary(binary) do insert(%Trie{}, binary) end @doc """ Inserts a binary or list of binaries into an existing trie. ## Examples Retrieval.new |> Retrieval.insert("apple") %Retrieval.Trie{...} Retrieval.new(~w/apple apply ape ample/) |> Retrieval.insert(~w/zebra corgi/) %Retrieval.Trie{...} """ def insert(%Trie{trie: trie}, binaries) when is_list(binaries) do %Trie{trie: Enum.reduce(binaries, trie, &_insert(&2, &1))} end def insert(%Trie{trie: trie}, binary) when is_binary(binary) do %Trie{trie: _insert(trie, binary)} end defp _insert(trie, <>) do case Map.has_key?(trie, next) do true -> Map.put(trie, next, _insert(trie[next], rest)) false -> Map.put(trie, next, _insert(%{}, rest)) end end defp _insert(trie, <<>>) do Map.put(trie, :mark, :mark) end @doc """ Returns whether or not a trie contains a given binary key. ## Examples Retrieval.new(~w/apple apply ape ample/) |> Retrieval.contains?("apple") true Retrieval.new(~w/apple apply ape ample/) |> Retrieval.contains?("zebra") false """ def contains?(%Trie{trie: trie}, binary) when is_binary(binary) do _contains?(trie, binary) end defp _contains?(trie, <>) do case Map.has_key?(trie, next) do true -> _contains?(trie[next], rest) false -> false end end defp _contains?(%{mark: :mark}, <<>>) do true end @doc """ Collects all binaries that begin with a given prefix. ## Examples Retrieval.new(~w/apple apply ape ample/) |> Retrieval.prefix("ap") ["apple", "apply", "ape"] Retrieval.new(~w/apple apply ape ample/) |> Retrieval.prefix("z") [] """ def prefix(%Trie{trie: trie}, binary) when is_binary(binary) do _prefix(trie, binary, binary) end defp _prefix(trie, <>, acc) do case Map.has_key?(trie, next) do true -> _prefix(trie[next], rest, acc) false -> [] end end # An interesting discovery I made here is that treating the accumulator as a binary is actually quicker # than converting the prefix to a char list, prepending to it, reversing when a word is found, and converting # to a binary. defp _prefix(trie, <<>>, acc) do Enum.flat_map(trie, fn {:mark, :mark} -> [acc] {ch, sub_trie} -> _prefix(sub_trie, <<>>, acc <> <>) end) end @doc """ Collects all binaries match a given pattern. Returns either a list of matches or an error in the form `{:error, reason}`. ## Patterns `*` - Wildcard, matches any character. `[...]` - Inclusion group, matches any character between brackets. `[^...]` - Exclusion group, matches any character not between brackets. `{...}` - Capture group, must be named and can be combined with an inclusion or exclusion group, otherwise treated as a wildcard. All future instances of same name captures are swapped with the value of the initial capture. ## Examples Retrieval.new(~w/apple apply ape ample/) |> Retrieval.pattern("a{1}{1}**") ["apple", "apply"] Retrieval.new(~w/apple apply ape ample/) |> Retrieval.pattern("*{1[^p]}{1}**") [] Retrieval.new(~w/apple apply zebra house/) |> Retrieval.pattern("[hz]****") ["house", "zebra"] Retrieval.new(~w/apple apply zebra house/) |> Retrieval.pattern("[hz]***[^ea]") [] Retrieval.new(~w/apple apply zebra house/) |> Retrieval.pattern("[hz]***[^ea") {:error, "Dangling group (exclusion) starting at column 8, expecting ]"} """ def pattern(%Trie{trie: trie}, pattern) when is_binary(pattern) do _pattern(trie, %{}, pattern, <<>>, :parse) end defp _pattern(trie, capture_map, pattern, acc, :parse) do case PatternParser.parse(pattern) do {:error, message} -> {:error, message} parsed_pattern -> _pattern(trie, capture_map, parsed_pattern, acc) end end defp _pattern(trie, capture_map, [{:character, ch}|rest], acc) do case Map.has_key?(trie, ch) do true -> _pattern(trie[ch], capture_map, rest, acc <> <>) false -> [] end end defp _pattern(trie, capture_map, [:wildcard|rest], acc) do Enum.flat_map(trie, fn {:mark, :mark} -> [] {ch, sub_trie} -> _pattern(sub_trie, capture_map, rest, acc <> <>) end) end defp _pattern(trie, capture_map, [{:exclusion, exclusions}|rest], acc) do pruned_trie = Enum.filter(trie, fn({k, _v}) -> !(Map.has_key?(exclusions, k)) end) Enum.flat_map(pruned_trie, fn {:mark, :mark} -> [] {ch, sub_trie} -> _pattern(sub_trie, capture_map, rest, acc <> <>) end) end defp _pattern(trie, capture_map, [{:inclusion, inclusions}|rest], acc) do pruned_trie = Enum.filter(trie, fn({k, _v}) -> Map.has_key?(inclusions, k) end) Enum.flat_map(pruned_trie, fn {:mark, :mark} -> [] {ch, sub_trie} -> _pattern(sub_trie, capture_map, rest, acc <> <>) end) end defp _pattern(trie, capture_map, [{:capture, name}|rest], acc) do case Map.has_key?(capture_map, name) do true -> match = capture_map[name] case Map.has_key?(trie, match) do true -> _pattern(trie[match], capture_map, rest, acc <> <>) false -> [] end false -> Enum.flat_map(trie, fn {:mark, :mark} -> [] {ch, sub_trie} -> capture_map = Map.put(capture_map, name, ch) _pattern(sub_trie, capture_map, rest, acc <> <>) end) end end defp _pattern(trie, capture_map, [{:capture, name, :exclusion, exclusions}|rest], acc) do case Map.has_key?(capture_map, name) do true -> match = capture_map[name] case Map.has_key?(trie, match) do true -> _pattern(trie[match], capture_map, rest, acc <> <>) false -> [] end false -> pruned_trie = Enum.filter(trie, fn({k, _v}) -> !(Map.has_key?(exclusions, k)) end) Enum.flat_map(pruned_trie, fn {:mark, :mark} -> [] {ch, sub_trie} -> capture_map = Map.put(capture_map, name, ch) _pattern(sub_trie, capture_map, rest, acc <> <>) end) end end defp _pattern(trie, capture_map, [{:capture, name, :inclusion, inclusions}|rest], acc) do case Map.has_key?(capture_map, name) do true -> match = capture_map[name] case Map.has_key?(trie, match) do true -> _pattern(trie[match], capture_map, rest, acc <> <>) false -> [] end false -> pruned_trie = Enum.filter(trie, fn({k, _v}) -> Map.has_key?(inclusions, k) end) Enum.flat_map(pruned_trie, fn {:mark, :mark} -> [] {ch, sub_trie} -> capture_map = Map.put(capture_map, name, ch) _pattern(sub_trie, capture_map, rest, acc <> <>) end) end end defp _pattern(trie, _capture_map, [], acc) do case Map.has_key?(trie, :mark) do true -> [acc] false -> [] end end end