feat(fuzzy): lazy get lua properties

Saghen · Saghen · commit 5cc63f0f298a · 2024-10-23T20:51:43.000-04:00
Significantly improver performance but much of the sorting and filtering
logic should likely be moved to the lua side for better flexibility and
performance
diff --git a/lua/blink/cmp/fuzzy/frecency.rs b/lua/blink/cmp/fuzzy/frecency.rs
@@ -1,4 +1,4 @@
-use crate::fuzzy::LspItem;
+use crate::lsp_item::LspItem;
 use heed::types::*;
 use heed::{Database, Env, EnvOpenOptions};
 use mlua::Result as LuaResult;
diff --git a/lua/blink/cmp/fuzzy/fuzzy.rs b/lua/blink/cmp/fuzzy/fuzzy.rs
@@ -1,62 +1,15 @@
 // TODO: refactor this heresy
 
 use crate::frecency::FrecencyTracker;
+use crate::lsp_item::LspItem;
 use mlua::prelude::*;
 use mlua::FromLua;
 use mlua::Lua;
-use serde::{Deserialize, Serialize};
+use std::cell::LazyCell;
 use std::cmp::Reverse;
 use std::collections::HashSet;
 
-#[derive(Clone, Serialize, Deserialize, Debug)]
-pub struct LspItem {
-    pub label: String,
-    #[serde(rename = "sortText")]
-    pub sort_text: Option<String>,
-    #[serde(rename = "filterText")]
-    pub filter_text: Option<String>,
-    pub kind: u32,
-    pub score_offset: Option<i32>,
-    pub source_id: String,
-}
-
-impl FromLua<'_> for LspItem {
-    fn from_lua(value: LuaValue<'_>, _lua: &'_ Lua) -> LuaResult<Self> {
-        if let Some(tab) = value.as_table() {
-            let label: String = tab.get("label").unwrap_or_default();
-            let sort_text: Option<String> = tab.get("sortText").ok();
-            let filter_text: Option<String> = tab.get("filterText").ok();
-            let kind: u32 = tab.get("kind").unwrap_or_default();
-            let score_offset: Option<i32> = tab.get("score_offset").ok();
-            let source_id: String = tab.get("source_id").unwrap_or_default();
-
-            Ok(LspItem {
-                label,
-                sort_text,
-                filter_text,
-                kind,
-                score_offset,
-                source_id,
-            })
-        } else {
-            Err(mlua::Error::FromLuaConversionError {
-                from: "LuaValue",
-                to: "LspItem",
-                message: None,
-            })
-        }
-    }
-}
-
-#[derive(Clone, Serialize, Deserialize)]
-pub struct MatchedLspItem {
-    label: String,
-    kind: u32,
-    index: u32,
-    score: i32,
-}
-
-#[derive(Clone, Serialize, Deserialize, Hash)]
+#[derive(Clone, Hash)]
 pub struct FuzzyOptions {
     use_typo_resistance: bool,
     use_frecency: bool,
@@ -99,30 +52,33 @@ impl FromLua<'_> for FuzzyOptions {
 
 pub fn fuzzy(
     needle: String,
-    haystack: Vec<LspItem>,
+    haystack_labels: Vec<String>,
+    haystack: Vec<LuaTable>,
     frecency: &FrecencyTracker,
     opts: FuzzyOptions,
 ) -> Vec<usize> {
+    let haystack = haystack
+        .into_iter()
+        .map(|item| LazyCell::new(move || -> LspItem { item.into() }))
+        .collect::<Vec<_>>();
+
     let nearby_words: HashSet<String> = HashSet::from_iter(opts.nearby_words.unwrap_or_default());
 
     // Fuzzy match with fzrs
-    let haystack_labels = haystack
-        .iter()
-        .map(|s| {
-            if let Some(filter_text) = &s.filter_text {
-                filter_text.as_str()
-            } else {
-                s.label.as_str()
-            }
-        })
-        .collect::<Vec<_>>();
     let options = frizbee::Options {
         prefilter: !opts.use_typo_resistance,
         min_score: opts.min_score,
         stable_sort: false,
         ..Default::default()
     };
-    let mut matches = frizbee::match_list(&needle, &haystack_labels, options);
+    let mut matches = frizbee::match_list(
+        &needle,
+        &haystack_labels
+            .iter()
+            .map(|s| s.as_str())
+            .collect::<Vec<_>>(),
+        options,
+    );
 
     // Sort by scores
     let match_scores = matches
@@ -135,13 +91,13 @@ pub fn fuzzy(
             };
             let nearby_words_score = if opts.use_proximity {
                 nearby_words
-                    .get(&haystack[mtch.index_in_haystack].label)
+                    .get(&haystack_labels[mtch.index_in_haystack])
                     .map(|_| 2)
                     .unwrap_or(0)
             } else {
                 0
             };
-            let score_offset = haystack[mtch.index_in_haystack].score_offset.unwrap_or(0);
+            let score_offset = haystack[mtch.index_in_haystack].score_offset;
 
             (mtch.score as i32) + frecency_score + nearby_words_score + score_offset
         })
