diff --git a/lib/retrieval.ex b/lib/retrieval.ex index ad7ce4a..4d0aa91 100644 --- a/lib/retrieval.ex +++ b/lib/retrieval.ex @@ -34,6 +34,10 @@ defmodule Trieval do insert(%Trie{}, binary) end + def new(binary_and_payload) when is_tuple(binary_and_payload) do + insert(%Trie{}, binary_and_payload) + end + @doc """ Inserts a binary or list of binaries into an existing trie. @@ -55,6 +59,21 @@ defmodule Trieval do %Trie{trie: _insert(trie, binary)} end + def insert(%Trie{trie: trie}, binary_and_payload) when is_tuple(binary_and_payload) do + %Trie{trie: _insert(trie, binary_and_payload)} + end + + defp _insert(trie, {<>, payload}) do + case Map.has_key?(trie, next) do + true -> Map.put(trie, next, _insert(trie[next], {rest, payload})) + false -> Map.put(trie, next, _insert(%{}, {rest, payload})) + end + end + + defp _insert(trie, {<<>>, payload}) do + Map.put(trie, :mark, payload) + end + defp _insert(trie, <>) do case Map.has_key?(trie, next) do true -> Map.put(trie, next, _insert(trie[next], rest)) @@ -90,7 +109,7 @@ defmodule Trieval do end end - defp _contains?(%{mark: :mark}, <<>>) do + defp _contains?(%{mark: _}, <<>>) do true end @@ -128,6 +147,7 @@ defmodule Trieval do defp _prefix(trie, <<>>, acc) do Enum.flat_map(trie, fn {:mark, :mark} -> [acc] + {:mark, payload} -> [{acc, payload}] {ch, sub_trie} -> _prefix(sub_trie, <<>>, acc <> <>) end) end @@ -188,7 +208,7 @@ defmodule Trieval do defp _pattern(trie, capture_map, [:wildcard|rest], acc) do Enum.flat_map(trie, fn - {:mark, :mark} -> [] + {:mark, _} -> [] {ch, sub_trie} -> _pattern(sub_trie, capture_map, rest, acc <> <>) end) end @@ -196,7 +216,7 @@ defmodule Trieval do 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} -> [] + {:mark, _} -> [] {ch, sub_trie} -> _pattern(sub_trie, capture_map, rest, acc <> <>) end) end @@ -204,7 +224,7 @@ defmodule Trieval do 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} -> [] + {:mark, _} -> [] {ch, sub_trie} -> _pattern(sub_trie, capture_map, rest, acc <> <>) end) end @@ -219,7 +239,7 @@ defmodule Trieval do end false -> Enum.flat_map(trie, fn - {:mark, :mark} -> [] + {:mark, _} -> [] {ch, sub_trie} -> capture_map = Map.put(capture_map, name, ch) _pattern(sub_trie, capture_map, rest, acc <> <>) @@ -238,7 +258,7 @@ defmodule Trieval do false -> pruned_trie = Enum.filter(trie, fn({k, _v}) -> !(Map.has_key?(exclusions, k)) end) Enum.flat_map(pruned_trie, fn - {:mark, :mark} -> [] + {:mark, _} -> [] {ch, sub_trie} -> capture_map = Map.put(capture_map, name, ch) _pattern(sub_trie, capture_map, rest, acc <> <>) @@ -257,7 +277,7 @@ defmodule Trieval do false -> pruned_trie = Enum.filter(trie, fn({k, _v}) -> Map.has_key?(inclusions, k) end) Enum.flat_map(pruned_trie, fn - {:mark, :mark} -> [] + {:mark, _} -> [] {ch, sub_trie} -> capture_map = Map.put(capture_map, name, ch) _pattern(sub_trie, capture_map, rest, acc <> <>) @@ -267,7 +287,10 @@ defmodule Trieval do defp _pattern(trie, _capture_map, [], acc) do case Map.has_key?(trie, :mark) do - true -> [acc] + true -> case Map.get(trie, :mark) do + :mark -> [acc] + payload -> [{acc, payload}] + end false -> [] end end diff --git a/test/retrieval_test.exs b/test/retrieval_test.exs index 29b3e0d..574fddd 100644 --- a/test/retrieval_test.exs +++ b/test/retrieval_test.exs @@ -8,6 +8,11 @@ defmodule TrievalTest do @test_trie Trieval.new(@test_data) + @payload_trie Enum.map(@test_data, fn item = <> -> + {item, <>} + end) |> Trieval.new + + test "empty trie" do assert Trieval.new == %Trieval.Trie{} end @@ -40,4 +45,17 @@ defmodule TrievalTest do assert Trieval.pattern(@test_trie, "{1[^okjh]}x[tnm]{1}*{2}{1}{2}") == ["extended"] end + test "payload prefix" do + assert Trieval.prefix(@payload_trie, "app") == [{"apple", "a"}, {"apply", "a"}] + assert Trieval.prefix(@payload_trie, "n") == [{"negative", "n"}] + assert Trieval.prefix(@payload_trie, "abc") == [] + end + + test "payload pattern" do + assert Trieval.pattern(@payload_trie, "*{1}{1}**") == [{"apple", "a"}, {"apply", "a"}] + assert Trieval.pattern(@payload_trie, "[^abc]{1}{1}**") == [] + assert Trieval.pattern(@payload_trie, "[co]**") == [{"cat", "c"}, {"out", "o"}] + assert Trieval.pattern(@payload_trie, "{1[^okjh]}x[tnm]{1}*{2}{1}{2}") == [{"extended", "e"}] + end + end