From b61051ccb4c23958d43d285b8b801af11620264f Mon Sep 17 00:00:00 2001
From: Riley Bruins <ribru17@hotmail.com>
Date: Sun, 1 Sep 2024 16:54:30 -0700
Subject: feat(func): allow manual cache invalidation for _memoize

This commit also adds some tests for the existing memoization
functionality.
---
 runtime/lua/vim/func/_memoize.lua | 56 +++++++++++++++++++++++++++------------
 1 file changed, 39 insertions(+), 17 deletions(-)

(limited to 'runtime/lua/vim/func/_memoize.lua')

diff --git a/runtime/lua/vim/func/_memoize.lua b/runtime/lua/vim/func/_memoize.lua
index 6e557905a7..c46f878067 100644
--- a/runtime/lua/vim/func/_memoize.lua
+++ b/runtime/lua/vim/func/_memoize.lua
@@ -1,5 +1,7 @@
 --- Module for private utility functions
 
+--- @alias vim.func.MemoObj { _hash: (fun(...): any), _weak: boolean?, _cache: table<any> }
+
 --- @param argc integer?
 --- @return fun(...): any
 local function concat_hash(argc)
@@ -33,29 +35,49 @@ local function resolve_hash(hash)
   return hash
 end
 
+--- @param weak boolean?
+--- @return table
+local create_cache = function(weak)
+  return setmetatable({}, {
+    __mode = weak ~= false and 'kv',
+  })
+end
+
 --- @generic F: function
 --- @param hash integer|string|fun(...): any
 --- @param fn F
---- @param strong? boolean
+--- @param weak? boolean
 --- @return F
-return function(hash, fn, strong)
+return function(hash, fn, weak)
   vim.validate('hash', hash, { 'number', 'string', 'function' })
   vim.validate('fn', fn, 'function')
+  vim.validate('weak', weak, 'boolean', true)
 
-  ---@type table<any,table<any,any>>
-  local cache = {}
-  if not strong then
-    setmetatable(cache, { __mode = 'kv' })
-  end
-
-  hash = resolve_hash(hash)
+  --- @type vim.func.MemoObj
+  local obj = {
+    _cache = create_cache(weak),
+    _hash = resolve_hash(hash),
+    _weak = weak,
+    --- @param self vim.func.MemoObj
+    clear = function(self, ...)
+      if select('#', ...) == 0 then
+        self._cache = create_cache(self._weak)
+        return
+      end
+      local key = self._hash(...)
+      self._cache[key] = nil
+    end,
+  }
 
-  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
+  return setmetatable(obj, {
+    --- @param self vim.func.MemoObj
+    __call = function(self, ...)
+      local key = self._hash(...)
+      local cache = self._cache
+      if cache[key] == nil then
+        cache[key] = vim.F.pack_len(fn(...))
+      end
+      return vim.F.unpack_len(cache[key])
+    end,
+  })
 end
-- 
cgit