Module:FunList/Iterators

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

-- Helper Functions --
local function isType(obj, class)
    local mt = getmetatable(obj)
    while mt do
        if mt == class then
            return true
        end
        mt = getmetatable(mt)
    end
    return false
end

-- BASE ENUMERATOR CLASS --
--enumerable = {}
local Enumerator = {}
local Enumerator_mt = {
	__index = Enumerator,
	__pairs = function(t) return t:overridePairs() end,
	__ipairs = function(t) return t:overridePairs(0) end
}

function Enumerator.new()
	local self = setmetatable({}, Enumerator_mt)
	self.current = nil
	self.index = nil
	-- Assume by default we are not dealing with a simple array
	self.isArray = false
	return self
end

function Enumerator:moveNext()
	error('Not implemented in base class.')
end

function Enumerator:getEnumerator()
	error('Not implemented in base class.')
end

-- Hooks the moveNext function into the Lua 'pairs' function
function Enumerator:overridePairs(startIndex)
	-- Get or create clean enumerator. This ensures the state is 0.
	local enum = self:getEnumerator()
	enum.current = nil
	enum.index = startIndex
	enum.isArray = startIndex == 0
	local function iterator(t, k)
		if enum:moveNext() == true then
			return enum.index, enum.current
		end
		return nil, nil
	end
	
	return iterator, enum, enum.index
end

-- TABLE ENUMERATOR CLASS --
-- This is essentially a wrapper for the table object, 
-- since it provides no state machine esque iteration out of the box
local TableEnumerator = setmetatable({}, { __index = Enumerator })
TableEnumerator.__index = TableEnumerator
TableEnumerator.__pairs =  Enumerator_mt.__pairs
TableEnumerator.__ipairs = Enumerator_mt.__ipairs

function TableEnumerator.new(tbl)
    local self = setmetatable(Enumerator.new(), TableEnumerator)
    self.tbl = tbl or {} -- Allow creation of empty enumerable
    self.state = 0

    return self
end

function TableEnumerator:moveNext()
    if self.state == 0 then
        self.state = 1
    end

    if self.isArray == true then
        -- Iterate using ipairs, starting from index 1
        self.index = self.index + 1
        self.current = self.tbl[self.index]
        return self.current ~= nil
    else
        -- Iterate using pairs
        local key = self.index
        key = next(self.tbl, key)
        self.index = key
        if key ~= nil then
            self.current = self.tbl[key]
            return true
        end
    end

    return false
end

-- startIndex is used to determine if the underlying table should be treated
-- as an array or as a mixed table. It is ignored in the other enumerators as
-- they just call moveNext on the enumerator instead.
function TableEnumerator:getEnumerator(isArray)
	local instance = nil
	if self.state == 0 then
		instance = self
	else
		instance = TableEnumerator.new(self.tbl)
	end
	instance.isArray = isArray
	return instance
end

-- SELECT ENUMERATOR --
local SelectEnumerator = setmetatable({}, { __index = Enumerator })
SelectEnumerator.__index = SelectEnumerator
SelectEnumerator.__pairs = Enumerator_mt.__pairs
SelectEnumerator.__ipairs = Enumerator_mt.__ipairs

function SelectEnumerator.new(source, selector)
    local self = setmetatable(Enumerator.new(), SelectEnumerator)
    self.state = 0
    self.source = source
    self.selector = selector
    self.enumerator = nil
    self.position = 0
    return self	
end

function SelectEnumerator:moveNext()
	if self.state == 0 then
		self.state = 1
		self.position = 0
		self.enumerator = self.source:getEnumerator(self.isArray)
	end
	
	if self.state == 1 then
		local enumerator = self.enumerator
		if enumerator:moveNext() == true then
			local sourceElement = enumerator.current
			self.index = enumerator.index -- Preserve index
			self.position = self.position + 1
			self.current = self.selector(sourceElement, self.position)
			assert(self.current, 'Selected value must be non-nil')
			return true
		end
	end
	return false
end

function SelectEnumerator:getEnumerator()
	if self.state == 0 then
		return self
	else
		return SelectEnumerator.new(self.source, self.selector)
	end
end

-- WHERE ENUMERATOR --
local WhereEnumerator = setmetatable({}, { __index = Enumerator })
WhereEnumerator.__index = WhereEnumerator
WhereEnumerator.__pairs = Enumerator_mt.__pairs
WhereEnumerator.__ipairs = Enumerator_mt.__ipairs

function WhereEnumerator.new(source, predicate)
    local self = setmetatable(Enumerator.new(), WhereEnumerator)
    self.state = 0
    self.source = source
    self.predicate = predicate
    self.enumerator = nil
    return self	
end

function WhereEnumerator:moveNext()
	if self.state == 0 then
		self.state = 1
		self.position = 0
		self.enumerator = self.source:getEnumerator(self.isArray)
	end
	
	if self.state == 1 then
		local enumerator = self.enumerator
		while enumerator:moveNext() == true do
			local sourceElement = enumerator.current
			if self.predicate(sourceElement) == true then
				self.index = enumerator.index
				self.current = sourceElement
				return true
			end
		end
	end
	
	return false
end

function WhereEnumerator:getEnumerator()
	if self.state == 0 then
		return self
	else
		return WhereEnumerator.new(self.source, self.predicate)
	end
end

-- WHERE ENUMERATOR --
local SelectManyEnumerator = setmetatable({}, { __index = Enumerator })
SelectManyEnumerator.__index = SelectManyEnumerator
SelectManyEnumerator.__pairs = Enumerator_mt.__pairs
SelectManyEnumerator.__ipairs = Enumerator_mt.__ipairs

function SelectManyEnumerator.new(source, selector)
    local self = setmetatable(Enumerator.new(), SelectManyEnumerator)
    self.state = 0
    self.source = source
    self.selector = selector
    self.position = 0
    self.enumerator = nil		 -- Enumerator of the source Enumerable
    self.sourceEnumerator = nil  -- Enumerator of the nested Enumerable
    return self	
end

function SelectManyEnumerator:moveNext()
    while true do
        -- Setup state
        if self.state == 0 then
            self.position = 0
            self.enumerator = self.source:getEnumerator(self.isArray)
            self.state = -3 -- signal to get (first) nested enumerator
        end

        -- Grab next value from nested enumerator		
        if self.state == -4 then
            if self.sourceEnumerator:moveNext() then
                self.current = self.sourceEnumerator.current
                self.index = self.sourceEnumerator.index
                self.state = -4 -- signal to get next item
                return true
            else
                self.state = -3 -- signal to get next enumerator
            end
        end
		
		-- Grab nest nested enumerator
        if self.state == -3 then
            if self.enumerator:moveNext() then
                local current = self.enumerator.current
                self.position = self.position + 1

                local sourceTable = self.selector(current, self.position)
                if not isType(sourceTable, Enumerator) then
                	-- We need to turn the nested table into an enumerator
                	self.sourceEnumerator = TableEnumerator.new(sourceTable)
                		:getEnumerator(self.isArray)
                else
                	self.sourceEnumerator = sourceTable
                end
                self.state = -4 -- signal to get next item
            else
            	-- enumerator doesn't have any more nested enumerators.
                return false
            end
        end
    end
end


function SelectManyEnumerator:getEnumerator()
	if self.state == 0 then
		return self
	else
		return SelectManyEnumerator.new(self.source, self.predicate)
	end
end

return {
    Enumerator = Enumerator,
    TableEnumerator = TableEnumerator,
    SelectEnumerator = SelectEnumerator,
    WhereEnumerator = WhereEnumerator,
    SelectManyEnumerator = SelectManyEnumerator
}