find refs: incrementally update word_index after file changes

kclvm/tools/src/LSP/src/find_refs.rs

@@ -2,11 +2,16 @@ use crate::from_lsp::kcl_pos;
 use crate::goto_def::goto_definition;
 use crate::util::{parse_param_and_compile, Param};
 use anyhow;
-use lsp_types::Location;
+use lsp_types::{Location, Url};
+use parking_lot::RwLock;
+use ra_ap_vfs::Vfs;
+use std::clone;
 use std::collections::HashMap;
+use std::sync::Arc;
 
 pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
-    word_index_map: HashMap<String, HashMap<String, Vec<Location>>>,
+    vfs: Option<Arc<RwLock<Vfs>>>,
+    word_index_map: HashMap<Url, HashMap<String, Vec<Location>>>,
     def_loc: Location,
     name: String,
     cursor_path: String,
@@ -29,7 +34,7 @@ pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
                         Param {
                             file: file_path.clone(),
                         },
-                        None,
+                        vfs.clone(),
                     ) {
                         Ok((prog, scope, _)) => {
                             let ref_pos = kcl_pos(&file_path, ref_loc.range.start);
@@ -62,9 +67,8 @@ pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
 mod tests {
     use super::find_refs;
     use crate::util::build_word_index;
-    use lsp_types::{Location, Position, Range};
+    use lsp_types::{Location, Position, Range, Url};
     use std::collections::HashMap;
-    use std::ops::Index;
     use std::path::PathBuf;
 
     fn logger(msg: String) -> Result<(), anyhow::Error> {
@@ -85,9 +89,9 @@ mod tests {
         }
     }
 
-    fn setup_word_index_map(root: &str) -> HashMap<String, HashMap<String, Vec<Location>>> {
+    fn setup_word_index_map(root: &str) -> HashMap<Url, HashMap<String, Vec<Location>>> {
         HashMap::from([(
-            "default".to_string(),
+            Url::from_file_path(root).unwrap(),
             build_word_index(root.to_string()).unwrap(),
         )])
     }
@@ -140,6 +144,7 @@ mod tests {
                 check_locations_match(
                     expect,
                     find_refs(
+                        None,
                         setup_word_index_map(path),
                         def_loc,
                         "a".to_string(),
@@ -193,6 +198,7 @@ mod tests {
                 check_locations_match(
                     expect,
                     find_refs(
+                        None,
                         setup_word_index_map(path),
                         def_loc,
                         "Name".to_string(),

kclvm/tools/src/LSP/src/notification.rs

@@ -1,11 +1,18 @@
 use lsp_types::notification::{
-    DidChangeTextDocument, DidChangeWatchedFiles, DidCloseTextDocument, DidOpenTextDocument,
-    DidSaveTextDocument,
+    DidChangeTextDocument,
+    DidChangeWatchedFiles,
+    DidCloseTextDocument,
+    DidOpenTextDocument,
+    DidSaveTextDocument, //DidDeleteFiles, DidRenameFiles, DidCreateFiles, //todo more
 };
+use std::path::Path;
 
 use crate::{
-    dispatcher::NotificationDispatcher, from_lsp, state::LanguageServerState,
+    dispatcher::NotificationDispatcher,
+    from_lsp,
+    state::LanguageServerState,
     util::apply_document_changes,
+    util::{build_word_index_for_file_content, word_index_add, word_index_subtract},
 };
 
 impl LanguageServerState {
@@ -71,14 +78,46 @@ impl LanguageServerState {
         let path = from_lsp::abs_path(&text_document.uri)?;
         self.log_message(format!("on did_change file: {:?}", path));
 
+        // update vfs
         let vfs = &mut *self.vfs.write();
         let file_id = vfs
             .file_id(&path.clone().into())
             .ok_or(anyhow::anyhow!("Already checked that the file_id exists!"))?;
 
         let mut text = String::from_utf8(vfs.file_contents(file_id).to_vec())?;
+        let old_text = text.clone();
         apply_document_changes(&mut text, content_changes);
-        vfs.set_file_contents(path.into(), Some(text.into_bytes()));
+        vfs.set_file_contents(path.into(), Some(text.clone().into_bytes()));
+
+        // update word index
+        let old_word_index = build_word_index_for_file_content(old_text, &text_document.uri);
+        let new_word_index = build_word_index_for_file_content(text.clone(), &text_document.uri);
+        let binding = text_document.uri.path();
+        let file_path = Path::new(binding); //todo rename
+        for (key, value) in &mut self.word_index_map {
+            let workspace_folder_path = Path::new(key.path());
+            if file_path.starts_with(workspace_folder_path) {
+                word_index_subtract(value, old_word_index.clone());
+                word_index_add(value, new_word_index.clone());
+            }
+        }
+        // let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+        // let url = lsp_types::Url::from_file_path(root.clone()).unwrap();
+        // let mm = self.word_index_map.get(&url).unwrap();
+        // println!("word_index_map: {:?}", mm);
+
+        // let file = from_lsp::file_path_from_url(&text_document.uri)?;
+        // let old_word_index = build_word_index_for_file_content(old_text, &text_document.uri);
+        // let new_word_index = build_word_index_for_file_content(text.clone(), &text_document.uri);
+
+        // let file_path = Path::new(&text_document.uri.path());
+        // for (key, mut value) in &self.word_index_map {
+        //     let workspace_folder_path = Path::new(key.path());
+        //     if file_path.starts_with(workspace_folder_path) {
+        //         value = &word_index_subtract(value, old_word_index.clone());
+        //         value = &word_index_add(value, new_word_index.clone());
+        //     }
+        // }
 
         Ok(())
     }

kclvm/tools/src/LSP/src/request.rs

@@ -164,7 +164,14 @@ pub(crate) fn handle_reference(
                     },
                     None => None,
                 } {
-                    return find_refs(word_index_map, def_loc, def_name, file, log);
+                    return find_refs(
+                        Some(snapshot.vfs),
+                        word_index_map,
+                        def_loc,
+                        def_name,
+                        file,
+                        log,
+                    );
                 }
             }
             _ => return Ok(None),

kclvm/tools/src/LSP/src/state.rs

@@ -6,6 +6,7 @@ use crate::util::{build_word_index, get_file_name, parse_param_and_compile, to_j
 use crossbeam_channel::{select, unbounded, Receiver, Sender};
 use indexmap::IndexSet;
 use lsp_server::{ReqQueue, Response};
+use lsp_types::Url;
 use lsp_types::{
     notification::{Notification, PublishDiagnostics},
     Diagnostic, InitializeParams, Location, PublishDiagnosticsParams,
@@ -69,7 +70,7 @@ pub(crate) struct LanguageServerState {
     pub vfs_handle: Box<dyn ra_ap_vfs::loader::Handle>,
 
     /// The word index map
-    pub word_index_map: HashMap<String, HashMap<String, Vec<Location>>>,
+    pub word_index_map: HashMap<Url, HashMap<String, Vec<Location>>>,
 }
 
 /// A snapshot of the state of the language server
@@ -82,7 +83,7 @@ pub(crate) struct LanguageServerSnapshot {
     /// Documents that are currently kept in memory from the client
     pub opened_files: IndexSet<FileId>,
     /// The word index map
-    pub word_index_map: HashMap<String, HashMap<String, Vec<Location>>>,
+    pub word_index_map: HashMap<Url, HashMap<String, Vec<Location>>>,
 }
 
 #[allow(unused)]
@@ -106,13 +107,13 @@ impl LanguageServerState {
             for folder in workspace_folders {
                 let path = folder.uri.path();
                 if let Ok(word_index) = build_word_index(path.to_string()) {
-                    word_index_map.insert(folder.name, word_index);
+                    word_index_map.insert(folder.uri, word_index);
                 }
             }
         } else if let Some(root_uri) = initialize_params.root_uri {
             let path = root_uri.path();
             if let Ok(word_index) = build_word_index(path.to_string()) {
-                word_index_map.insert("default".to_string(), word_index);
+                word_index_map.insert(root_uri, word_index);
             }
         }
         LanguageServerState {

kclvm/tools/src/LSP/src/tests.rs

@@ -1257,7 +1257,7 @@ fn lsp_invalid_subcommand_test() {
 }
 
 #[test]
-fn test_find_refs() {
+fn find_refs_test() {
     let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
     let mut path = root.clone();
     path.push("src/test_data/find_refs_test/main.k");
@@ -1341,3 +1341,100 @@ fn test_find_refs() {
         .unwrap()
     );
 }
+
+#[test]
+fn find_refs_with_file_change_test() {
+    let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    let mut path = root.clone();
+    path.push("src/test_data/find_refs_test/main.k");
+
+    let path = path.to_str().unwrap();
+    let src = std::fs::read_to_string(path.clone()).unwrap();
+    let mut initialize_params = InitializeParams::default();
+    initialize_params.workspace_folders = Some(vec![WorkspaceFolder {
+        uri: Url::from_file_path(root.clone()).unwrap(),
+        name: "test".to_string(),
+    }]);
+    let server = Project {}.server(initialize_params);
+    let url = Url::from_file_path(path).unwrap();
+
+    // Mock open file
+    server.notification::<lsp_types::notification::DidOpenTextDocument>(
+        lsp_types::DidOpenTextDocumentParams {
+            text_document: TextDocumentItem {
+                uri: url.clone(),
+                language_id: "KCL".to_string(),
+                version: 0,
+                text: src,
+            },
+        },
+    );
+    // Mock change file content
+    server.notification::<lsp_types::notification::DidChangeTextDocument>(
+        lsp_types::DidChangeTextDocumentParams {
+            text_document: lsp_types::VersionedTextDocumentIdentifier {
+                uri: url.clone(),
+                version: 1,
+            },
+            content_changes: vec![lsp_types::TextDocumentContentChangeEvent {
+                range: None,
+                range_length: None,
+                text: r#"a = "demo"
+
+schema Name:
+    name: str
+
+schema Person:
+    n: Name
+
+p2 = Person {
+    n: Name{
+        name: a
+    }
+}"#
+                .to_string(),
+            }],
+        },
+    );
+    let id = server.next_request_id.get();
+    server.next_request_id.set(id.wrapping_add(1));
+    // Mock trigger find references
+    let r: Request = Request::new(
+        id.into(),
+        "textDocument/references".to_string(),
+        ReferenceParams {
+            text_document_position: TextDocumentPositionParams {
+                text_document: TextDocumentIdentifier { uri: url.clone() },
+                position: Position::new(0, 1),
+            },
+            work_done_progress_params: Default::default(),
+            partial_result_params: Default::default(),
+            context: ReferenceContext {
+                include_declaration: true,
+            },
+        },
+    );
+
+    // Send request and wait for it's response
+    let res = server.send_and_receive(r);
+    assert_eq!(
+        res.result.unwrap(),
+        to_json(vec![
+            Location {
+                uri: url.clone(),
+                range: Range {
+                    start: Position::new(0, 0),
+                    end: Position::new(0, 1),
+                },
+            },
+            Location {
+                uri: url.clone(),
+                range: Range {
+                    start: Position::new(10, 14),
+                    end: Position::new(10, 15),
+                },
+            },
+        ])
+        .unwrap()
+    );
+}