initial new set support

Signed-off-by: Asbjørn Sloth Tønnesen <ast@2e8.dk>

3 files changed, 142 insertions, 6 deletions

diff --git a/README.rst b/README.rst
index 2083b83..f243fab 100644
--- a/README.rst
+++ b/README.rst
@@ -31,12 +31,12 @@ API
 API                     Description
 ======================= =====================================================
 ``inet(...)``           Parse address and build ``inet4`` or ``inet6`` table
-``inet.is(foo)``        is ``foo`` an ``inet*`` table
-``inet.is4(foo)``       is ``foo`` an ``inet4`` table
-``inet.is6(foo)``       is ``foo`` an ``inet6`` table
-``inet.set4()``         get new empty IPv4 set
-``inet.set6()``         get new empty IPv6 set
-``inet.mixed_networks`` IPv6 mixed notation set
+``inet.is(foo)``        is ``foo`` an ``inet*`` table?
+``inet.is4(foo)``       is ``foo`` an ``inet4`` table?
+``inet.is6(foo)``       is ``foo`` an ``inet6`` table?
+``inet.is_set(foo)``    is ``set`` table?
+``inet.set()``          get new empty ``set`` instance.
+``inet.mixed_networks`` IPv6 mixed notation ``set``
 ``inet.version``        API version (currently ``1``)
 ======================= =====================================================
 
@@ -324,31 +324,75 @@ Like ``tostring(foo)``, but never uses mixed notation.
 Sets
 ----
 
+::
+
+  local foo = inet.set()
+
 ``set:list()``
 ~~~~~~~~~~~~~~
 
 list networks in set
 
+::
+
+  foo:list() -- returns {}
+
 ``set:add(foo)``
 ~~~~~~~~~~~~~~~~
 
 add network to set
 
+::
+
+  foo:add(inet('2001:db8::/48')) -- returns true
+  foo:list() -- returns { inet('2001:db8::/48') }
+  foo:add(inet('2001:db8:1::/48')) -- returns true
+  foo:list() -- returns { inet('2001:db8::/47') }
+  foo:add(inet('192.0.2.0/24')) -- returns nil, 'invalid family'
+
 ``set:remove(foo)``
 ~~~~~~~~~~~~~~~~~~~
 
 remove network from set
 
+::
+
+  foo:remove(inet('2001:db8:1::/48')) -- returns true
+  foo:remove(inet('2001:db8:1::/48')) -- returns false
+  foo:list() -- returns { inet('2001:db8::/48') }
+
+  foo:remove(inet('2001:db8:0:4200::/56')) -- returns true
+  foo:list() -- returns {
+    inet('2001:db8::/50'),
+    inet('2001:db8:0:4000::/55'),
+    inet('2001:db8:0:4300::/56'),
+    inet('2001:db8:0:4400::/54'),
+    inet('2001:db8:0:4800::/53'),
+    inet('2001:db8:0:5000::/52'),
+    inet('2001:db8:0:6000::/51'),
+    inet('2001:db8:0:8000::/49'),
+  }
+
 ``set:contains(foo)``
 ~~~~~~~~~~~~~~~~~~~~~
 
 is network contained in set?
 
+::
+
+  foo:contains(inet('2001:db8::'))           -- returns true
+  foo:contains(inet('2001:db8::/32'))        -- returns false
+  foo:contains(inet('2001:db8:1:2:3:4:5:6')) -- returns false
+
 ``set:flush()``
 ~~~~~~~~~~~~~~~
 
 Empties the set.
 
+::
+  foo:flush() -- returns true
+  foo:list() -- returns {}
+
 History
 =======
 
diff --git a/lua/inet/init.lua b/lua/inet/init.lua
index 55ce0e4..ea35f80 100644
--- a/lua/inet/init.lua
+++ b/lua/inet/init.lua
@@ -1,4 +1,5 @@
 local core = require 'inet.core'
+local set  = require 'inet.set'
 
 local new_inet = core.new_inet
 
@@ -13,4 +14,6 @@ M.is4 = core.is_inet4
 M.is6 = core.is_inet6
 M.is  = core.is_inet
 
+M.set = set.new
+
 return setmetatable(M, mt)
diff --git a/lua/inet/set.lua b/lua/inet/set.lua
index 5999f65..a353f30 100644
--- a/lua/inet/set.lua
+++ b/lua/inet/set.lua
@@ -1,3 +1,11 @@
+local core = require 'inet.core'
+
+local is_inet = core.is_inet
+
+local insert = table.insert
+local remove = table.remove
+local sort = table.sort
+
 local M = {}
 
 local function table_compact(t, n)
@@ -95,6 +103,7 @@ function M.iterator(set)
 	return iter
 end
 
+-- iterate over the set repeatedly
 function M.loopiterator(set)
 	local orig_iter = M.iterator(set)
 	local function iter()
@@ -108,4 +117,84 @@ function M.loopiterator(set)
 	return iter
 end
 
+local inet_set = {}
+inet_set.__index = inet_set
+
+local function check_net(self, net)
+	if not is_inet(net) then return nil, 'invalid net' end
+	local set_fam = self.family
+	local net_fam = net:family()
+	if set_fam then
+		if set_fam ~= net_fam then
+			return nil, 'invalid family'
+		end
+	else
+		self.family = net_fam
+	end
+	return true
+end
+
+function inet_set:add(net)--, aggregate)
+	if not check_net(self, net) then return check_net(self, net) end
+	insert(self.nets, net)
+	M.aggregate(self.nets)
+	return true
+end
+
+function inet_set:remove(net)
+	if not check_net(self, net) then return check_net(self, net) end
+	local nets = self.nets
+	local ret = false
+	for i=1,#nets do
+		local n = nets[i]
+		if n == net then
+			ret = true
+			remove(nets, i)
+		elseif n:contains(net) then
+			ret = true
+			remove(nets, i)
+			for j=#n+1,#net do
+				local sn = (net / j):network()
+				if sn:contains(net) then
+					sn = sn:flip()
+				end
+				self:add(sn)
+			end
+		end
+	end
+	return ret
+end
+
+function inet_set:contains(net)
+	if not check_net(self, net) then return check_net(self, net) end
+	local nets = self.nets
+	for i=1,#nets do
+		local n = nets[i]
+		if n == net or n:contains(net) then
+			return true
+		end
+	end
+	return false
+end
+
+function inet_set:list()
+	local nets = self.nets
+	sort(nets)
+	return nets
+end
+
+function inet_set:flush()
+	self.nets = {}
+	self.family = nil
+	return true
+end
+
+local function new_set()
+	return setmetatable({
+		nets = {},
+	}, inet_set)
+end
+
+M.new = new_set
+
 return M