Module:StringSet

From Star of Providence Wiki
Jump to navigation Jump to search

Details

The StringSet module offers a Lua class that makes it easier to work with sets of strings. It's intended to be used by other Lua modules instead of templates.

It has no dependencies and exports a StringSet "class" you can use directly. It provides annotations documenting the various functions that are understood by Lua Language Server.


--- Module:StringSet
--- by User:LtDk, licensed CC-BY-SA 4.0
--- Simple Lua class for parsing sets of strings.

local StringSet = {}

--- Set of strings.
--- @class StringSet
local StringSetProto = {}
local StringSetMeta = { __index = StringSetProto }

--- Checks if a value is a StringSet.
--- @param value any
--- @return boolean
function StringSet.isStringSet(value)
    return type(value) == 'table' and rawequal(getmetatable(value), StringSetMeta)
end

--- Iterates over the keys in the set.
--- @return fun(keys: { [string]: true }, string?): string?, string?
--- @return { [string]: true }
--- @return nil
function StringSetMeta:__pairs()
    return function (keys, key)
        key = next(keys, key)
        return key, key
    end, rawget(self, '__keys'), nil
end

--- Iterates over the keys in the set.
--- @return fun(keys: { [string]: true }, string?): string?, string?
--- @return { [string]: true }
--- @return nil
function StringSetMeta:__ipairs()
    return pairs(self)
end

--- Gets the number of keys in the set in constant time.
--- @return integer
function StringSetMeta:__len()
    return rawget(self, '__len')
end

--- Gets the number of keys in the set in constant time.
--- @return integer
function StringSetProto:len()
    return getmetatable(self).__len(self)
end

--- Iterates over the keys in the set.
--- @return fun(keys: { [string]: true }, string?): string?
--- @return { [string]: true }
--- @return nil
function StringSetProto:keys()
    return pairs(self)
end

--- Returns a sorted array of the keys in the set.
--- @return string[]
function StringSetProto:sortedKeys()
    local keys = {}
    for key in self:keys() do
        table.insert(keys, key)
    end
    table.sort(keys)
    return keys
end

--- Adds a key to the set, returning whether it was actually added.
--- @param key string
--- @return boolean
function StringSetProto:add(key)
    key = tostring(key)
    local keys = rawget(self, '__keys')
    local ret = not keys[key]
    keys[key] = true
    if ret then
        rawset(self, '__len', self:len() + 1)
    end
    return ret
end

--- Adds a set of keys to a set.
--- @param set StringSet|string[]
function StringSetProto:addAll(set)
    if set ~= nil then
        for _, key in ipairs(set) do
            self:add(key)
        end
    end
end

--- Removes a key to the set, returning whether it was actually removed.
--- @param key string
--- @return boolean
function StringSetProto:remove(key)
    key = tostring(key)
    local keys = rawget(self, '__keys')
    local ret = keys[key]
    keys[key] = nil
    if ret then
        rawset(self, '__len', self:len() - 1)
    end
    return ret
end

--- Removes a set of keys from a set.
--- @param set StringSet|string[]
function StringSetProto:removeAll(set)
    if set ~= nil then
        for _, key in ipairs(set) do
            self:remove(key)
        end
    end
end

--- Checks if a key is in the set.
--- @param key string
--- @return boolean
function StringSetProto:has(key)
    key = tostring(key)
    local keys = rawget(self, '__keys')
    return keys[key]
end

--- Subsets are always ordered before sets, but this is not a total ordering.
function StringSetMeta.__le(lhs, rhs)
    if not StringSet.isStringSet(rhs) then
        error('cannot compare StringSet with non-StringSet')
    end
    for key in lhs:keys() do
        if not rhs:has(key) then
            return false
        end
    end
    return true
end

--- Subsets are always ordered before sets, but this is not a total ordering.
function StringSetMeta.__lt(lhs, rhs)
    return lhs:len() ~= rhs:len() and lhs <= rhs
end

function StringSetMeta.__eq(lhs, rhs)
    if not StringSet.isStringSet(rhs) then
        return false
    end
    return lhs:len() == rhs:len() and lhs <= rhs and rhs <= lhs
end

function StringSetMeta:__newindex(k, v)
    error('read-only; use methods instead')
end

--- The string representation of a set always has its keys sorted, which can be expensive.
--- This is mostly intended for debugging.
function StringSetMeta:__tostring()
    local ret = '{'
    local first = true
    for _, key in ipairs(self:sortedKeys()) do
        if first then
            first = false
        else
            ret = ret .. ', '
        end
        ret = ret .. string.format('%q', key)
    end
    return ret .. '}'
end

--- Creates a new StringSet, potentially from another set.
--- @param set (StringSet|string[])?
--- @return StringSet
function StringSet.new(set)
    local ss = setmetatable({ __keys = {}, __len = 0 }, StringSetMeta)
    ss:addAll(set)
    return ss
end

return StringSet