improve reliability of indexing while editing, make on_reload async

lua/cmp_buffer/buffer.lua

@@ -1,9 +1,12 @@
+local timer = require('cmp_buffer.timer')
+
 ---@class cmp_buffer.Buffer
 ---@field public bufnr number
 ---@field public opts cmp_buffer.Options
 ---@field public regex any
----@field public timer any|nil
+---@field public timer cmp_buffer.Timer
 ---@field public lines_count number
+---@field public timer_current_line number
 ---@field public lines_words table<number, string[]>
 ---@field public unique_words_curr_line table<string, boolean>
 ---@field public unique_words_other_lines table<string, boolean>
@@ -30,7 +33,7 @@ function buffer.new(bufnr, opts)
   local self = setmetatable({}, { __index = buffer })
 
   self.bufnr = bufnr
-  self.timer = nil
+  self.timer = timer.new()
   self.closed = false
   self.on_close_cb = nil
 
@@ -38,6 +41,7 @@ function buffer.new(bufnr, opts)
   self.regex = vim.regex(self.opts.keyword_pattern)
 
   self.lines_count = 0
+  self.timer_current_line = -1
   self.lines_words = {}
 
   self.unique_words_curr_line = {}
@@ -58,8 +62,11 @@ end
 function buffer.close(self)
   self.closed = true
   self:stop_indexing_timer()
+  self.timer:close()
+  self.timer = nil
 
   self.lines_count = 0
+  self.timer_current_line = -1
   self.lines_words = {}
 
   self.unique_words_curr_line = {}
@@ -79,11 +86,8 @@ function buffer.close(self)
 end
 
 function buffer.stop_indexing_timer(self)
-  if self.timer and not self.timer:is_closing() then
-    self.timer:stop()
-    self.timer:close()
-  end
-  self.timer = nil
+  self.timer:stop()
+  self.timer_current_line = -1
 end
 
 function buffer.mark_all_lines_dirty(self)
@@ -91,23 +95,7 @@ function buffer.mark_all_lines_dirty(self)
   self.unique_words_other_lines_dirty = true
   self.last_edit_first_line = 0
   self.last_edit_last_line = 0
-end
-
----Indexing buffer
-function buffer.index(self)
-  self.lines_count = vim.api.nvim_buf_line_count(self.bufnr)
-  -- NOTE: Pre-allocating self.lines_words here somehow wastes more memory, and
-  -- not doing that doesn't have a visible effect on performance. Win-win.
-  -- for i = 1, self.lines_count do
-  --   self.lines_words[i] = {}
-  -- end
-
-  if self.opts.indexing_interval < 1 then
-    self:index_range(0, self.lines_count)
-    self:mark_all_lines_dirty()
-  else
-    self:index_buffer_async()
-  end
+  self.words_distances_dirty = true
 end
 
 --- Workaround for https://github.com/neovim/neovim/issues/16729
@@ -119,7 +107,6 @@ function buffer.safe_buf_call(self, callback)
   end
 end
 
---- sync algorithm
 function buffer.index_range(self, range_start, range_end, skip_already_indexed)
   self:safe_buf_call(function()
     local chunk_size = self.GET_LINES_CHUNK_SIZE
@@ -137,55 +124,50 @@ function buffer.index_range(self, range_start, range_end, skip_already_indexed)
   end)
 end
 
---- async algorithm
-function buffer.index_buffer_async(self)
-  local chunk_start = 0
+function buffer.start_indexing_timer(self)
+  self.lines_count = vim.api.nvim_buf_line_count(self.bufnr)
+  self.timer_current_line = 0
 
-  -- This flag prevents vim.schedule() callbacks from piling up in the queue
-  -- when the indexing interval is very short.
-  local scheduled = false
-  self.timer = vim.loop.new_timer()
   -- Negative values result in an integer overflow in luv (vim.loop), and zero
   -- disables timer repeat, so only intervals larger than 1 are valid.
   local interval = math.max(1, self.opts.indexing_interval)
   self.timer:start(0, interval, function()
-    if scheduled then
+    if self.closed then
+      self:stop_indexing_timer()
       return
     end
-    scheduled = true
-    vim.schedule(function()
-      scheduled = false
-      if self.closed then
-        return
-      end
 
-      -- Note that the async indexer is designed to not break even if the user
-      -- is editing the file while it is in the process of being indexed.
-      -- Because the indexing in watcher must use the synchronous algorithm, we
-      -- assume that the data already present in self.lines_words to be correct
-      -- and doesn't need refreshing here because even if we do receive text
-      -- from nvim_buf_get_lines different from what the watcher has seen, it
-      -- will catch up on the next on_lines event.
+    -- Note that the async indexer is designed to not break even if the user is
+    -- editing the file while it is in the process of being indexed. Because
+    -- the indexing in watcher must use the synchronous algorithm, we assume
+    -- that the data already present in self.lines_words to be correct and
+    -- doesn't need refreshing here because even if we do receive text from
+    -- nvim_buf_get_lines different from what the watcher has seen so far, it
+    -- (the watcher) will catch up on the next on_lines event.
 
-      local line_count = vim.api.nvim_buf_line_count(self.bufnr)
-      -- Skip over the already indexed lines
-      while chunk_start < line_count and self.lines_words[chunk_start + 1] do
-        chunk_start = chunk_start + 1
-      end
-      local chunk_end = math.min(chunk_start + self.opts.indexing_chunk_size, line_count)
-      if chunk_end >= line_count then
-        self:stop_indexing_timer()
-      end
-      self:index_range(chunk_start, chunk_end, true)
-      chunk_start = chunk_end
-      self:mark_all_lines_dirty()
-      self.words_distances_dirty = true
-    end)
+    -- Skip over the already indexed lines
+    while self.lines_words[self.timer_current_line + 1] do
+      self.timer_current_line = self.timer_current_line + 1
+    end
+
+    local chunk_start = self.timer_current_line
+    local chunk_size = self.opts.indexing_chunk_size
+    -- NOTE: self.lines_count may be modified by the indexer.
+    local chunk_end = chunk_size >= 1 and math.min(chunk_start + chunk_size, self.lines_count) or self.lines_count
+    if chunk_end >= self.lines_count then
+      self:stop_indexing_timer()
+    end
+
+    self:index_range(chunk_start, chunk_end, true)
+    self.timer_current_line = chunk_end
+    self:mark_all_lines_dirty()
   end)
 end
 
 --- watch
 function buffer.watch(self)
