Rebuild word tables incrementally when editing happens on a single line

lua/cmp_buffer/buffer.lua

@@ -8,8 +8,12 @@
 ---@field public timer any|nil
 ---@field public lines_count number
 ---@field public lines_words table<number, string[]>
----@field public unique_words table<string, boolean>
----@field public unique_words_dirty boolean
+---@field public unique_words_curr_line table<string, boolean>
+---@field public unique_words_other_lines table<string, boolean>
+---@field public unique_words_curr_line_dirty boolean
+---@field public unique_words_other_lines_dirty boolean
+---@field public last_edit_first_line number
+---@field public last_edit_last_line number
 ---@field public closed boolean
 ---@field public on_close_cb fun()|nil
 local buffer = {}
@@ -21,19 +25,28 @@ local buffer = {}
 ---@return cmp_buffer.Buffer
 function buffer.new(bufnr, length, pattern)
   local self = setmetatable({}, { __index = buffer })
+
   self.bufnr = bufnr
+  self.timer = nil
+  self.closed = false
+  self.on_close_cb = nil
+
   self.regex = vim.regex(pattern)
   self.length = length
   self.pattern = pattern
   self.indexing_chunk_size = 1000
   self.indexing_interval = 200
-  self.timer = nil
+
   self.lines_count = 0
   self.lines_words = {}
-  self.unique_words = {}
-  self.unique_words_dirty = true
-  self.closed = false
-  self.on_close_cb = nil
+
+  self.unique_words_curr_line = {}
+  self.unique_words_other_lines = {}
+  self.unique_words_curr_line_dirty = true
+  self.unique_words_other_lines_dirty = true
+  self.last_edit_first_line = 0
+  self.last_edit_last_line = 0
+
   return self
 end
 
@@ -41,10 +54,17 @@ end
 function buffer.close(self)
   self.closed = true
   self:stop_indexing_timer()
+
   self.lines_count = 0
   self.lines_words = {}
-  self.unique_words = {}
-  self.unique_words_dirty = false
+
+  self.unique_words_curr_line = {}
+  self.unique_words_other_lines = {}
+  self.unique_words_curr_line_dirty = false
+  self.unique_words_other_lines_dirty = false
+  self.last_edit_first_line = 0
+  self.last_edit_last_line = 0
+
   if self.on_close_cb then
     self.on_close_cb()
   end
@@ -58,6 +78,13 @@ function buffer.stop_indexing_timer(self)
   self.timer = nil
 end
 
+function buffer.mark_all_lines_dirty(self)
+  self.unique_words_curr_line_dirty = true
+  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)
@@ -98,7 +125,7 @@ function buffer.index_range_async(self, range_start, range_end)
         end
       end)
       chunk_start = chunk_end
-      self.unique_words_dirty = true
+      self:mark_all_lines_dirty()
 
       if chunk_end >= range_end then
         self:stop_indexing_timer()
@@ -168,7 +195,15 @@ function buffer.watch(self)
 
       -- replace lines
       self:index_range(first_line, new_last_line)
-      self.unique_words_dirty = true
+
+      if first_line == self.last_edit_first_line and old_last_line == self.last_edit_last_line and new_last_line == self.last_edit_last_line then
+        self.unique_words_curr_line_dirty = true
+      else
+        self.unique_words_curr_line_dirty = true
+        self.unique_words_other_lines_dirty = true
+      end
+      self.last_edit_first_line = first_line
+      self.last_edit_last_line = new_last_line
     end,
 
     on_reload = function(_, _)
@@ -191,7 +226,7 @@ function buffer.watch(self)
       self.lines_count = new_lines_count
 
       self:index_range(0, self.lines_count)
-      self.unique_words_dirty = true
+      self:mark_all_lines_dirty()
     end,
 
     on_detach = function(_, _)
@@ -234,22 +269,33 @@ function buffer.get_words(self)
   -- NOTE: unique_words are rebuilt on-demand because it is common for the
   -- watcher callback to be fired VERY frequently, and a rebuild needs to go
   -- over ALL lines, not just the changed ones.
-  if self.unique_words_dirty then
-    self:rebuild_unique_words()
+  if self.unique_words_other_lines_dirty then
+    local words = self.unique_words_other_lines
+    for w, _ in pairs(words) do
+      words[w] = nil
+    end
+    self:rebuild_unique_words(words, 0, self.last_edit_first_line)
+    self:rebuild_unique_words(words, self.last_edit_last_line, self.lines_count)
+    self.unique_words_other_lines_dirty = false
   end
-  return self.unique_words
+  if self.unique_words_curr_line_dirty then
+    local words = self.unique_words_curr_line
+    for w, _ in pairs(words) do
+      words[w] = nil
+    end
+    self:rebuild_unique_words(words, self.last_edit_first_line, self.last_edit_last_line)
+    self.unique_words_curr_line_dirty = false
+  end
+  return { self.unique_words_other_lines, self.unique_words_curr_line }
 end
 
-function buffer.rebuild_unique_words(self)
-  for w, _ in pairs(self.unique_words) do
-    self.unique_words[w] = nil
-  end
-  for _, line in ipairs(self.lines_words) do
-    for _, w in ipairs(line) do
-      self.unique_words[w] = true
+--- rebuild_unique_words
+function buffer.rebuild_unique_words(self, words_table, range_start, range_end)
+  for i = range_start + 1, range_end do
+    for _, w in ipairs(self.lines_words[i]) do
+      words_table[w] = true
     end
   end
-  self.unique_words_dirty = false
 end
 
 return buffer

lua/cmp_buffer/init.lua

@@ -47,13 +47,15 @@ source.complete = function(self, params, callback)
     local items = {}
     local words = {}
     for _, buf in ipairs(bufs) do
-      for word, _ in pairs(buf:get_words()) do
-        if not words[word] and input ~= word then
-          words[word] = true
-          table.insert(items, {
-            label = word,
-            dup = 0,
-          })
+      for _, word_list in ipairs(buf:get_words()) do
+        for word, _ in pairs(word_list) do
+          if not words[word] and input ~= word then
+            words[word] = true
+            table.insert(items, {
+              label = word,
+              dup = 0,
+            })
+          end
         end
       end
     end