aboutsummaryrefslogtreecommitdiff
path: root/lib/tree_bitmap.ex
blob: 0da87fa50ecea79e6080b504dfeaf33fc5231d10 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
defmodule TreeBitmap do
  alias TreeBitmap.NIF
  defstruct [:i4, :i6, :ets]

  @opaque t() :: %__MODULE__{}
  @type masklen :: non_neg_integer()

  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

  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

  @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