+  self.lines_count = vim.api.nvim_buf_line_count(self.bufnr)
+
   -- NOTE: As far as I know, indexing in watching can't be done asynchronously
   -- because even built-in commands generate multiple consequent `on_lines`
   -- events, and I'm not even mentioning plugins here. To get accurate results
@@ -250,6 +232,25 @@ function buffer.watch(self)
       end
       self.lines_count = new_lines_count
 
+      -- This branch is support code for handling cases when the user is
+      -- editing the buffer while the async indexer is running. It solves the
+      -- problem that if new lines are inserted or old lines are deleted, the
+      -- indexes of each subsequent line will change, and so the indexer
+      -- current position must be adjusted to not accidentally skip any lines.
+      if self.timer:is_active() then
+        if first_line <= self.timer_current_line and self.timer_current_line < old_last_line then
+          -- The indexer was in the area of the current text edit. We will
+          -- synchronously index this area it in a moment, so the indexer
+          -- should resume from right after the edit range.
+          self.timer_current_line = new_last_line
+        elseif self.timer_current_line >= old_last_line then
+          -- The indexer was somewhere past the current text edit. This means
+          -- that the line numbers could have changed, and the indexing
+          -- position must be adjusted accordingly.
+          self.timer_current_line = self.timer_current_line + delta
+        end
+      end
+
       -- replace lines
       self:index_range(first_line, new_last_line)
 
@@ -270,25 +271,13 @@ function buffer.watch(self)
         return true
       end
 
-      -- The logic for adjusting lines list on buffer reloads is much simpler
-      -- because tables of all lines can be assumed to be fresh.
-      local new_lines_count = vim.api.nvim_buf_line_count(self.bufnr)
-      if new_lines_count > self.lines_count then -- append
-        -- Again, no need to pre-allocate, index_line will append new lines
-        -- itself.
-        -- for i = self.lines_count + 1, new_lines_count do
-        --   self.lines_words[i] = {}
-        -- end
-      elseif new_lines_count < self.lines_count then -- remove
-        for i = self.lines_count, new_lines_count + 1, -1 do
-          self.lines_words[i] = nil
-        end
+      -- clear all lines
+      for i = self.lines_count, 1, -1 do
+        self.lines_words[i] = nil
       end
-      self.lines_count = new_lines_count
 
-      self:index_range(0, self.lines_count)
-      self:mark_all_lines_dirty()
-      self.words_distances_dirty = true
+      self:stop_indexing_timer()
+      self:start_indexing_timer()
     end,
 
     on_detach = function(_, _)

lua/cmp_buffer/source.lua

@@ -50,7 +50,7 @@ source.complete = function(self, params, callback)
   local processing = false
   local bufs = self:_get_buffers(opts)
   for _, buf in ipairs(bufs) do
-    if buf.timer then
+    if buf.timer:is_active() then
       processing = true
       break
     end
@@ -90,7 +90,7 @@ source._get_buffers = function(self, opts)
       new_buf.on_close_cb = function()
         self.buffers[bufnr] = nil
       end
-      new_buf:index()
+      new_buf:start_indexing_timer()
       new_buf:watch()
       self.buffers[bufnr] = new_buf
     end

lua/cmp_buffer/timer.lua

@@ -0,0 +1,48 @@
+---@class cmp_buffer.Timer
+---@field public handle any
+---@field private callback_wrapper_instance fun()|nil
+local timer = {}
+
+function timer.new()
+  local self = setmetatable({}, { __index = timer })
+  self.handle = vim.loop.new_timer()
+  self.callback_wrapper_instance = nil
+  return self
+end
+
+---@param timeout_ms number
+---@param repeat_ms number
+---@param callback fun()
+function timer:start(timeout_ms, repeat_ms, callback)
+  local scheduled = false
+  local function callback_wrapper()
+    if scheduled then
+      return
+    end
+    scheduled = true
+    vim.schedule(function()
+      scheduled = false
+      if self.callback_wrapper_instance ~= callback_wrapper then
+        return
+      end
+      callback()
+    end)
+  end
+  self.handle:start(timeout_ms, repeat_ms, callback_wrapper)
+  self.callback_wrapper_instance = callback_wrapper
+end
+
+function timer:stop()
+  self.handle:stop()
+  self.callback_wrapper_instance = nil
+end
+
+function timer:is_active()
+  return self.handle:is_active()
+end
+
+function timer:close()
+  self.handle:close()
+end
+
+return timer