Math Set in Pure Lua (Apache 2.0)

set.lua

local Set = {}
package.loaded["Set"] = Set

Set.__index = Set

Set.__eq = function (a, b)
    for e in a:iter() do
        if not b:contains(e) then
            return false
        end
    end
    
    for e in b:iter() do
        if not a:contains(e) then
            return false
        end
    end
    
    return true
end

function Set:new()
    self = {}
    self._set = {}
    setmetatable(self, Set)
    return self
end

function Set:fromData(data)
    local s = self:new()
    
    for _, e in ipairs(data) do
        s:add(e)
    end
    
    return s
end

function Set:contains(e)
    return self._set[e] == true
end

function Set:add(e)
    self._set[e] = true
end

function Set:remove(e)
    if self._set[e] == true then
        self._set[e] = nil
    end
end

function Set:iter()
    local out = {}
    
    for k, _ in pairs(self._set) do
        table.insert(out, k)
    end
    
    local n = 0
    
    return function ()
        n = n + 1
        
        if n <= #out then
            return out[n]
        else
            return nil
        end
    end
end
        

Set.union = function (a, b)
    local out = Set:new()

    for e in a:iter() do
        out:add(e)
    end
    
    for e in b:iter() do
        out:add(e)
    end

    return out
end

Set.intersection = function (a, b)
    local out = Set:new()
    
    for e in a:iter() do
        if b:contains(e) then
            out:add(e)
        end
    end
    
    return out
end

Set.difference = function (a, b)
    local out = Set:new()
    
    for e in a:iter() do
        if not b:contains(e) then
            out:add(e)
        end
    end
    
    return out
end

return Set

set_test.lua

local set = require("set")

do
    local s = set:new()
    assert(s:contains("a") == false)
    
    s:add("a")
    assert(s:contains("a") == true)
    
    s:remove("a")
    assert(s:contains("a") == false)
    
    s:remove("b")
    assert(s:contains("b") == false)
end

do
    local a = set:fromData({"a", "b", "c"})
    local b = set:fromData({"c", "d", "e"})
    
    local s = set.union(a, b)
    
    assert(s == set:fromData({"a", "b", "c", "d", "e"}))
end

do
    local a = set:fromData({"a", "b", "c"})
    local b = set:fromData({"c", "d", "e"})
    
    local s = set.intersection(a, b)
    
    assert(s == set:fromData({"c"}))
end

do
    local a = set:fromData({"a", "b", "c"})
    local b = set:fromData({"c", "d", "e"})
    
    local s = set.difference(a, b)
    
    assert(s == set:fromData({"a", "b"}))
end