diff --git a/lua/blink/cmp/fuzzy/init.lua b/lua/blink/cmp/fuzzy/init.lua
@@ -35,12 +35,15 @@ function fuzzy.filter_items(needle, haystack)
   local cursor_row = vim.api.nvim_win_get_cursor(0)[1]
   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_words =
-    fuzzy.rust.get_words(table.concat(vim.api.nvim_buf_get_lines(0, start_row, end_row, false), '\n'))
+  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 {}
 
   -- perform fuzzy search
-  local filtered_items = {}
-  local matched_indices = fuzzy.rust.fuzzy(needle, haystack, {
+  local haystack_filter_text = {}
+  for _, item in ipairs(haystack) do
+    table.insert(haystack_filter_text, item.label)
+  end
+  local matched_indices = fuzzy.rust.fuzzy(needle, haystack_filter_text, haystack, {
     -- each matching char is worth 4 points and it receives a bonus for capitalization, delimiter and prefix
     -- so this should generally be good
     -- TODO: make this configurable
@@ -53,10 +56,10 @@ function fuzzy.filter_items(needle, haystack)
     nearby_words = nearby_words,
   })
 
+  local filtered_items = {}
   for _, idx in ipairs(matched_indices) do
     table.insert(filtered_items, haystack[idx + 1])
   end
-
   return filtered_items
 end
 
diff --git a/lua/blink/cmp/fuzzy/lib.rs b/lua/blink/cmp/fuzzy/lib.rs
@@ -1,5 +1,6 @@
 use crate::frecency::FrecencyTracker;
-use crate::fuzzy::{FuzzyOptions, LspItem};
+use crate::fuzzy::FuzzyOptions;
+use crate::lsp_item::LspItem;
 use lazy_static::lazy_static;
 use mlua::prelude::*;
 use regex::Regex;
@@ -8,6 +9,7 @@ use std::sync::RwLock;
 
 mod frecency;
 mod fuzzy;
+mod lsp_item;
 
 lazy_static! {
     static ref REGEX: Regex = Regex::new(r"[A-Za-z][A-Za-z0-9_\\-]{2,32}").unwrap();
@@ -52,7 +54,12 @@ pub fn access(_: &Lua, item: LspItem) -> LuaResult<bool> {
 
 pub fn fuzzy(
     _lua: &Lua,
-    (needle, haystack, opts): (String, Vec<LspItem>, FuzzyOptions),
+    (needle, haystack_filter_text, haystack, opts): (
+        String,
+        Vec<String>,
+        Vec<LuaTable>,
+        FuzzyOptions,
+    ),
 ) -> LuaResult<Vec<u32>> {
     let mut frecency_handle = FRECENCY.write().map_err(|_| {
         mlua::Error::RuntimeError("Failed to acquire lock for frecency".to_string())
@@ -61,10 +68,12 @@ pub fn fuzzy(
         mlua::Error::RuntimeError("Attempted to use frencecy before initialization".to_string())
     })?;
 
-    Ok(fuzzy::fuzzy(needle, haystack, frecency, opts)
-        .into_iter()
-        .map(|i| i as u32)
-        .collect())
+    Ok(
+        fuzzy::fuzzy(needle, haystack_filter_text, haystack, frecency, opts)
+            .into_iter()
+            .map(|i| i as u32)
+            .collect(),
+    )
 }
 
 pub fn get_words(_: &Lua, text: String) -> LuaResult<Vec<String>> {
@@ -76,13 +85,15 @@ pub fn get_words(_: &Lua, text: String) -> LuaResult<Vec<String>> {
         .collect())
 }
 
-#[mlua::lua_module]
+// NOTE: skip_memory_check greatly improves performance
+// https://github.com/mlua-rs/mlua/issues/318
+#[mlua::lua_module(skip_memory_check)]
 fn blink_cmp_fuzzy(lua: &Lua) -> LuaResult<LuaTable> {
     let exports = lua.create_table()?;
     exports.set("fuzzy", lua.create_function(fuzzy)?)?;
     exports.set("get_words", lua.create_function(get_words)?)?;
     exports.set("init_db", lua.create_function(init_db)?)?;
     exports.set("destroy_db", lua.create_function(destroy_db)?)?;
-    exports.set("access", lua.create_function(access)?)?;
+    // exports.set("access", lua.create_function(access)?)?;
     Ok(exports)
 }
diff --git a/lua/blink/cmp/fuzzy/lsp_item.rs b/lua/blink/cmp/fuzzy/lsp_item.rs
@@ -0,0 +1,90 @@
+use mlua::prelude::*;
+
+#[derive(Debug)]
+pub struct LspItem {
+    pub label: String,
+    pub kind: u32,
+    pub score_offset: i32,
+    pub source_id: String,
+}
+
+impl From<LuaTable<'_>> for LspItem {
+    fn from(tab: LuaTable) -> Self {
+        let label: String = tab.get("label").unwrap_or_default();
+        let kind: u32 = tab.get("kind").unwrap_or_default();
+        let score_offset: i32 = tab.get("score_offset").unwrap_or(0);
+        let source_id: String = tab.get("source_id").unwrap_or_default();
+
+        LspItem {
+            label,
+            kind,
+            score_offset,
+            source_id,
+        }
+    }
+}
+
+// #[derive(Debug)]
+// pub struct LspItem<'lua> {
+//     table: LuaTable<'lua>,
+//     pub label: OnceCell<String>,
+//     pub sort_text: OnceCell<Option<String>>,
+//     pub filter_text: OnceCell<String>,
+//     pub kind: OnceCell<u32>,
+//     pub score_offset: OnceCell<i32>,
+//     pub source_id: OnceCell<String>,
+// }
+//
+// impl LspItem<'_> {
+//     pub fn label(&self) -> String {
+//         self.label
+//             .get_or_init(|| self.table.get::<_, String>("label").unwrap_or_default())
+//             .to_owned()
+//     }
+//     pub fn sort_text(&self) -> Option<String> {
+//         self.sort_text
+//             .get_or_init(|| self.table.get("sortText").ok())
+//             .to_owned()
+//     }
+//     pub fn filter_text(&self) -> String {
+//         self.filter_text
+//             .get_or_init(|| self.table.get("filterText").unwrap_or(self.label()))
+//             .to_owned()
+//     }
+//     pub fn kind(&self) -> u32 {
+//         self.kind
+//             .get_or_init(|| self.table.get("kind").unwrap_or_default())
+//             .to_owned()
+//     }
+//     pub fn score_offset(&self) -> i32 {
+//         self.score_offset
+//             .get_or_init(|| self.table.get("score_offset").unwrap_or(0))
+//             .to_owned()
+//     }
+//     pub fn source_id(&self) -> String {
+//         self.source_id
+//             .get_or_init(|| self.table.get("source_id").unwrap_or_default())
+//             .to_owned()
+//     }
+// }
+//
+// impl<'lua> FromLua<'lua> for LspItem<'lua> {
+//     fn from_lua(value: LuaValue<'lua>, _lua: &'lua Lua) -> LuaResult<Self> {
+//         match value {
+//             LuaValue::Table(tab) => Ok(LspItem {
+//                 table: tab,
+//                 label: OnceCell::new(),
+//                 sort_text: OnceCell::new(),
+//                 filter_text: OnceCell::new(),
+//                 kind: OnceCell::new(),
+//                 score_offset: OnceCell::new(),
+//                 source_id: OnceCell::new(),
+//             }),
+//             _ => Err(mlua::Error::FromLuaConversionError {
+//                 from: "LuaValue",
+//                 to: "LspItem",
+//                 message: None,
+//             }),
+//         }
+//     }
+// }

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-use crate::fuzzy::LspItem;`
	`1`	`+use crate::lsp_item::LspItem;`
`2`	`2`	`use heed::types::*;`
`3`	`3`	`use heed::{Database, Env, EnvOpenOptions};`
`4`	`4`	`use mlua::Result as LuaResult;`