neovim/runtime/lua/vim/func/_memoize.lua
Lewis Russell 877d04d0fb feat(lua): add vim.func._memoize
Memoizes a function, using a custom function to hash the arguments.

Private for now until:

- There are other places in the codebase that could benefit from this
  (e.g. LSP), but might require other changes to accommodate.
- Invalidation of the cache needs to be controllable. Using weak tables
  is an acceptable invalidation policy, but it shouldn't be the only
  one.
- I don't think the story around `hash_fn` is completely thought out. We
  may be able to have a good default hash_fn by hashing each argument,
  so basically a better 'concat'.
2023-09-20 13:42:41 +01:00

60 lines
1.3 KiB
Lua

--- Module for private utility functions
--- @param argc integer?
--- @return fun(...): any
local function concat_hash(argc)
return function(...)
return table.concat({ ... }, '%%', 1, argc)
end
end
--- @param idx integer
--- @return fun(...): any
local function idx_hash(idx)
return function(...)
return select(idx, ...)
end
end
--- @param hash integer|string|fun(...): any
--- @return fun(...): any
local function resolve_hash(hash)
if type(hash) == 'number' then
hash = idx_hash(hash)
elseif type(hash) == 'string' then
local c = hash == 'concat' or hash:match('^concat%-(%d+)')
if c then
hash = concat_hash(tonumber(c))
else
error('invalid value for hash: ' .. hash)
end
end
--- @cast hash -integer
return hash
end
--- @generic F: function
--- @param hash integer|string|fun(...): any
--- @param fn F
--- @return F
return function(hash, fn)
vim.validate({
hash = { hash, { 'number', 'string', 'function' } },
fn = { fn, 'function' },
})
---@type table<any,table<any,any>>
local cache = setmetatable({}, { __mode = 'kv' })
hash = resolve_hash(hash)
return function(...)
local key = hash(...)
if cache[key] == nil then
cache[key] = vim.F.pack_len(fn(...))
end
return vim.F.unpack_len(cache[key])
end
end