feat: rename symbol

kclvm/tools/src/LSP/Cargo.toml

@@ -19,6 +19,7 @@ log = "0.4.14"
 im-rc = "15.0.0"
 rustc_lexer = "0.1.0"
 clap = "4.3.0"
+regex = "1.4"
 
 kclvm-tools = { path = "../../../tools" }
 kclvm-error = { path = "../../../error" }

kclvm/tools/src/LSP/src/capabilities.rs

@@ -39,6 +39,7 @@ pub fn server_capabilities(client_caps: &ClientCapabilities) -> ServerCapabiliti
         document_formatting_provider: Some(OneOf::Left(true)),
         document_range_formatting_provider: Some(OneOf::Left(true)),
         references_provider: Some(OneOf::Left(true)),
+        rename_provider: Some(OneOf::Left(true)),
         ..Default::default()
     }
 }

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

@@ -1,23 +1,64 @@
 use crate::from_lsp::kcl_pos;
-use crate::goto_def::goto_definition;
+use crate::goto_def::{find_def, goto_definition};
+use crate::to_lsp::lsp_location;
 use crate::util::{parse_param_and_compile, Param};
 use anyhow;
+use kclvm_ast::ast::{Program, Stmt};
+use kclvm_error::Position as KCLPos;
+use kclvm_sema::resolver::scope::ProgramScope;
 use lsp_types::{Location, Url};
 use parking_lot::RwLock;
 use ra_ap_vfs::Vfs;
 use std::collections::HashMap;
 use std::sync::Arc;
 
 pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
+    program: &Program,
+    kcl_pos: &KCLPos,
+    prog_scope: &ProgramScope,
+    word_index_map: HashMap<Url, HashMap<String, Vec<Location>>>,
+    vfs: Option<Arc<RwLock<Vfs>>>,
+    logger: F,
+) -> Result<Vec<Location>, String> {
+    // find the definition at the position
+    let def = match program.pos_to_stmt(kcl_pos) {
+        Some(node) => match node.node {
+            Stmt::Import(_) => None,
+            _ => find_def(node.clone(), kcl_pos, prog_scope),
+        },
+        None => None,
+    };
+    if def.is_none() {
+        return Err(String::from(
+            "Definition item not found, result in no reference",
+        ));
+    }
+    let def = def.unwrap();
+    if def.get_positions().len() != 1 {
+        return Err(String::from(
+            "Found more than one definitions, reference not supported",
+        ));
+    }
+    let (start, end) = def.get_positions().iter().next().unwrap().clone();
+    let def_loc = lsp_location(start.filename.clone(), &start, &end);
+    // find all the refs of the def
+    Ok(find_refs_from_def(
+        vfs,
+        word_index_map,
+        def_loc,
+        def.get_name(),
+        logger,
+    ))
+}
+
+pub(crate) fn find_refs_from_def<F: Fn(String) -> Result<(), anyhow::Error>>(
     vfs: Option<Arc<RwLock<Vfs>>>,
     word_index_map: HashMap<Url, HashMap<String, Vec<Location>>>,
     def_loc: Location,
     name: String,
-    cursor_path: String,
     logger: F,
-) -> anyhow::Result<Option<Vec<Location>>> {
+) -> Vec<Location> {
     let mut ref_locations = vec![];
-
     for (_, word_index) in word_index_map {
         if let Some(locs) = word_index.get(name.as_str()).cloned() {
             let matched_locs: Vec<Location> = locs
@@ -47,7 +88,8 @@ pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
                             }
                         }
                         Err(_) => {
-                            let _ = logger(format!("{cursor_path} compilation failed"));
+                            let file_path = def_loc.uri.path();
+                            let _ = logger(format!("{file_path} compilation failed"));
                             return false;
                         }
                     }
@@ -56,12 +98,12 @@ pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
             ref_locations.extend(matched_locs);
         }
     }
