diff --git a/lua/spec/fn_util_spec.lua b/lua/spec/fn_util_spec.lua index b4ee99c2656..aeeaae764b3 100644 --- a/lua/spec/fn_util_spec.lua +++ b/lua/spec/fn_util_spec.lua @@ -26,6 +26,59 @@ describe('function utils', function() end) end) + describe('memoize2', function() + it('check', function() + local calledCount = 0 + local multiply = FnUtil.memoize2(function(x, y) + calledCount = calledCount + 1 + if x ~= nil and y ~= nil then + return x * y + end + if x == nil and y == nil then + return nil + end + return 0 + end) + + assert.are_equal(0, calledCount) + assert.are_equal(4, multiply(2, 2)) + assert.are_equal(1, calledCount) + assert.are_equal(4, multiply(2, 2)) + assert.are_equal(1, calledCount) + assert.are_equal(9, multiply(3, 3)) + assert.are_equal(2, calledCount) + assert.are_equal(9, multiply(3, 3)) + assert.are_equal(2, calledCount) + assert.is_nil(multiply(nil, nil)) + assert.is_nil(multiply(nil, nil)) + assert.is_nil(multiply(nil, nil)) + assert.are_equal(3, calledCount) + -- Mixed nil inputs are memoized independently + assert.are_equal(0, multiply(2, nil)) + assert.are_equal(4, calledCount) + assert.are_equal(0, multiply(2, nil)) + assert.are_equal(4, calledCount) + assert.are_equal(0, multiply(nil, 3)) + assert.are_equal(5, calledCount) + assert.are_equal(0, multiply(nil, 3)) + assert.are_equal(5, calledCount) + -- Argument order produces separate cache entries + assert.are_equal(6, multiply(2, 3)) + assert.are_equal(6, calledCount) + assert.are_equal(6, multiply(3, 2)) + assert.are_equal(7, calledCount) + assert.are_equal(6, multiply(2, 3)) + assert.are_equal(7, calledCount) + assert.are_equal(6, multiply(3, 2)) + assert.are_equal(7, calledCount) + -- Previously computed values remain cached after new computations + assert.are_equal(4, multiply(2, 2)) + assert.are_equal(7, calledCount) + assert.are_equal(9, multiply(3, 3)) + assert.are_equal(7, calledCount) + end) + end) + describe('memoize Y', function() it('check', function() local calledCount = 0 diff --git a/lua/wikis/commons/FnUtil.lua b/lua/wikis/commons/FnUtil.lua index c9853c5b969..321c334b731 100644 --- a/lua/wikis/commons/FnUtil.lua +++ b/lua/wikis/commons/FnUtil.lua @@ -11,6 +11,7 @@ local FnUtil = {} ---@alias memoizableFunction fun(input: any): any +---@alias memoizableFunction2 fun(input: any, input2: any): any ---Creates a memoized copy of a 0 or 1 param function. ---@generic T:memoizableFunction @@ -38,6 +39,34 @@ function FnUtil.memoize(func) end end +---Creates a memoized copy of a 2 param function. +---@generic T:memoizableFunction2 +---@param func T +---@return T +function FnUtil.memoize2(func) + local called = {} + local results = {} + local xNil = FnUtil.memoize(function(y) return func(nil, y) end) + local yNil = FnUtil.memoize(function(x) return func(x, nil) end) + return function(x, y) + if y == nil then + return yNil(x) + elseif x == nil then + return xNil(y) + else + if not called[x] then + called[x] = {} + results[x] = {} + end + if not called[x][y] then + called[x][y] = true + results[x][y] = func(x, y) + end + return results[x][y] + end + end +end + --[[ Memoized variant of the Y combinator. Useful for caching results of recursive functions, so that previously computed inputs are not recomputed.