set: basic support for looping over sets

Signed-off-by: Asbjørn Sloth Tønnesen <ast@fiberby.net>

1 files changed, 56 insertions, 0 deletions

diff --git a/lua/inet/set.lua b/lua/inet/set.lua
index 27c9388..5999f65 100644
--- a/lua/inet/set.lua
+++ b/lua/inet/set.lua
@@ -52,4 +52,60 @@ function M.aggregate(t)
 	table_compact(t, n)
 end
 
+local function has(set, addr)
+	assert(set)
+	for i=1,#set do
+		local elem = set[i]
+		if elem >= addr then
+			local exclude = set.exclude
+			if exclude then
+				return not has(exclude, addr)
+			else
+				return true
+			end
+		end
+	end
+	return false
+end
+
+function M.iterator(set)
+	local excl = set.exclude
+	local i = 1
+	if #set < 1 then return nil, 'empty set' end
+	local addr = set[i]
+	local net = set[i]:network()
+	local function iter()
+		if not addr then return end
+		local ret
+		ret = addr/32
+		addr = addr + 1
+		if addr:network() ~= net then
+			i = i + 1
+			addr = set[i]
+			if addr then
+				net = set[i]:network()
+			end
+		end
+		if has(excl, ret) then
+			return iter()
+		end
+		return ret
+	end
+
+	return iter
+end
+
+function M.loopiterator(set)
+	local orig_iter = M.iterator(set)
+	local function iter()
+		local addr = orig_iter()
+		if not addr then
+			orig_iter = M.iterator(set)
+			addr = orig_iter()
+		end
+		return addr
+	end
+	return iter
+end
+
 return M