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'.

1 files changed, 59 insertions, 0 deletions

diff --git a/runtime/lua/vim/func/_memoize.lua b/runtime/lua/vim/func/_memoize.lua
new file mode 100644
index 0000000000..835bf64c93
--- /dev/null
+++ b/runtime/lua/vim/func/_memoize.lua
@@ -0,0 +1,59 @@
+--- 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