Module:Enum

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

-- <nowiki> awawa
local libraryUtil = require('libraryUtil')
local checkType = libraryUtil.checkType
local checkTypeMulti = libraryUtil.checkTypeMulti
local p = {}
p.__index = p

setmetatable(p, {
	__call = function(_, enum)
		return p.new(enum)
	end
})

function p.__tostring(enum)
	local dumpObject = require('Module:Logger').dumpObject
	setmetatable(enum, nil)
	return dumpObject(enum, {clean=true, collapseLimit=100})
end

function p.__concat(lhs, rhs)
	if type(lhs) == 'table' and type(rhs) == 'table' then
		return p.insert(lhs, rhs)
	else
		return tostring(lhs) .. tostring(rhs)
	end
end

function p.__unm(enum)
	return p.map(enum, function(x) return -x end)
end

local function mathTemplate(lhs, rhs, funName, fun)
	checkTypeMulti('Module:Enum.' .. funName, 1, lhs, {'number', 'table'})
	checkTypeMulti('Module:Enum.' .. funName, 2, rhs, {'number', 'table'})
	local res = setmetatable({}, getmetatable(lhs) or getmetatable(rhs))

	if type(lhs) == 'number' or type(rhs) == 'number' then
		local scalar, vector
		if type(lhs) == 'number' then
			scalar = lhs
			vector = rhs
		else
			scalar = rhs
			vector = lhs
		end

		for i = 1, #vector do
			res[i] = fun(vector[i], scalar)
		end
	else
		assert(#lhs == #rhs, 'Tables are not equal length')
		for i = 1, #lhs do
			res[i] = fun(lhs[i], rhs[i])
		end
	end

	return res
end

function p.__add(lhs, rhs)
	return mathTemplate(lhs, rhs, '__add', function(x, y) return x + y end)
end

function p.__sub(lhs, rhs)
	return mathTemplate(lhs, rhs, '__sub', function(x, y) return x - y end)
end

function p.__mul(lhs, rhs)
	return mathTemplate(lhs, rhs, '__mul', function(x, y) return x * y end)
end

function p.__div(lhs, rhs)
	return mathTemplate(lhs, rhs, '__div', function(x, y) return x / y end)
end

function p.__pow(lhs, rhs)
	return mathTemplate(lhs, rhs, '__pow', function(x, y) return x ^ y end)
end

function p.__lt(lhs, rhs)
	for i = 1, math.min(#lhs, #rhs) do
		if lhs[i] >= rhs[i] then
			return false
		end
	end
	return true
end

function p.__le(lhs, rhs)
	for i = 1, math.min(#lhs, #rhs) do
		if lhs[i] > rhs[i] then
			return false
		end
	end
	return true
end

function p.__eq(lhs, rhs)
	if #lhs ~= #rhs then
		return false
	end
	for i = 1, #lhs do
		if lhs[i] ~= rhs[i] then
			return false
		end
	end
	return true
end

function p.all(enum, fn, clone)
	checkType('Module:Enum.all', 1, enum, 'table')
	checkType('Module:Enum.all', 2, fn, 'function', true)
	checkType('Module:Enum.all', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	fn = fn or function(item) return item end
	local i = 1
	while enum[i] ~= nil do
		if not fn(enum[i], i) then
			return false
		end
		i = i + 1
	end
	return true
end

function p.any(enum, fn, clone)
	checkType('Module:Enum.any', 1, enum, 'table')
	checkType('Module:Enum.any', 2, fn, 'function', true)
	checkType('Module:Enum.any', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	fn = fn or function(item) return item end
	local i = 1
	while enum[i] ~= nil do
		if fn(enum[i], i) then
			return true
		end
		i = i + 1
	end
	return false
end

function p.contains(enum, elem, clone)
	checkType('Module:Enum.contains', 1, enum, 'table')
	checkType('Module:Enum.contains', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum); elem = mw.clone(elem) end
	return p.any(enum, function(item) return item == elem end)
end

function p.each(enum, fn, clone)
	checkType('Module:Enum.each', 1, enum, 'table')
	checkType('Module:Enum.each', 2, fn, 'function')
	checkType('Module:Enum.each', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	local i = 1
	while enum[i] ~= nil do
		fn(enum[i], i)
		i = i + 1
	end
end

function p.filter(enum, fn, clone)
	checkType('Module:Enum.filter', 1, enum, 'table')
	checkType('Module:Enum.filter', 2, fn, 'function', true)
	checkType('Module:Enum.filter', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	fn = fn or function(item) return item end
	local r = setmetatable({}, getmetatable(enum))
	local len = 0
	local i = 1
	while enum[i] ~= nil do
		if fn(enum[i], i) then
			len = len + 1
			r[len] = enum[i]
		end
		i = i + 1
	end
	return r
end

function p.find(enum, fn, default, clone)
	checkType('Module:Enum.find', 1, enum, 'table')
	checkType('Module:Enum.find', 2, fn, 'function')
	checkType('Module:Enum.find', 4, clone, 'boolean', true)
	if clone then enum = mw.clone(enum); default = mw.clone(default) end
	local i = 1
	while enum[i] ~= nil do
		if fn(enum[i], i) then
			return enum[i], i
		end
		i = i + 1
	end
	return default
end

function p.find_index(enum, fn, default, clone)
	checkType('Module:Enum.find_index', 1, enum, 'table')
	checkType('Module:Enum.find_index', 2, fn, 'function')
	checkType('Module:Enum.find_index', 4, clone, 'boolean', true)
	if clone then enum = mw.clone(enum); default = mw.clone(default) end
	fn = fn or function(item) return item end
	local i = 1
	while enum[i] ~= nil do
		if fn(enum[i], i) then
			return i
		end
		i = i + 1
	end
	return default
end

function p.newIncrementor(start, step)
	checkType('Module:Enum.newIncrementor', 1, start, 'number', true)
	checkType('Module:Enum.newIncrementor', 2, step, 'number', true)
	step = step or 1
	local n = (start or 1) - step
	local obj = {}
	return setmetatable(obj, {
		__call = function() n = n + step return n end,
		__tostring = function() return n end,
		__index = function() return n end,
		__newindex = function(self, k, v)
			if k == 'step' and type(v) == 'number' then
				step = v
			elseif type(v) == 'number' then
				n = v
			end
		end,
		__concat = function(x, y) return tostring(x) .. tostring(y) end
	})
end

function p.intersect(enum1, enum2, clone)
	checkType('Module:Enum.intersect', 1, enum1, 'table')
	checkType('Module:Enum.intersect', 2, enum2, 'table')
	checkType('Module:Enum.intersect', 3, clone, 'boolean', true)
	if clone then enum1 = mw.clone(enum1); enum2 = mw.clone(enum2) end
	local enum2Elements = {}
	local res = setmetatable({}, getmetatable(enum1) or getmetatable(enum2))
	local len = 0
	p.each(enum2, function(item) enum2Elements[item] = true end)
	p.each(enum1, function(item)
		if enum2Elements[item] then
			len = len + 1
			res[len] = item
		end
	end)
	return res
end

function p.intersects(enum1, enum2, clone)
	checkType('Module:Enum.intersects', 1, enum1, 'table')
	checkType('Module:Enum.intersects', 2, enum2, 'table')
	checkType('Module:Enum.intersects', 3, clone, 'boolean', true)
	if clone then enum1 = mw.clone(enum1); enum2 = mw.clone(enum2) end
	local small = {}
	local large
	if #enum1 <= #enum2 then
		p.each(enum1, function(item) small[item] = true end)
		large = enum2
	else
		p.each(enum2, function(item) small[item] = true end)
		large = enum1
	end
	return p.any(large, function(item) return small[item] end)
end

function p.insert(enum1, enum2, index, clone)
	checkType('Module:Enum.insert', 1, enum1, 'table')
	checkType('Module:Enum.insert', 2, enum2, 'table')
	checkType('Module:Enum.insert', 3, index, 'number', true)
	checkType('Module:Enum.insert', 4, clone, 'boolean', true)
	if clone then enum1 = mw.clone(enum1); enum2 = mw.clone(enum2) end
	local len1 = #enum1
	local len2 = #enum2
	index = index or (len1 + 1)
	local res = setmetatable({}, getmetatable(enum1) or getmetatable(enum2))

	for i = 1, (len1 + len2) do
		if i < index then
			res[i] = enum1[i]
		elseif i < (index + len2) then
			res[i] = enum2[i - index + 1]
		else
			res[i] = enum1[i - len2]
		end
	end

	return res
end

function p.map(enum, fn, clone)
	checkType('Module:Enum.map', 1, enum, 'table')
	checkType('Module:Enum.map', 2, fn, 'function')
	checkType('Module:Enum.map', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	local len = 0
	local r = setmetatable({}, getmetatable(enum))
	local i = 1
	while enum[i] ~= nil do
		local tmp = fn(enum[i], i)
		if tmp ~= nil then
			len = len + 1
			r[len] = tmp
		end
		i = i + 1
	end
	return r
end

function p.max_by(enum, fn, clone)
	checkType('Module:Enum.max_by', 1, enum, 'table')
	checkType('Module:Enum.max_by', 2, fn, 'function')
	checkType('Module:Enum.max_by', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	return unpack(p.reduce(enum, function(new, old)
		local y = fn(new)
		return y > old[2] and {new, y} or old
	end, {enum[1], -math.huge}))
end

function p.new(enum)
	for _, v in pairs(enum) do
		if type(v) == 'table' then
			p.new(v)
		end
	end

	if getmetatable(enum) == nil then
		setmetatable(enum, p)
	end

	return enum
end

function p.range(start, stop, step)
	checkType('Module:Enum.range', 1, start, 'number')
	checkType('Module:Enum.range', 2, stop, 'number', true)
	checkType('Module:Enum.range', 3, step, 'number', true)
	local array = setmetatable({}, p)
	local len = 0
	if not stop then
		stop = start
		start = 1
	end
	for i = start, stop, step or 1 do
		len = len + 1
		array[len] = i
	end
	return array
end

function p.reduce(enum, fn, accumulator, clone)
	checkType('Module:Enum.reduce', 1, enum, 'table')
	checkType('Module:Enum.reduce', 2, fn, 'function')
	checkType('Module:Enum.reduce', 4, clone, 'boolean', true)
	if clone then enum = mw.clone(enum); accumulator = mw.clone(accumulator) end
	local acc = accumulator
	local i = 1
	while enum[i] ~= nil do
		if i == 1 and not accumulator then
			acc = enum[i]
		else
			acc = fn(enum[i], acc)
		end
		i = i + 1
	end
	return acc
end

function p.reject(enum, fn, clone)
	checkType('Module:Enum.reject', 1, enum, 'table')
	checkTypeMulti('Module:Enum.reject', 2, fn, {'function', 'table', 'nil'})
	checkType('Module:Enum.reject', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	fn = fn or function(item) return item end
	local r = setmetatable({}, getmetatable(enum))
	local len = 0
	if type(fn) == 'function' then
		local i = 1
		while enum[i] ~= nil do
			if not fn(enum[i], i) then
				len = len + 1
				r[len] = enum[i]
			end
			i = i + 1
		end
	else
		local rejectMap = {}
		p.each(fn, function(item) rejectMap[item] = true end)
		local i = 1
		while enum[i] ~= nil do
			if not rejectMap[enum[i]] then
				len = len + 1
				r[len] = enum[i]
			end
			i = i + 1
		end
	end
	return r
end

function p.rep(val, n, clone)
	checkType('Module:Enum.rep', 2, n, 'number')
	checkType('Module:Enum.reject', 3, clone, 'boolean', true)
	if clone then val = mw.clone(val) end
	local r = setmetatable({}, p)
	for i = 1, n do
		r[i] = val
	end
	return r
end

function p.scan(enum, fn, accumulator, clone)
	checkType('Module:Enum.scan', 1, enum, 'table')
	checkType('Module:Enum.scan', 2, fn, 'function')
	checkType('Module:Enum.scan', 4, clone, 'boolean', true)
	if clone then enum = mw.clone(enum); accumulator = mw.clone(accumulator) end
	local acc = accumulator
	local r = setmetatable({}, getmetatable(enum))
	local i = 1
	while enum[i] ~= nil do
		if i == 1 and not accumulator then
			acc = enum[i]
		else
			acc = fn(enum[i], acc)
		end
		r[i] = acc
		i = i + 1
	end
	return r
end

function p.slice(enum, start, finish, clone)
	checkType('Module:Enum.slice', 1, enum, 'table')
	checkType('Module:Enum.slice', 2, start, 'number', true)
	checkType('Module:Enum.slice', 3, finish, 'number', true)
	checkType('Module:Enum.slice', 4, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	start = start or 1
	finish = finish or #enum
	local r = setmetatable({}, getmetatable(enum))
	local len = 0
	for i = start, finish do
		len = len + 1
		r[len] = enum[i]
	end
	return r
end

function p.split(enum, count, clone)
	checkType('Module:Enum.split', 1, enum, 'table')
	checkType('Module:Enum.split', 2, count, 'number')
	checkType('Module:Enum.split', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	if #enum < count then
		return enum, {}
	elseif count < 1 then
		return {}, enum
	end

	local x = setmetatable({}, getmetatable(enum))
	local y = setmetatable({}, getmetatable(enum))

	for i = 1, #enum do
		table.insert(i <= count and x or y, enum[i])
	end
	return x, y
end

function p.sum(enum, clone)
	checkType('Module:Enum.sum', 1, enum, 'table')
	checkType('Module:Enum.sum', 2, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	return p.reduce(enum, function(x, y) return x + y end)
end

function p.take(enum, count, clone)
	checkType('Module:Enum.take', 1, enum, 'table')
	checkType('Module:Enum.take', 2, count, 'number')
	checkType('Module:Enum.take', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	local x, _ = p.split(enum, count)
	return x
end

function p.take_every(enum, n, clone)
	checkType('Module:Enum.take_every', 1, enum, 'table')
	checkType('Module:Enum.take_every', 2, n, 'number')
	checkType('Module:Enum.take_every', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	local r = setmetatable({}, getmetatable(enum))
	local len = 0
	local i = 1
	while enum[i] ~= nil do
		if (i - 1) % n == 0 then
			len = len + 1
			r[len] = enum[i]
		end
		i = i + 1
	end
	return r
end

function p.unique(enum, fn, clone)
	checkType('Module:Enum.unique', 1, enum, 'table')
	checkType('Module:Enum.unique', 2, fn, 'function', true)
	checkType('Module:Enum.unique', 3, clone, 'boolean', true)
	if clone then enum = mw.clone(enum) end
	fn = fn or function(item) return item end
	local r = setmetatable({}, getmetatable(enum))
	local len = 0
	local hash = {}
	local i = 1
	while enum[i] ~= nil do
		local id = fn(enum[i])
		if not hash[id] then
			len = len + 1
			r[len] = enum[i]
			hash[id] = true
		end
		i = i + 1
	end
	return r
end

function p.zip(enums, clone)
	checkType('Module:Enum.zip', 1, enums, 'table')
	checkType('Module:Enum.zip', 2, clone, 'boolean', true)
	if clone then enums = mw.clone(enums) end
	local r = setmetatable({}, getmetatable(enums))
	local _, longest = p.max_by(enums, function(enum) return #enum end)
	for i = 1, longest do
		local q = {}
		for j = 1, #enums do
			table.insert(q, enums[j][i])
		end
		table.insert(r, q)
	end
	return r
end

return p
-- </nowiki>