fix(memoize): release cache with file, reset on re-parse

lua/orgmode/files/file.lua

@@ -32,6 +32,7 @@ local clean_empty_line = vim.fn.has('nvim-0.13') == 1 or vim.fn.has('nvim-0.12.3
 ---@field metadata OrgFileMetadata
 ---@field parser vim.treesitter.LanguageTree
 ---@field root TSNode
+---@field memoize_cache? table Memoized method results, released with the file
 local OrgFile = {}
 
 -- The root node id only changes on a re-parse, so a buffer edited since the

lua/orgmode/utils/memoize.lua

@@ -5,9 +5,7 @@
 ---@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' }),
-}
+local Memoize = {}
 Memoize.__index = Memoize
 
 ---@return fun(method: string): boolean
@@ -70,20 +68,31 @@ end
 ---@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] = {
+  -- The cache is stored on the file so that it is released together with it.
+  -- A module-level weak table cannot do this: LuaJIT has no ephemeron support,
+  -- and cached values reference the file they were built from, so every entry
+  -- stayed reachable through its own value for the whole session.
+  local file = memoize_key.file
+  local version_key = file.metadata.mtime
+  -- Lookups after a re-parse use the ids of the new tree, so the buckets
+  -- built for the previous tree are dead. Drop them as a whole instead of
+  -- leaving a set behind on every edit.
+  local generation = file.root and file.root:id() or ''
+  local entry = file.memoize_cache
+
+  if not entry or entry.__version ~= version_key or entry.__generation ~= generation then
+    entry = {
       __version = version_key,
+      __generation = generation,
     }
+    file.memoize_cache = entry
   end
 
-  if not self.cache[filename][id] then
-    self.cache[filename][id] = {}
+  if not entry[id] then
+    entry[id] = {}
   end
 
-  return self.cache[filename][id]
+  return entry[id]
 end
 
 return Memoize

tests/plenary/utils/memoize_spec.lua

@@ -0,0 +1,63 @@
+local OrgFile = require('orgmode.files.file')
+
+describe('Memoize', function()
+  ---@return OrgFile
+  local new_file = function(content, filename)
+    content = content or {}
+    filename = filename or vim.fn.tempname() .. '.org'
+    return OrgFile:new({ filename = filename, lines = content })
+  end
+
+  ---Number of per-node buckets held for a file. Bookkeeping keys are not
+  ---buckets.
+  ---@param file OrgFile
+  local function buckets(file)
+    local count = 0
+    for key in pairs(file.memoize_cache or {}) do
+      if not tostring(key):match('^__') then
+        count = count + 1
+      end
+    end
+    return count
+  end
+
+  local function collect()
+    collectgarbage('collect')
+    collectgarbage('collect')
+  end
+
+  it('does not keep a file alive once nothing else references it', function()
+    local weak = setmetatable({}, { __mode = 'v' })
+
+    do
+      local file = new_file({ '* Headline 1', '  body' })
+      weak.file = file
+      file:get_headlines()
+      assert.is_not_nil(weak.file)
+    end
+
+    collect()
+
+    assert.is_nil(weak.file)
+  end)
+
+  it('does not accumulate buckets when the tree is replaced', function()
+    local file = new_file({ '* Headline 1', '  body' })
+    file:get_headlines()
+
+    local baseline = buckets(file)
+
+    -- Rewrite the body five times. The headline count never changes, so the
+    -- number of live buckets must not change either.
+    -- Dropping the parser forces a fresh tree on each round, which is what an
+    -- edit in a loaded buffer does; without it the parser is reused and the
+    -- node ids stay identical.
+    for i = 1, 5 do
+      file.parser = nil
+      file:_update_lines({ '* Headline 1', '  body ' .. i })
+      file:get_headlines()
+    end
+
+    assert.are.same(baseline, buckets(file))
+  end)
+end)