feat: lua fuzzy matching implementation

Author: Saghen

Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [lib]
-path = "lua/blink/cmp/fuzzy/lib.rs"
+path = "lua/blink/cmp/fuzzy/rust/lib.rs"
 crate-type = ["cdylib"]
 
 [dependencies]

lua/blink/cmp/fuzzy/init.lua

@@ -2,25 +2,28 @@ local config = require('blink.cmp.config')
 
 --- @class blink.cmp.Fuzzy
 local fuzzy = {
-  rust = require('blink.cmp.fuzzy.rust'),
+  provider_type = 'lua',
+  provider = require('blink.cmp.fuzzy.lua'),
   haystacks_by_provider_cache = {},
   has_init_db = false,
 }
 
 function fuzzy.init_db()
   if fuzzy.has_init_db then return end
 
-  fuzzy.rust.init_db(vim.fn.stdpath('data') .. '/blink/cmp/fuzzy.db', config.use_unsafe_no_lock)
+  fuzzy.provider.init_db(vim.fn.stdpath('data') .. '/blink/cmp/fuzzy.db', config.use_unsafe_no_lock)
 
   vim.api.nvim_create_autocmd('VimLeavePre', {
-    callback = fuzzy.rust.destroy_db,
+    callback = fuzzy.provider.destroy_db,
   })
 
   fuzzy.has_init_db = true
 end
 
 ---@param item blink.cmp.CompletionItem
 function fuzzy.access(item)
+  if fuzzy.provider_type ~= 'rust' then return end
+
   fuzzy.init_db()
 
   -- writing to the db takes ~10ms, so schedule writes in another thread
@@ -33,14 +36,14 @@ function fuzzy.access(item)
 end
 
 ---@param lines string
-function fuzzy.get_words(lines) return fuzzy.rust.get_words(lines) end
+function fuzzy.get_words(lines) return fuzzy.provider.get_words(lines) end
 
 --- @param line string
 --- @param cursor_col number
 --- @param haystack string[]
 --- @param range blink.cmp.CompletionKeywordRange
 function fuzzy.fuzzy_matched_indices(line, cursor_col, haystack, range)
-  return fuzzy.rust.fuzzy_matched_indices(line, cursor_col, haystack, range == 'full')
+  return fuzzy.provider.fuzzy_matched_indices(line, cursor_col, haystack, range == 'full')
 end
 
 --- @param line string
@@ -55,7 +58,7 @@ function fuzzy.fuzzy(line, cursor_col, haystacks_by_provider, range)
     -- set the provider items once since Lua <-> Rust takes the majority of the time
     if fuzzy.haystacks_by_provider_cache[provider_id] ~= haystack then
       fuzzy.haystacks_by_provider_cache[provider_id] = haystack
-      fuzzy.rust.set_provider_items(provider_id, haystack)
+      fuzzy.provider.set_provider_items(provider_id, haystack)
     end
   end
 
@@ -64,30 +67,28 @@ function fuzzy.fuzzy(line, cursor_col, haystacks_by_provider, range)
   local start_row = math.max(0, cursor_row - 30)
   local end_row = math.min(cursor_row + 30, vim.api.nvim_buf_line_count(0))
   local nearby_text = table.concat(vim.api.nvim_buf_get_lines(0, start_row, end_row, false), '\n')
-  local nearby_words = #nearby_text < 10000 and fuzzy.rust.get_words(nearby_text) or {}
+  local nearby_words = #nearby_text < 10000 and fuzzy.provider.get_words(nearby_text) or {}
 
-  local keyword_start_col, keyword_end_col =
-    require('blink.cmp.fuzzy').get_keyword_range(line, cursor_col, config.completion.keyword.range)
+  local keyword_start_col, keyword_end_col = fuzzy.get_keyword_range(line, cursor_col, config.completion.keyword.range)
   local keyword_length = keyword_end_col - keyword_start_col
   local keyword = line:sub(keyword_start_col, keyword_end_col)
 
   local filtered_items = {}
   for provider_id, haystack in pairs(haystacks_by_provider) do
     -- perform fuzzy search