-    anyhow::Ok(Some(ref_locations))
+    ref_locations
 }
 
 #[cfg(test)]
 mod tests {
-    use super::find_refs;
+    use super::find_refs_from_def;
     use crate::util::build_word_index;
     use lsp_types::{Location, Position, Range, Url};
     use std::collections::HashMap;
@@ -72,17 +114,8 @@ mod tests {
         anyhow::Ok(())
     }
 
-    fn check_locations_match(expect: Vec<Location>, actual: anyhow::Result<Option<Vec<Location>>>) {
-        match actual {
-            Ok(act) => {
-                if let Some(locations) = act {
-                    assert_eq!(expect, locations)
-                } else {
-                    assert!(false, "got empty result. expect: {:?}", expect)
-                }
-            }
-            Err(_) => assert!(false),
-        }
+    fn check_locations_match(expect: Vec<Location>, actual: Vec<Location>) {
+        assert_eq!(expect, actual)
     }
 
     fn setup_word_index_map(root: &str) -> HashMap<Url, HashMap<String, Vec<Location>>> {
@@ -139,12 +172,11 @@ mod tests {
                 ];
                 check_locations_match(
                     expect,
-                    find_refs(
+                    find_refs_from_def(
                         None,
                         setup_word_index_map(path),
                         def_loc,
                         "a".to_string(),
-                        path.to_string(),
                         logger,
                     ),
                 );
@@ -193,12 +225,11 @@ mod tests {
                 ];
                 check_locations_match(
                     expect,
-                    find_refs(
+                    find_refs_from_def(
                         None,
                         setup_word_index_map(path),
                         def_loc,
                         "Name".to_string(),
-                        path.to_string(),
                         logger,
                     ),
                 );
@@ -240,12 +271,11 @@ mod tests {
                 ];
                 check_locations_match(
                     expect,
-                    find_refs(
+                    find_refs_from_def(
                         None,
                         setup_word_index_map(path),
                         def_loc,
                         "name".to_string(),
-                        path.to_string(),
                         logger,
                     ),
                 );

kclvm/tools/src/LSP/src/goto_def.rs

@@ -19,13 +19,12 @@ use kclvm_sema::resolver::scope::{
     builtin_scope, ProgramScope, Scope, ScopeObject, ScopeObjectKind,
 };
 use kclvm_sema::ty::{DictType, SchemaType};
-use lsp_types::{GotoDefinitionResponse, Url};
-use lsp_types::{Location, Range};
+use lsp_types::GotoDefinitionResponse;
 use std::cell::RefCell;
 use std::path::Path;
 use std::rc::Rc;
 
-use crate::to_lsp::lsp_pos;
+use crate::to_lsp::lsp_location;
 use crate::util::{
     fix_missing_identifier, get_pkg_scope, get_pos_from_real_path, get_real_path_from_external,
     inner_most_expr_in_stmt,
@@ -408,24 +407,16 @@ fn positions_to_goto_def_resp(
         0 => None,
         1 => {
             let (start, end) = positions.iter().next().unwrap().clone();
-            Some(lsp_types::GotoDefinitionResponse::Scalar(Location {
-                uri: Url::from_file_path(start.filename.clone()).unwrap(),
-                range: Range {
-                    start: lsp_pos(&start),
-                    end: lsp_pos(&end),
-                },
-            }))
+            Some(lsp_types::GotoDefinitionResponse::Scalar(lsp_location(
+                start.filename.clone(),
+                &start,
+                &end,
+            )))
         }
         _ => {
             let mut res = vec![];
             for (start, end) in positions {
-                res.push(Location {
-                    uri: Url::from_file_path(start.filename.clone()).unwrap(),
-                    range: Range {
-                        start: lsp_pos(start),
-                        end: lsp_pos(end),
-                    },
-                })
+                res.push(lsp_location(start.filename.clone(), &start, &end))
             }
             Some(lsp_types::GotoDefinitionResponse::Array(res))
         }

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

@@ -1,8 +1,8 @@
-use anyhow::Ok;
+use anyhow::{anyhow, Ok};
 use crossbeam_channel::Sender;
-use kclvm_ast::ast::Stmt;
 use lsp_types::{Location, TextEdit};
 use ra_ap_vfs::VfsPath;
+use std::collections::HashMap;
 use std::time::Instant;
 
 use crate::{
@@ -13,9 +13,10 @@ use crate::{
     find_refs::find_refs,
     formatting::format,
     from_lsp::{self, file_path_from_url, kcl_pos},
-    goto_def::{find_def, goto_definition},
+    goto_def::goto_definition,
     hover, quick_fix,
     state::{log_message, LanguageServerSnapshot, LanguageServerState, Task},
+    util,
 };
 
 impl LanguageServerState {
@@ -52,6 +53,7 @@ impl LanguageServerState {
             .on::<lsp_types::request::CodeActionRequest>(handle_code_action)?
             .on::<lsp_types::request::Formatting>(handle_formatting)?
             .on::<lsp_types::request::RangeFormatting>(handle_range_formatting)?
+            .on::<lsp_types::request::Rename>(handle_rename)?
             .finish();
 
         Ok(())
@@ -141,49 +143,25 @@ pub(crate) fn handle_reference(
     params: lsp_types::ReferenceParams,
     sender: Sender<Task>,
 ) -> anyhow::Result<Option<Vec<Location>>> {
-    // 1. find definition of current token
     let file = file_path_from_url(&params.text_document_position.text_document.uri)?;
     let path = from_lsp::abs_path(&params.text_document_position.text_document.uri)?;
     let db = snapshot.get_db(&path.clone().into())?;
     let pos = kcl_pos(&file, params.text_document_position.position);
-    let word_index_map = snapshot.word_index_map.clone();
-
     let log = |msg: String| log_message(msg, &sender);
-
-    if let Some(def_resp) = goto_definition(&db.prog, &pos, &db.scope) {
-        match def_resp {
-            lsp_types::GotoDefinitionResponse::Scalar(def_loc) => {
-                // get the def location
-                if let Some(def_name) = match db.prog.pos_to_stmt(&pos) {
-                    Some(node) => match node.node {
-                        Stmt::Import(_) => None,
-                        _ => match find_def(node.clone(), &pos, &db.scope) {
-                            Some(def) => Some(def.get_name()),
-                            None => None,
-                        },
-                    },
-                    None => None,
-                } {
-                    return find_refs(
-                        Some(snapshot.vfs),
-                        word_index_map,
-                        def_loc,
-                        def_name,
-                        file,
-                        log,
-                    );
-                }
-            }
-            _ => return Ok(None),
+    match find_refs(
+        &db.prog,
+        &pos,
+        &db.scope,
+        snapshot.word_index_map.clone(),
+        Some(snapshot.vfs.clone()),
+        log,
+    ) {
+        core::result::Result::Ok(locations) => Ok(Some(locations)),
+        Err(msg) => {
+            log(format!("Find references failed: {msg}"))?;
+            Ok(None)
         }
-    } else {
-        log_message(
-            "Definition item not found, result in no reference".to_string(),
-            &sender,
-        )?;
     }
-
-    return Ok(None);
 }
 
 /// Called when a `Completion` request was received.
@@ -239,3 +217,66 @@ pub(crate) fn handle_document_symbol(
     }
     Ok(res)
 }
+
+/// Called when a `Rename` request was received.
+pub(crate) fn handle_rename(
+    snapshot: LanguageServerSnapshot,
+    params: lsp_types::RenameParams,
+    sender: Sender<Task>,
+) -> anyhow::Result<Option<lsp_types::WorkspaceEdit>> {
+    // 1. check the new name validity
+    let new_name = params.new_name;
+    if !util::is_valid_kcl_name(new_name.clone()) {
+        return Err(anyhow!("Can not rename to: {new_name}, invalid name"));
+    }
+
+    // 2. find all the references of the symbol
+    let file = file_path_from_url(&params.text_document_position.text_document.uri)?;
+    let path = from_lsp::abs_path(&params.text_document_position.text_document.uri)?;
+    let db = snapshot.get_db(&path.into())?;
+    let kcl_pos = kcl_pos(&file, params.text_document_position.position);
+    let log = |msg: String| log_message(msg, &sender);
+    let references = find_refs(
+        &db.prog,
+        &kcl_pos,
+        &db.scope,
+        snapshot.word_index_map.clone(),
+        Some(snapshot.vfs.clone()),
+        log,
+    );
+    match references {
+        Result::Ok(locations) => {
+            match locations.len() {
+                0 => {
+                    let _ = log("Symbol not found".to_string());
+                    anyhow::Ok(None)
+                }
+                _ => {
+                    // 3. return the workspaceEdit to rename all the references with the new name
+                    let mut workspace_edit = lsp_types::WorkspaceEdit::default();
+
+                    let changes =
+                        locations
+                            .into_iter()
+                            .fold(HashMap::new(), |mut map, location| {
+                                let uri = location.uri;
+                                map.entry(uri.clone())
+                                    .or_insert_with(Vec::new)
+                                    .push(TextEdit {
+                                        range: location.range,
+                                        new_text: new_name.clone(),
+                                    });
+                                map
+                            });
+                    workspace_edit.changes = Some(changes);
+                    anyhow::Ok(Some(workspace_edit))
+                }
+            }
+        }
+        Err(msg) => {
+            let err_msg = format!("Can not rename symbol: {msg}");
+            log(err_msg.clone())?;
+            return Err(anyhow!(err_msg));
+        }
+    }
+}

kclvm/tools/src/LSP/src/test_data/rename_test/main.k

@@ -0,0 +1,6 @@
+import .pkg.vars
+
+Bob = vars.Person {
+    name: "Bob"
+    age: 30
+}