refactor: Update memoization to clean up unused cache

lua/orgmode/files/file.lua

@@ -7,6 +7,7 @@ local config = require('orgmode.config')
 local Block = require('orgmode.files.elements.block')
 local Link = require('orgmode.org.hyperlinks.link')
 local Range = require('orgmode.files.elements.range')
+local Memoize = require('orgmode.utils.memoize')
 
 ---@class OrgFileMetadata
 ---@field mtime number
@@ -26,12 +27,11 @@ local Range = require('orgmode.files.elements.range')
 ---@field root TSNode
 local OrgFile = {}
 
-local memoize = utils.memoize(OrgFile, function(self)
-  return table.concat({
-    self.filename,
-    self.root and self.root:id() or '',
-    self.metadata.mtime,
-  }, '_')
+local memoize = Memoize:new(OrgFile, function(self)
+  return {
+    file = self,
+    id = table.concat({ 'file', self.root and self.root:id() or '' }, '_'),
+  }
 end)
 
 ---Constructor function, should not be used directly
@@ -740,8 +740,10 @@ function OrgFile:_update_lines(lines, bufnr)
   self:parse()
   if bufnr then
     self.metadata.changedtick = vim.api.nvim_buf_get_changedtick(bufnr)
-  else
-    self.metadata.mtime = vim.loop.fs_stat(self.filename).mtime.nsec
+  end
+  local stat = vim.loop.fs_stat(self.filename)
+  if stat then
+    self.metadata.mtime = stat.mtime.nsec
   end
   return self
 end

lua/orgmode/files/headline.lua

@@ -7,6 +7,7 @@ local PriorityState = require('orgmode.objects.priority_state')
 local indent = require('orgmode.org.indent')
 local Logbook = require('orgmode.files.elements.logbook')
 local OrgId = require('orgmode.org.id')
+local Memoize = require('orgmode.utils.memoize')
 
 ---@alias OrgPlanDateTypes 'DEADLINE' | 'SCHEDULED' | 'CLOSED'
 
@@ -15,13 +16,12 @@ local OrgId = require('orgmode.org.id')
 ---@field file OrgFile
 local Headline = {}
 
-local memoize = utils.memoize(Headline, function(self)
+local memoize = Memoize:new(Headline, function(self)
   ---@cast self OrgHeadline
-  return table.concat({
-    self.file.filename,
-    self.headline:id(),
-    self.file.metadata.mtime,
-  }, '_')
+  return {
+    file = self.file,
+    id = table.concat({ 'headline', self.headline:id() }, '_'),
+  }
 end)
 
 ---@param headline_node TSNode tree sitter headline node

lua/orgmode/utils/init.lua

@@ -547,50 +547,6 @@ function utils.edit_file(filename)
   }
 end
 
-function utils.memoize(class, key_getter)
-  local memoizedMethods = {}
-  local methodsToMemoize = {}
-  local cache = setmetatable({}, { __mode = 'k' })
-
-  local function memoizedIndex(_, key)
-    local method = class[key]
-
-    if type(method) == 'function' and methodsToMemoize[key] and not memoizedMethods[key] then
-      memoizedMethods[key] = function(self, ...)
-        local top_key = key_getter(self)
-        local arg_key = key .. '_' .. table.concat({ ... }, '_')
-
-        if not cache[top_key] then
-          cache[top_key] = {}
-        end
-
-        if not cache[top_key][arg_key] then
-          local value = vim.F.pack_len(method(self, ...))
-          cache[top_key][arg_key] = value
-        end
-
-        local cached_value = cache[top_key][arg_key]
-
-        if cached_value then
-          local result = { pcall(vim.F.unpack_len, cached_value) }
-          if result[1] then
-            return unpack(result, 2)
-          end
-        end
-      end
-    end
-
-    return memoizedMethods[key] or method
-  end
-
-  class.__index = memoizedIndex
-
-  return function(method)
-    methodsToMemoize[method] = true
-    return true
-  end
-end
-
 function utils.has_version_10()
   local v = vim.version()
   return not vim.version.lt({ v.major, v.minor, v.patch }, { 0, 10, 0 })

lua/orgmode/utils/memoize.lua

@@ -0,0 +1,89 @@
+---@alias MemoizeKey { file: OrgFile, id: string }
+
+---@class OrgMemoize
+---@field class table
+---@field key_getter fun(self: table): MemoizeKey
+---@field memoized_methods table<string, fun(self: table, ...): any>
+---@field methods_to_memoize table<string, boolean>
+local Memoize = {
+  cache = setmetatable({}, { __mode = 'k' }),
+}
+Memoize.__index = Memoize
+
+---@return fun(method: string): boolean
+function Memoize:new(class, key_getter)
+  local this = setmetatable({
+    class = class,
+    key_getter = key_getter,
+    memoized_methods = {},
+    methods_to_memoize = {},
+  }, Memoize)
+
+  this:setup()
+
+  return function(method)
+    this.methods_to_memoize[method] = true
+    return true
+  end
+end
+
+function Memoize:setup()
+  self.class.__index = function(_, key)
+    local method = self.class[key]
+
+    -- Not memoizable or not required to be memoized
+    if type(method) ~= 'function' or not self.methods_to_memoize[key] then
+      return method
+    end
+
+    -- Already memoized
+    if self.memoized_methods[key] then
+      return self.memoized_methods[key]
+    end
+
+    self.memoized_methods[key] = function(method_self, ...)
+      local memoize_key = self.key_getter(method_self)
+      local cache = self:_get_cache_for_key(memoize_key)
+      local arg_key = key .. '_' .. table.concat({ ... }, '_')
+
+      if not cache[arg_key] then
+        local value = vim.F.pack_len(method(method_self, ...))
+        cache[arg_key] = value
+      end
+
+      local cached_value = cache[arg_key]
+
+      if cached_value then
+        local result = { pcall(vim.F.unpack_len, cached_value) }
+        if result[1] then
+          return unpack(result, 2)
+        end
+      end
+    end
+
+    return self.memoized_methods[key]
+  end
+end
+
+---@private
+---@param memoize_key MemoizeKey
+---@return string
+function Memoize:_get_cache_for_key(memoize_key)
+  local id = memoize_key.id
+  local filename = memoize_key.file.filename
+  local version_key = memoize_key.file.metadata.mtime
+
+  if not self.cache[filename] or self.cache[filename].__version ~= version_key then
+    self.cache[filename] = {
+      __version = version_key,
+    }
+  end
+
+  if not self.cache[filename][id] then
+    self.cache[filename][id] = {}
+  end
+
+  return self.cache[filename][id]
+end
+
+return Memoize