adventofcode-2021/set.lua

local function add(self, value)
	if not self.values[value] then
		self.values[value] = true
		self.size = self.size + 1
	end
end

local function has(self, value)
	return self.values[value] ~= nil
end

local function remove(self, value)
	if self.values[value] then
		self.values[value] = nil
		self.size = self.size - 1
	end
end

local function iter(self)
	local index, item
	return function()
		repeat
			index, item = next(self.values, index)
		until index == nil or item ~= nil
		if index == nil then
			return nil
		else
			return index
		end
	end
end

local function keeponly(self, iterable)
	local tokeep = set(iterable)
	for v in self:iter() do
		if not tokeep:has(v) then
			self:remove(v)
		end
	end
end

local function removeall(self, iterable)
	for v in iterable do
		self:remove(v)
	end
end

local function addall(self, iterable)
	for v in iterable do
		self:add(v)
	end
end

local function equal(self, other)
	for c in self:iter() do if not other:has(c) then return false end end
	for c in other:iter() do if not self:has(c) then return false end end
	return true
end

local function write(self)
	io.write('{')
	for v in self:iter() do
		io.write(' '..v)
	end
	io.write(' }\n')
end

function set(t)
	local values = {}
	local size = 0
	if type(t) == "table" then
		for k, v in pairs(t) do
			size = size + 1
			values[v] = true
		end
	elseif type(t) == "function" then
		for v in t do
			size = size + 1
			values[v] = true
		end
	end
	return {
		values = values,
		size = size,
		add = add,
		has = has,
		remove = remove,
		iter = iter,
		keeponly = keeponly,
		removeall = removeall,
		addall = addall,
		equal = equal,
		write = write,
	}
end