Module:FunList/Lookup

Documentation for this module may be created at Module:FunList/Lookup/doc

local Enumerable = require('Module:Enumerable')
local TableEnumerator = require('Module:TableEnumerator')

---@class Grouping : Enumerable
---@field key any
---@field elements table
---@field count integer
local Grouping = setmetatable({}, { __index = Enumerable })
Grouping.__index = Grouping

---@return Grouping
function Grouping.new(key)
    assert(key, 'key may not be nil')
	local self = setmetatable({}, Grouping)
    self.key = key
	self.elements = {}
    self.count = 0

    return self
end

function Grouping:add(item)
    assert(item)
    table.insert(self.elements, item)
    self.count = self.count + 1
end

function Grouping:getElement(index)
    assert(tonumber(index))
    if index > self.count then
        error('Index exceeds elements in grouping.')
    end
    return self.elements[index]
end

---@return Enumerator
function Grouping:getEnumerator(isArray)
    return TableEnumerator.createForArray(self.elements)
end

---@class Lookup : Enumerable
---@field _set table<integer, Grouping>
---@field count integer
local Lookup = setmetatable({}, { __index = Enumerable })
Lookup.__index = Lookup

---@param source any
---@return Lookup
function Lookup.new(source, keySelector, elementSelector, isArray)
    assert(source, 'source may not be nil.')
    assert(keySelector, 'keySelector may not be nil.')

    local self = setmetatable({}, Lookup)
    self.count = 0
    self._set = {}
    elementSelector = elementSelector or function(x) return x end

    local enumerator = nil
    if Enumerable.isEnumerable(source) == true then
        enumerator = source:getEnumerator(isArray)
    else
        enumerator = TableEnumerator.create(source, isArray)
    end

    -- Add data to the Lookup
    while enumerator:moveNext() == true do
        local current = enumerator.current
        local index = enumerator.index
        local key = keySelector(current, index)
        local grouping = self:getGrouping(key)

        -- Add new grouping if this key doesn't exist.
        if grouping == nil then
            grouping = Grouping.new(key)
            self._set[key] = grouping
        end

        grouping:add(elementSelector(current, index))
        self.count = self.count + 1
    end

    enumerator:finalize()
    return self
end

---@param key any
---@return Grouping
function Lookup:getGrouping(key)
    assert(key, 'Key may not be nil.')

    return self._set[key]
end

---@return Enumerator
function Lookup:getEnumerator()
    return TableEnumerator.createForHash(self._set)
end

return Lookup