diff --git a/crates/lsp-ai/src/memory_backends/file_store.rs b/crates/lsp-ai/src/memory_backends/file_store.rs index 9d55bca..a893c18 100644 --- a/crates/lsp-ai/src/memory_backends/file_store.rs +++ b/crates/lsp-ai/src/memory_backends/file_store.rs @@ -1,7 +1,7 @@ use anyhow::Context; use indexmap::IndexSet; use lsp_types::TextDocumentPositionParams; -use parking_lot::Mutex; +use parking_lot::{Mutex, RwLock}; use ropey::Rope; use serde_json::Value; use std::{collections::HashMap, io::Read}; @@ -49,7 +49,7 @@ impl File { pub(crate) struct FileStore { params: AdditionalFileStoreParams, - file_map: Mutex>, + file_map: RwLock>, accessed_files: Mutex>, crawl: Option>, } @@ -65,7 +65,7 @@ impl FileStore { .map(|x| Mutex::new(Crawl::new(x, config.clone()))); let s = Self { params: AdditionalFileStoreParams::default(), - file_map: Mutex::new(HashMap::new()), + file_map: RwLock::new(HashMap::new()), accessed_files: Mutex::new(IndexSet::new()), crawl, }; @@ -86,7 +86,7 @@ impl FileStore { .map(|x| Mutex::new(Crawl::new(x, config.clone()))); let s = Self { params, - file_map: Mutex::new(HashMap::new()), + file_map: RwLock::new(HashMap::new()), accessed_files: Mutex::new(IndexSet::new()), crawl, }; @@ -111,7 +111,7 @@ impl FileStore { None }; self.file_map - .lock() + .write() .insert(uri.to_string(), File::new(Rope::from_str(&contents), tree)); self.accessed_files.lock().insert(uri.to_string()); } @@ -130,7 +130,7 @@ impl FileStore { } // This means it has been opened before let insert_uri = format!("file:///{path}"); - if self.file_map.lock().contains_key(&insert_uri) { + if self.file_map.read().contains_key(&insert_uri) { return Ok(true); } // Open the file and see if it is small enough to read @@ -163,7 +163,7 @@ impl FileStore { let current_document_uri = position.text_document.uri.to_string(); let mut rope = self .file_map - .lock() + .read() .get(¤t_document_uri) .context("Error file not found")? .rope @@ -181,7 +181,7 @@ impl FileStore { if needed == 0 || !pull_from_multiple_files { break; } - let file_map = self.file_map.lock(); + let file_map = self.file_map.read(); let r = &file_map.get(file).context("Error file not found")?.rope; let slice_max = needed.min(r.len_chars() + 1); let rope_str_slice = r @@ -202,7 +202,7 @@ impl FileStore { ) -> anyhow::Result { let rope = self .file_map - .lock() + .read() .get(position.text_document.uri.as_str()) .context("Error file not found")? .rope @@ -275,19 +275,16 @@ impl FileStore { }) } - pub(crate) fn file_map(&self) -> &Mutex> { + pub fn file_map(&self) -> &RwLock> { &self.file_map } - pub(crate) fn contains_file(&self, uri: &str) -> bool { - self.file_map.lock().contains_key(uri) + pub fn contains_file(&self, uri: &str) -> bool { + self.file_map.read().contains_key(uri) } - pub(crate) fn position_to_byte( - &self, - position: &TextDocumentPositionParams, - ) -> anyhow::Result { - let file_map = self.file_map.lock(); + pub fn position_to_byte(&self, position: &TextDocumentPositionParams) -> anyhow::Result { + let file_map = self.file_map.read(); let uri = position.text_document.uri.to_string(); let file = file_map .get(&uri) @@ -307,7 +304,7 @@ impl MemoryBackend for FileStore { fn get_filter_text(&self, position: &TextDocumentPositionParams) -> anyhow::Result { let rope = self .file_map - .lock() + .read() .get(position.text_document.uri.as_str()) .context("Error file not found")? .rope @@ -351,7 +348,7 @@ impl MemoryBackend for FileStore { params: lsp_types::DidChangeTextDocumentParams, ) -> anyhow::Result<()> { let uri = params.text_document.uri.to_string(); - let mut file_map = self.file_map.lock(); + let mut file_map = self.file_map.write(); let file = file_map .get_mut(&uri) .with_context(|| format!("Trying to get file that does not exist {uri}"))?; @@ -450,7 +447,7 @@ impl MemoryBackend for FileStore { #[instrument(skip(self))] fn renamed_files(&self, params: lsp_types::RenameFilesParams) -> anyhow::Result<()> { for file_rename in params.files { - let mut file_map = self.file_map.lock(); + let mut file_map = self.file_map.write(); if let Some(rope) = file_map.remove(&file_rename.old_uri) { file_map.insert(file_rename.new_uri, rope); } @@ -537,7 +534,7 @@ mod tests { file_store.opened_text_document(params)?; let file = file_store .file_map - .lock() + .read() .get("file:///filler/") .unwrap() .clone(); @@ -563,7 +560,7 @@ mod tests { let file = file_store .file_map - .lock() + .read() .get("file:///filler2/") .unwrap() .clone(); @@ -604,7 +601,7 @@ mod tests { file_store.changed_text_document(params)?; let file = file_store .file_map - .lock() + .read() .get("file:///filler/") .unwrap() .clone(); @@ -624,7 +621,7 @@ mod tests { file_store.changed_text_document(params)?; let file = file_store .file_map - .lock() + .read() .get("file:///filler/") .unwrap() .clone(); @@ -881,7 +878,7 @@ fn main() { }], }; file_store.changed_text_document(params)?; - let file = file_store.file_map.lock().get(uri).unwrap().clone(); + let file = file_store.file_map.read().get(uri).unwrap().clone(); assert_eq!(file.tree.unwrap().root_node().to_sexp(), "(source_file (attribute_item (attribute (identifier) arguments: (token_tree (identifier)))) (struct_item name: (type_identifier) body: (field_declaration_list (field_declaration name: (field_identifier) type: (primitive_type)) (field_declaration name: (field_identifier) type: (primitive_type)))) (impl_item type: (type_identifier) body: (declaration_list (function_item name: (identifier) parameters: (parameters (self_parameter (self))) return_type: (primitive_type) body: (block (binary_expression left: (field_expression value: (self) field: (field_identifier)) right: (field_expression value: (self) field: (field_identifier))))))) (function_item name: (identifier) parameters: (parameters) body: (block (let_declaration pattern: (identifier) value: (struct_expression name: (type_identifier) body: (field_initializer_list (field_initializer field: (field_identifier) value: (integer_literal)) (field_initializer field: (field_identifier) value: (integer_literal))))) (expression_statement (macro_invocation macro: (identifier) (token_tree (string_literal (string_content)) (identifier) (identifier) (token_tree)))))))"); // Test delete @@ -906,7 +903,7 @@ fn main() { }], }; file_store.changed_text_document(params)?; - let file = file_store.file_map.lock().get(uri).unwrap().clone(); + let file = file_store.file_map.read().get(uri).unwrap().clone(); assert_eq!(file.tree.unwrap().root_node().to_sexp(), "(source_file (function_item name: (identifier) parameters: (parameters) body: (block (let_declaration pattern: (identifier) value: (struct_expression name: (type_identifier) body: (field_initializer_list (field_initializer field: (field_identifier) value: (integer_literal)) (field_initializer field: (field_identifier) value: (integer_literal))))) (expression_statement (macro_invocation macro: (identifier) (token_tree (string_literal (string_content)) (identifier) (identifier) (token_tree)))))))"); // Test replace @@ -922,7 +919,7 @@ fn main() { }], }; file_store.changed_text_document(params)?; - let file = file_store.file_map.lock().get(uri).unwrap().clone(); + let file = file_store.file_map.read().get(uri).unwrap().clone(); assert_eq!(file.tree.unwrap().root_node().to_sexp(), "(source_file (function_item name: (identifier) parameters: (parameters) body: (block)))"); Ok(()) diff --git a/crates/lsp-ai/src/memory_backends/postgresml/mod.rs b/crates/lsp-ai/src/memory_backends/postgresml/mod.rs index 88da39c..2168101 100644 --- a/crates/lsp-ai/src/memory_backends/postgresml/mod.rs +++ b/crates/lsp-ai/src/memory_backends/postgresml/mod.rs @@ -50,7 +50,7 @@ async fn split_and_upsert_file( let chunks = { file_store .file_map() - .lock() + .read() .get(uri) .map(|f| splitter.split(f)) }; @@ -184,7 +184,7 @@ impl PostgresML { let chunks: Vec> = match file_uris .iter() .map(|uri| { - let file_store = task_file_store.file_map().lock(); + let file_store = task_file_store.file_map().read(); let file = file_store .get(uri) .with_context(|| format!("getting file for splitting: {uri}"))?; diff --git a/crates/lsp-ai/src/memory_backends/vector_store.rs b/crates/lsp-ai/src/memory_backends/vector_store.rs index 52c6765..0bf9339 100644 --- a/crates/lsp-ai/src/memory_backends/vector_store.rs +++ b/crates/lsp-ai/src/memory_backends/vector_store.rs @@ -226,7 +226,8 @@ impl VS { if acc.is_empty() { acc.insert(score, chunk); } else if acc.first_key_value().unwrap().0 < &score { - if acc.len() == limit { + // We want to get limit + 1 here in case the limit is 1 and then we filter the chunk out later + if acc.len() == limit + 1 { acc.pop_first(); } acc.insert(score, chunk); @@ -395,7 +396,7 @@ impl MemoryBackend for VectorStore { let uri = params.text_document.uri.to_string(); self.file_store.opened_text_document(params)?; - let file_map = self.file_store.file_map().lock(); + let file_map = self.file_store.file_map().read(); let file = file_map.get(&uri).context("file not found")?; let chunks = self.splitter.split(file); self.upsert_chunks(&uri, chunks); @@ -411,9 +412,11 @@ impl MemoryBackend for VectorStore { let uri = params.text_document.uri.to_string(); self.file_store.changed_text_document(params.clone())?; - let file_map = self.file_store.file_map().lock(); - let file = file_map.get(&uri).context("file not found")?; - let chunks = self.splitter.split(file); + let chunks = { + let file_map = self.file_store.file_map().read(); + let file = file_map.get(&uri).context("file not found")?; + self.splitter.split(file) + }; let chunks_size = chunks.len(); // This is not as efficient as it could be, but it is ok for now