Module:memoize

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

local math_module = "Module:math"
local table_pack_module = "Module:table/pack"

local require = require
local select = select
local unpack = unpack or table.unpack -- Lua 5.2 compatibility

-- table.pack: in Lua 5.2+, this is a function that wraps the parameters given
-- into a table with the additional key `n` that contains the total number of
-- parameters given. This is not available on Lua 5.1, so [[Module:table/pack]]
-- provides the same functionality.
local function pack(...)
	pack = require(table_pack_module)
	return pack(...)
end

local function sign(...)
	sign = require(math_module).sign
	return sign(...)
end

----- 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(x)
	if x == x then
		return x == nil and _nil or x == 0 and 1 / x < 0 and neg_0 or x
	end
	return sign(x) == 1 and pos_nan or 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

-- Used to catch the function output values instead of using a table directly,
-- since pack() returns a table with the key `n`, giving the number of return
-- values, even if they are nil. This ensures that any nil return values after
-- the last non-nil value will always be present (e.g. pack() gives {n = 0},
-- pack(nil) gives {n = 1}, pack(nil, "foo", nil) gives {[2] = "foo", n = 3}
-- 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 memoize_then_return(memo, _memo, ...)
	_memo[memo] = pack(...)
	return ...
end

return function(func, simple)
	local memo = {}

	if simple then

		return function(...)
			local key = get_key((...))
			local output = memo[key]
			if output == nil then
				output = func(...)
				memo[key] = output == nil and _nil or output
				return output
			elseif output == _nil then
				return nil
			end
			return output
		end

	end

	return function(...)
		local nargs = select("#", ...)
		-- Since all possible inputs need to be memoized (including true, false
		-- and nil), the memo table itself is used as a sentinel to ensure that
		-- the table of arguments will always have a unique key.
		local _memo = nargs == 0 and memo or get_memo(memo, 1, nargs, ...)
		local output = _memo[memo]
		-- If get_memo() returned nil, call `func` with the arguments and catch
		-- the output with memoize_then_return(); this packs the return values
		-- into a table to memoize them, then returns them. Since the return
		-- values are available to it as `...`, this avoids the need to call
		-- unpack() on the memoized table on the first call, as they can be
		-- returned directly.
		if output == nil then
			return memoize_then_return(memo, _memo, func(...))
		end
		-- Unpack from 1 to the original number of return values (memoized at
		-- key `n`); unpack() returns nil for any values not in output.
		return unpack(output, 1, output.n)
	end
end