-    local scores, matched_indices, exacts = fuzzy.rust.fuzzy(line, cursor_col, provider_id, {
+    local scores, matched_indices, exacts = fuzzy.provider.fuzzy(line, cursor_col, provider_id, {
       -- TODO: make this configurable
       max_typos = config.fuzzy.max_typos(keyword),
       use_frecency = config.fuzzy.use_frecency and keyword_length > 0,
       use_proximity = config.fuzzy.use_proximity and keyword_length > 0,
-      sorts = config.fuzzy.sorts,
       nearby_words = nearby_words,
       match_suffix = range == 'full',
       snippet_score_offset = config.snippets.score_offset,
     })
 
     for idx, item_index in ipairs(matched_indices) do
       local item = haystack[item_index + 1]
-      --TODO: maybe we should declare these fields in `blink.cmp.CompletionItem`?
+      -- TODO: maybe we should declare these fields in `blink.cmp.CompletionItem`?
       item.score = scores[idx]
       item.exact = exacts[idx]
       table.insert(filtered_items, item)

lua/blink/cmp/fuzzy/lua/init.lua

@@ -0,0 +1,73 @@
+local match = require('blink.cmp.fuzzy.lua.match')
+local match_indices = require('blink.cmp.fuzzy.lua.match_indices')
+local get_keyword_range = require('blink.cmp.fuzzy.lua.keyword').get_keyword_range
+local guess_keyword_range_from_item = require('blink.cmp.fuzzy.lua.keyword').guess_keyword_range_from_item
+
+--- @type blink.cmp.FuzzyProvider
+--- @diagnostic disable-next-line: missing-fields
+local fuzzy = {
+  provider_items = {},
+}
+
+function fuzzy.init_db() end
+function fuzzy.destroy_db() end
+function fuzzy.access() end
+
+local words_regex = vim.regex(
+  [[\%(-\?\d\+\%(\.\d\+\)\?\|\h\%(\w\|á\|Á\|é\|É\|í\|Í\|ó\|Ó\|ú\|Ú\)*\%(-\%(\w\|á\|Á\|é\|É\|í\|Í\|ó\|Ó\|ú\|Ú\)*\)*\)]]
+)
+
+--- Takes ~0.25ms for 1200 characters split over 40 lines
+function fuzzy.get_words(text)
+  local words = {}
+
+  while #text > 0 do
+    local match_start, match_end = words_regex:match_str(text)
+    if match_start == nil then break end
+
+    if match_end - match_start > 2 then
+      local word = text:sub(match_start + 1, match_end)
+      words[word] = true
+    end
+
+    text = text:sub(match_end + 1)
+  end
+
+  return vim.tbl_keys(words)
+end
+
+function fuzzy.set_provider_items(provider_id, items) fuzzy.provider_items[provider_id] = items end
+
+function fuzzy.fuzzy(line, cursor_col, provider_id, range)
+  local keyword_start, keyword_end = get_keyword_range(line, cursor_col, range == 'full')
+  local keyword = line:sub(keyword_start, keyword_end)
+
+  local scores = {}
+  local matched_indices = {}
+  local exacts = {}
+  for idx, item in ipairs(fuzzy.provider_items[provider_id] or {}) do
+    local score, exact = match(keyword, item.filterText or item.label)
+    if score ~= nil then
+      table.insert(scores, score)
+      table.insert(matched_indices, idx - 1)
+      table.insert(exacts, exact)
+    end
+  end
+
+  return scores, matched_indices, exacts
+end
+
+function fuzzy.fuzzy_matched_indices(line, cursor_col, haystack, range)
+  local keyword_start, keyword_end = get_keyword_range(line, cursor_col, range == 'full')
+  local keyword = line:sub(keyword_start, keyword_end)
+
+  return vim.tbl_map(function(text) return match_indices(keyword, text) end, haystack)
+end
+
+function fuzzy.get_keyword_range(line, col, range) return get_keyword_range(line, col, range == 'full') end
+
+function fuzzy.guess_edit_range(item, line, col, range)
+  return guess_keyword_range_from_item(item.insertText or item.label, line, col, range == 'full')
+end
+
+return fuzzy

lua/blink/cmp/fuzzy/lua/keyword.lua

@@ -0,0 +1,67 @@
+local keyword = {}
+
+local BACKWARD_REGEX = vim.regex([[\k*$]])
+local FORWARD_REGEX = vim.regex([[^\k+]])
+
+--- @generic T
+--- @generic Y
+--- @param cb fun(): T, Y
+--- @return T, Y
+function keyword.with_constant_is_keyword(cb)
+  local existing_is_keyword = vim.opt.iskeyword
+  local desired_is_keyword = '@,48-57,_,192-255'
+  if existing_is_keyword == desired_is_keyword then return cb() end
+
+  vim.opt.iskeyword = '@,48-57,_,192-255'
+  local success, a, b = pcall(cb)
+  vim.opt.iskeyword = existing_is_keyword
+
+  if success then return a, b end
+  error(a)
+end
+
+--- @param line string
+--- @param col number
+--- @param match_suffix boolean
+--- @return number, number
+function keyword.get_keyword_range(line, col, match_suffix)
+  return keyword.with_constant_is_keyword(function()
+    local before_match_start = BACKWARD_REGEX:match_str(line:sub(1, col))
+    if not match_suffix then return before_match_start or col, col end
+
+    local after_match_end = FORWARD_REGEX:match_str(line:sub(col + 1))
+    if after_match_end then after_match_end = after_match_end + col end
+    return before_match_start or col, after_match_end or col
+  end)
+end
+
+function keyword.guess_keyword_range_from_item(item_text, line, cursor_col, match_suffix)
+  local line_range_start, line_range_end = keyword.get_keyword_range(line, cursor_col, match_suffix)
+  local text_range_start, _ = keyword.get_keyword_range(item_text, #item_text, false)
+
+  local line_prefix = line:sub(1, line_range_start)
+  local text_prefix = item_text:sub(1, text_range_start)
+  if line_prefix:sub(-#text_prefix) == text_prefix then return line_range_start - #text_prefix, line_range_end end
+
+  return line_range_start, line_range_end
+end
+
+function keyword.guess_keyword_from_item(item_text, line, cursor_col, match_suffix)
+  local start, finish = keyword.guess_keyword_range_from_item(item_text, line, cursor_col, match_suffix)
+  return line:sub(start + 1, finish)
+end
+
+-- Example usage
+local function assert_range(expected, line, col, match_suffix)
+  local start, finish = keyword.get_keyword_range(line, col, match_suffix)
+  assert(
+    start == expected[1] and finish == expected[2],
+    'expected ' .. vim.inspect(expected) .. ' but got ' .. vim.inspect({ start, finish })
+  )
+end
+
+assert_range({ 1, 6 }, "'вest'", 6, false)
+assert_range({ 1, 4 }, "'вest'", 4, false)
+assert_range({ 1, 6 }, "'вest'", 6, true)
+
+return keyword

lua/blink/cmp/fuzzy/lua/match.lua

@@ -0,0 +1,78 @@
+local MATCH_SCORE = 7
+local GAP_PENALTY = -1
+
+-- bonus for matching the first character of the haystack
+local PREFIX_BONUS = 6
+-- bonus for matching character after a delimiter in the haystack (e.g. space, comma, underscore, slash, etc)
+local DELIMITER_BONUS = 4
+-- bonus for haystack == needle
+local EXACT_MATCH_BONUS = 4
+-- bonus for matching the case (upper or lower) of the haystack
+local MATCHING_CASE_BONUS = 1
+
+local DELIMITERS = {
+  [string.byte(' ', 1)] = true,
+  [string.byte('/', 1)] = true,
+  [string.byte('.', 1)] = true,
+  [string.byte(',', 1)] = true,
+  [string.byte('_', 1)] = true,
+  [string.byte('-', 1)] = true,
+  [string.byte(':', 1)] = true,
+}
+
+--- @param needle string
+--- @param haystack string
+--- @return number?, boolean?
+local function match(needle, haystack)
+  local score = 0
+  local haystack_idx = 1
+
+  for needle_idx = 1, #needle do
+    local needle_char = string.byte(needle, needle_idx)
+    local is_upper = needle_char >= 65 and needle_char <= 90
+    local is_lower = needle_char >= 97 and needle_char <= 122
+
+    local needle_lower_char = is_upper and needle_char + 32 or needle_char
+    local needle_upper_char = is_lower and needle_char - 32 or needle_char
+
+    local haystack_start_idx = haystack_idx
+    while haystack_idx <= (#haystack - #needle + needle_idx) do
+      local haystack_char = string.byte(haystack, haystack_idx)
+
+      if needle_lower_char == haystack_char or needle_upper_char == haystack_char then
+        score = score + MATCH_SCORE
+
+        -- gap penalty
+        if needle_idx ~= 1 then score = score + GAP_PENALTY * (haystack_idx - haystack_start_idx) end
+
+        -- bonuses
+        if needle_char == haystack_char then score = score + MATCHING_CASE_BONUS end
+        if haystack_idx == 1 then score = score + PREFIX_BONUS end
+        if DELIMITERS[string.byte(haystack, haystack_idx - 1)] then score = score + DELIMITER_BONUS end
+
+        haystack_idx = haystack_idx + 1
+        goto continue
+      end
+
+      haystack_idx = haystack_idx + 1
+    end
+
+    -- didn't find a match, so return nil
+    if true then return end
+
+    ::continue::
+  end
+
+  local exact = needle == haystack
+  if exact then score = score + EXACT_MATCH_BONUS end
+
+  return score, exact
+end
+
+assert(match('fbb', 'barbazfoobarbaz') == 20, 'fbb should match barbazfoobarbaz with score 18')
+assert(match('foo', '_foobar') == 28, 'foo should match foobar with score 29')
+assert(match('Foo', 'foobar') == 29, 'foo should match foobar with score 29')
+assert(match('foo', 'foobar') == 30, 'foo should match foobar with score 30')
+assert(match('foo', 'fobar') == nil, 'foo should not match fobar')
+
+return match

lua/blink/cmp/fuzzy/lua/match_indices.lua

@@ -0,0 +1,64 @@
+local MATCH_SCORE = 7
+local GAP_PENALTY = -1
+
+-- bonus for matching the first character of the haystack
+local PREFIX_BONUS = 6
+-- bonus for matching character after a delimiter in the haystack (e.g. space, comma, underscore, slash, etc)
+local DELIMITER_BONUS = 4
+-- bonus for haystack == needle
+local EXACT_MATCH_BONUS = 4
+-- bonus for matching the case (upper or lower) of the haystack
+local MATCHING_CASE_BONUS = 1
+
+local DELIMITERS = {
+  [string.byte(' ', 1)] = true,
+  [string.byte('/', 1)] = true,
+  [string.byte('.', 1)] = true,
+  [string.byte(',', 1)] = true,
+  [string.byte('_', 1)] = true,
+  [string.byte('-', 1)] = true,
+  [string.byte(':', 1)] = true,
+}
+
+--- @param needle string
+--- @param haystack string
+--- @return number[]
+local function match_indices(needle, haystack)
+  local matched_indices = {}
+
+  local haystack_idx = 1
+  for needle_idx = 1, #needle do
+    local needle_char = string.byte(needle, needle_idx)
+    local is_upper = needle_char >= 65 and needle_char <= 90
+    local is_lower = needle_char >= 97 and needle_char <= 122
+
+    local needle_lower_char = is_upper and needle_char + 32 or needle_char
+    local needle_upper_char = is_lower and needle_char - 32 or needle_char
+
+    while haystack_idx <= (#haystack - #needle + needle_idx) do
+      local haystack_char = string.byte(haystack, haystack_idx)
+
+      if needle_lower_char == haystack_char or needle_upper_char == haystack_char then
+        table.insert(matched_indices, haystack_idx - 1)
+        haystack_idx = haystack_idx + 1
+        break
+      end
+
+      haystack_idx = haystack_idx + 1
+    end
+  end
+
+  return matched_indices
+end
+
+local function assert_equal(expected, actual)
+  assert(vim.deep_equal(expected, actual), vim.inspect(expected) .. ' but got ' .. vim.inspect(actual))
+end
+
+assert_equal(match_indices('fbb', 'barbazfoobarbaz'), { 6, 9, 12 })
+assert(match_indices('foo', '_foobar'), { 1, 2, 3 })
+assert(match_indices('Foo', 'foobar'), { 0, 1, 2 })
+assert(match_indices('foo', 'foobar'), { 0, 1, 2 })
+assert(match_indices('foo', 'fobar'), { 0, 1 })
+
+return match_indices