Skip to content

Commit

Permalink
find refs: incrementally update word_index after file changes
Browse files Browse the repository at this point in the history
Signed-off-by: xiarui.xr <[email protected]>
  • Loading branch information
amyXia1994 committed Oct 12, 2023
1 parent 55040b7 commit 85771f4
Show file tree
Hide file tree
Showing 7 changed files with 279 additions and 43 deletions.
20 changes: 13 additions & 7 deletions kclvm/tools/src/LSP/src/find_refs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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);
Expand Down Expand Up @@ -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> {
Expand All @@ -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(),
)])
}
Expand Down Expand Up @@ -140,6 +144,7 @@ mod tests {
check_locations_match(
expect,
find_refs(
None,
setup_word_index_map(path),
def_loc,
"a".to_string(),
Expand Down Expand Up @@ -193,6 +198,7 @@ mod tests {
check_locations_match(
expect,
find_refs(
None,
setup_word_index_map(path),
def_loc,
"Name".to_string(),
Expand Down
19 changes: 11 additions & 8 deletions kclvm/tools/src/LSP/src/goto_def.rs
Original file line number Diff line number Diff line change
Expand Up @@ -350,14 +350,17 @@ pub(crate) fn resolve_var(
if let Some(ty) = STRING_MEMBER_FUNCTIONS.get(&func_name) {
match &ty.kind {
kclvm_sema::ty::TypeKind::Function(func_ty) => {
return Some(Definition::Object(ScopeObject {
name: func_name,
start: func_name_node.get_pos(),
end: func_name_node.get_end_pos(),
ty: Rc::new(ty.clone()),
kind: ScopeObjectKind::FunctionCall,
doc: Some(func_ty.doc.clone()),
}, func_name))
return Some(Definition::Object(
ScopeObject {
name: func_name.clone(),
start: func_name_node.get_pos(),
end: func_name_node.get_end_pos(),
ty: Rc::new(ty.clone()),
kind: ScopeObjectKind::FunctionCall,
doc: Some(func_ty.doc.clone()),
},
func_name,
))
}
// unreachable
_ => {}
Expand Down
47 changes: 43 additions & 4 deletions kclvm/tools/src/LSP/src/notification.rs
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down Expand Up @@ -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(())
}
Expand Down
9 changes: 8 additions & 1 deletion kclvm/tools/src/LSP/src/request.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
9 changes: 5 additions & 4 deletions kclvm/tools/src/LSP/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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)]
Expand All @@ -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 {
Expand Down
101 changes: 99 additions & 2 deletions kclvm/tools/src/LSP/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -771,7 +771,7 @@ fn hover_assign_in_lambda_test() {

let path = path.to_str().unwrap();
let src = std::fs::read_to_string(path.clone()).unwrap();
let server = Project {}.server();
let server = Project {}.server(InitializeParams::default());

// Mock open file
server.notification::<lsp_types::notification::DidOpenTextDocument>(
Expand Down Expand Up @@ -1310,7 +1310,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");
Expand Down Expand Up @@ -1394,3 +1394,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()
);
}
Loading

0 comments on commit 85771f4

Please sign in to comment.