Module:memoize

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

local format = string.format
local select = select
local unpack = unpack

----- M E M O I Z A T I O N-----
-- Memoizes a function or callable table.
-- Supports any number of arguments and return values.
-- If the optional parameter `simple` is set, then the memoizer will use a faster implementation, but this is only compatible with one argument and one return value. If `simple` is set, additional arguments will be accepted, but this should only be done if those arguments will always be the same.

-- Sentinels.
local nil_, neg_0, pos_nan, neg_nan

-- Certain values can't be used as table keys, so they require sentinels as well: e.g. f("foo", nil, "bar") would be memoized at memo["foo"][nil_]["bar"][memo]. These values are:
	-- nil.
	-- -0, which is equivalent to 0 in most situations, but becomes "-0" on conversion to string; it also behaves differently in some operations (e.g. 1/a evaluates to inf if a is 0, but -inf if a is -0).
	-- NaN and -NaN, which are the only values for which n == n is false; they only seem to differ on conversion to string ("nan" and "-nan").
local function get_key(input)
	-- nil
	if input == nil then
		if not nil_ then
			nil_ = {}
		end
		return nil_
	-- -0
	elseif input == 0 and 1 / input < 0 then
		if not neg_0 then
			neg_0 = {}
		end
		return neg_0
	-- Default
	elseif input == input then
		return input
	-- NaN
	elseif format("%f", input) == "nan" then
		if not pos_nan then
			pos_nan = {}
		end
		return pos_nan
	-- -NaN
	elseif not neg_nan then
		neg_nan = {}
	end
	return neg_nan
end

-- Return values are memoized as tables of return values, which are looked up using each input argument as a key, followed by `memo`. e.g. if the input arguments were (1, 2, 3), the memo would be located at t[1][2][3][memo]. `memo` is always used as the final lookup key so that (for example) the memo for f(1, 2, 3), f[1][2][3][memo], doesn't interfere with the memo for f(1, 2), f[1][2][memo].
local function get_memo(memo, n, nargs, key, ...)
	key = get_key(key)
	local next_memo = memo[key]
	if next_memo == nil then
		next_memo = {}
		memo[key] = next_memo
	end
	memo = next_memo
	return n == nargs and memo or get_memo(memo, n + 1, nargs, ...)
end

-- Catch the function output values, and return the hidden variable arg (which is {...}, and available when a function has ...). We do this instead of catching the output in a table directly, because arg also contains the key "n", which is equal to select("#", ...). i.e. it's the number of arguments in ..., including any nils returned after the last non-nil value (e.g. select("#", nil) == 1, select("#") == 0, select("#", nil, "foo", nil, nil) == 4 etc.). The distinction between nil and nothing affects some native functions (e.g. tostring() throws an error, but tostring(nil) returns "nil"), so it needs to be reconstructable from the memo.
local function catch_output(...)
	-- TODO uses arg; will not work if Scribunto is upgraded to Lua 5.2, 5.3, etc.
	return arg
end

return function(func, simple)
	local memo
	return simple and function(...)
		local key = get_key(...)
		if not memo then
			memo = {}
		end
		local output = memo[key]
		if output == nil then
			output = func(...)
			if output ~= nil then
				memo[key] = output
				return output
			elseif not nil_ then
				nil_ = {}
			end
			memo[key] = nil_
			return nil
		elseif output == nil_ then
			return nil
		end
		return output
	end or function(...)
		local nargs = select("#", ...)
		if not memo then
			memo = {}
		end
		-- Since all possible inputs need to be memoized (including true, false and nil), the memo table itself is used as the key for the arguments.
		local _memo = nargs == 0 and memo or get_memo(memo, 1, nargs, ...)
		local output = _memo[memo]
		if output == nil then
			output = catch_output(func(...))
			_memo[memo] = output
		end
		-- Unpack from 1 to the original number of return values (memoized as output.n); unpack returns nil for any values not in output.
		return unpack(output, 1, output.n)
	end
end