Add a distance-based comparator

README.md

@@ -2,7 +2,7 @@
 
 nvim-cmp source for buffer words.
 
-# Setup
+## Setup
 
 ```lua
 require'cmp'.setup {
@@ -12,7 +12,7 @@ require'cmp'.setup {
 }
 ```
 
-# Configuration
+## Configuration
 
 The below source configuration are available.
 
@@ -60,3 +60,29 @@ get_bufnrs = function()
   return vim.tbl_keys(bufs)
 end
 ```
+
+
+## Distance-based sorting
+
+This source also provides a comparator function which uses information from the word indexer
+to sort completion results based on the distance of the word from the cursor line. It will also
+sort completion results coming from other sources, such as Language Servers, which might improve
+accuracy of their suggestions too. The usage is as follows:
+
+```lua
+local cmp_buffer = require('cmp_buffer')
+
+cmp.setup({
+  sources = {
+    { name = 'buffer' }
+      -- The rest of your sources...
+  },
+
+  sorting = {
+    comparators = {
+      cmp_buffer.compare_word_distance,
+      -- The rest of your comparators...
+    }
+  }
+})
+```

lua/cmp_buffer/buffer.lua

@@ -15,6 +15,9 @@
 ---@field public last_edit_last_line number
 ---@field public closed boolean
 ---@field public on_close_cb fun()|nil
+---@field public words_distances table<string, number>
+---@field public words_distances_last_cursor_pos number
+---@field public words_distances_dirty boolean
 local buffer = {}
 
 ---Create new buffer object
@@ -44,6 +47,10 @@ function buffer.new(bufnr, opts)
   self.last_edit_first_line = 0
   self.last_edit_last_line = 0
 
+  self.words_distances = {}
+  self.words_distances_dirty = true
+  self.words_distances_last_cursor_pos = 0
+
   return self
 end
 
@@ -62,6 +69,10 @@ function buffer.close(self)
   self.last_edit_first_line = 0
   self.last_edit_last_line = 0
 
+  self.words_distances = {}
+  self.words_distances_dirty = false
+  self.words_distances_last_cursor_pos = 0
+
   if self.on_close_cb then
     self.on_close_cb()
   end
@@ -123,6 +134,7 @@ function buffer.index_range_async(self, range_start, range_end)
       end)
       chunk_start = chunk_end
       self:mark_all_lines_dirty()
+      self.words_distances_dirty = true
 
       if chunk_end >= range_end then
         self:stop_indexing_timer()
@@ -201,6 +213,8 @@ function buffer.watch(self)
       end
       self.last_edit_first_line = first_line
       self.last_edit_last_line = new_last_line
+
+      self.words_distances_dirty = true
     end,
 
     on_reload = function(_, _)
@@ -224,6 +238,7 @@ function buffer.watch(self)
 
       self:index_range(0, self.lines_count)
       self:mark_all_lines_dirty()
+      self.words_distances_dirty = true
     end,
 
     on_detach = function(_, _)
@@ -295,4 +310,24 @@ function buffer.rebuild_unique_words(self, words_table, range_start, range_end)
   end
 end
 
+---@param cursor_pos number 1-indexed line number
+---@return table<string, number>
+function buffer.get_words_distances(self, cursor_pos)
+  if self.words_distances_dirty or cursor_pos ~= self.words_distances_last_cursor_pos then
+    local distances = self.words_distances
+    for k in pairs(distances) do
+      distances[k] = nil
+    end
+    for linenr, line in ipairs(self.lines_words) do
+      for _, w in ipairs(line) do
+        local dist = math.abs(cursor_pos - linenr)
+        distances[w] = distances[w] and math.min(distances[w], dist) or dist
+      end
+    end
+    self.words_distances_last_cursor_pos = cursor_pos
+    self.words_distances_dirty = false
+  end
+  return self.words_distances
+end
+
 return buffer

lua/cmp_buffer/init.lua

@@ -94,4 +94,25 @@ source._get_buffers = function(self, opts)
   return buffers
 end
 
+source._get_distance_from_entry = function(entry)
+  if getmetatable(entry.source.source).__index == source then
+    local buf = entry.source.source.buffers[entry.context.bufnr]
+    if buf then
+      local distances = buf:get_words_distances(entry.context.cursor.line + 1)
+      return distances[entry.completion_item.filterText] or distances[entry.completion_item.label]
+    end
+  end
+end
+
+source.compare_word_distance = function(entry1, entry2)
+  if entry1.context ~= entry2.context then
+    return
+  end
+  local dist1 = source._get_distance_from_entry(entry1) or math.huge
+  local dist2 = source._get_distance_from_entry(entry2) or math.huge
+  if dist1 ~= dist2 then
+    return dist1 < dist2
+  end
+end
+
 return source