1 files changed, 141 insertions, 11 deletions
diff --git a/lib/tree_bitmap.ex b/lib/tree_bitmap.ex
index 6d1bc07..0da87fa 100644
--- a/lib/tree_bitmap.ex
+++ b/lib/tree_bitmap.ex
@@ -1,18 +1,148 @@
 defmodule TreeBitmap do
-  @moduledoc """
-  Documentation for `TreeBitmap`.
-  """
+  alias TreeBitmap.NIF
+  defstruct [:i4, :i6, :ets]
 
-  @doc """
-  Hello world.
+  @opaque t() :: %__MODULE__{}
+  @type masklen :: non_neg_integer()
 
-  ## Examples
+  def new(_opts \\ []) do
+    %__MODULE__{i4: NIF.new(), i6: NIF.new(), ets: :ets.new(__MODULE__, [:public])}
+  end
+
+  @spec add(t(), :inet.ip_address(), masklen(), any()) :: nil | any()
+  def add(tree, ip, masklen, value)
+
+  def add(tree, {a, b, c, d}, masklen, value) do
+    add(tree, tree.i4, {:inet4, a, b, c, d}, masklen, value)
+  end
+
+  def add(tree, {a, b, c, d, e, f, g, h}, masklen, value) do
+    add(tree, tree.i6, {:inet6, a, b, c, d, e, f, g, h}, masklen, value)
+  end
+
+  @spec remove(t(), :inet.ip_address(), masklen()) :: nil | any()
+  def remove(tree, ip, masklen)
+
+  def remove(tree, {a, b, c, d}, masklen) do
+    remove(tree, tree.i4, {:inet4, a, b, c, d}, masklen)
+  end
+
+  def remove(tree, {a, b, c, d, e, f, g, h}, masklen) do
+    remove(tree, tree.i6, {:inet6, a, b, c, d, e, f, g, h}, masklen)
+  end
+
+  @spec longest_match(t(), :inet.ip_address()) :: map() | nil
+  def longest_match(tree, ip)
+
+  def longest_match(tree, {a, b, c, d}) do
+    longest_match(tree, tree.i4, {:inet4, a, b, c, d})
+  end
+
+  def longest_match(tree, {a, b, c, d, e, f, g, h}) do
+    longest_match(tree, tree.i6, {:inet6, a, b, c, d, e, f, g, h})
+  end
+
+  @spec longest_match?(t(), :inet.ip_address()) :: boolean()
+  def longest_match?(tree, ip)
+
+  def longest_match?(tree, {a, b, c, d}) do
+    longest_match?(tree, tree.i4, {:inet4, a, b, c, d})
+  end
+
+  def longest_match?(tree, {a, b, c, d, e, f, g, h}) do
+    longest_match?(tree, tree.i6, {:inet6, a, b, c, d, e, f, g, h})
+  end
+
+  @spec exact_match(t(), :inet.ip_address(), masklen()) :: map() | nil
+  def exact_match(tree, ip, masklen)
+
+  def exact_match(tree, {a, b, c, d}, masklen) do
+    exact_match(tree, tree.i4, {:inet4, a, b, c, d}, masklen)
+  end
+
+  def exact_match(tree, {a, b, c, d, e, f, g, h}, masklen) do
+    exact_match(tree, tree.i6, {:inet6, a, b, c, d, e, f, g, h}, masklen)
+  end
+
+  @spec exact_match?(t(), :inet.ip_address(), masklen()) :: boolean()
+  def exact_match?(tree, ip, masklen)
+
+  def exact_match?(tree, {a, b, c, d}, masklen) do
+    exact_match?(tree, tree.i4, {:inet4, a, b, c, d}, masklen)
+  end
 
-      iex> TreeBitmap.hello()
-      :world
+  def exact_match?(tree, {a, b, c, d, e, f, g, h}, masklen) do
+    exact_match?(tree, tree.i6, {:inet6, a, b, c, d, e, f, g, h}, masklen)
+  end
+
+  @type tree_memory() :: {nodes :: non_neg_integer(), results :: non_neg_integer()}
+  @spec memory(t()) :: %{inet4: tree_memory(), inet6: tree_memory(), ets: non_neg_integer()}
+  def memory(tree) do
+    %{inet4: NIF.memory(tree.i4), inet6: NIF.memory(tree.i6), ets: :ets.info(tree.ets, :memory)}
+  end
 
-  """
-  def hello do
-    :world
+  @spec length(t()) :: %{inet4: non_neg_integer(), inet6: non_neg_integer(), ets: non_neg_integer()}
+  def length(tree) do
+    %{inet4: NIF.length(tree.i4), inet6: NIF.length(tree.i6), ets: :ets.info(tree.ets, :size)}
   end
+
+  defp add(tree, tbm, ip, masklen, value) do
+    id = :ets.update_counter(tree.ets, {__MODULE__, :counter}, 1, {{__MODULE__, :counter}, -1})
+    :ets.insert(tree.ets, {id, value})
+    {:ok, prev_id} = NIF.add(tbm, ip, masklen, id)
+    prev = if prev_id do
+      [{^prev_id, value}] = :ets.lookup(tree.ets, prev_id)
+      :ets.delete(tree.ets, prev_id)
+      value
+    end
+    prev
+  end
+
+  defp remove(tree, tbm, ip, masklen) do
+    {:ok, id} = NIF.remove(tbm, ip, masklen)
+    prev = if id do
+      [{^id, value}] = :ets.lookup(tree.ets, id)
+      :ets.delete(tree.ets, id)
+      value
+    end
+    prev
+  end
+
+  defp longest_match(tree, tbm, ip) do
+    case NIF.longest_match(tbm, ip) do
+      {:ok, prefix, masklen, id} ->
+        [{^id, value}] = :ets.lookup(tree.ets, id)
+        %{prefix: to_inet(prefix), len: masklen, value: value}
+      {:ok, nil} ->
+        nil
+    end
+  end
+
+  defp longest_match?(_, tbm, ip) do
+	  case NIF.longest_match(tbm, ip) do
+      {:ok, nil} -> false
+      {:ok, _, _, _} -> true
+    end
+  end
+
+  defp exact_match(tree, tbm, ip, masklen) do
+    case NIF.exact_match(tbm, ip, masklen) do
+      {:ok, nil} ->
+        nil
+      {:ok, id} ->
+        [{^id, value}] = :ets.lookup(tree.ets, id)
+        value
+    end
+  end
+
+  defp exact_match?(_, tbm, ip, masklen) do
+	  case NIF.exact_match(tbm, ip, masklen) do
+      {:ok, nil} -> false
+      {:ok, _} -> true
+    end
+  end
+
+  defp to_inet({:inet4, a, b, c, d}), do: {a, b, c, d}
+  defp to_inet({:inet6, a, b, c, d, e, f, g, h}), do: {a, b, c, d, e, f, g, h}
+
 end