enhance: build word index prune word in comments

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

@@ -126,7 +126,7 @@ mod tests {
     fn setup_word_index_map(root: &str) -> HashMap<Url, HashMap<String, Vec<Location>>> {
         HashMap::from([(
             Url::from_file_path(root).unwrap(),
-            build_word_index(root.to_string()).unwrap(),
+            build_word_index(root.to_string(), true).unwrap(),
         )])
     }
 

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

@@ -97,8 +97,9 @@ impl LanguageServerState {
         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 old_word_index = build_word_index_for_file_content(old_text, &text_document.uri, true);
+        let new_word_index =
+            build_word_index_for_file_content(text.clone(), &text_document.uri, true);
         let binding = text_document.uri.path();
         let file_path = Path::new(binding); //todo rename
         let word_index_map = &mut *self.word_index_map.write();

kclvm/tools/src/LSP/src/rename.rs

@@ -30,7 +30,7 @@ pub fn rename_symbol(
     match select_symbol(&symbol_spec) {
         Some((name, range)) => {
             // 3. build word index on file_paths, find refs within file_paths scope
-            let word_index = build_word_index_for_file_paths(file_paths)?;
+            let word_index = build_word_index_for_file_paths(file_paths, true)?;
             if let Some(locations) = word_index.get(&name) {
                 // 4. filter out the matched refs
                 // 4.1 collect matched words(names) and remove Duplicates of the file paths

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

@@ -125,7 +125,7 @@ impl LanguageServerState {
         let word_index_map = state.word_index_map.clone();
         state
             .thread_pool
-            .execute(move || build_word_index_map(word_index_map, initialize_params));
+            .execute(move || build_word_index_map(word_index_map, initialize_params, true));
 
         state
     }
@@ -334,17 +334,18 @@ pub(crate) fn log_message(message: String, sender: &Sender<Task>) -> anyhow::Res
 fn build_word_index_map(
     word_index_map: Arc<RwLock<HashMap<Url, HashMap<String, Vec<Location>>>>>,
     initialize_params: InitializeParams,
+    prune_comments: bool,
 ) {
     if let Some(workspace_folders) = initialize_params.workspace_folders {
         for folder in workspace_folders {
             let path = folder.uri.path();
-            if let Ok(word_index) = build_word_index(path.to_string()) {
+            if let Ok(word_index) = build_word_index(path.to_string(), prune_comments) {
                 word_index_map.write().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()) {
+        if let Ok(word_index) = build_word_index(path.to_string(), prune_comments) {
             word_index_map.write().insert(root_uri, word_index);
         }
     }

kclvm/tools/src/LSP/src/util.rs

@@ -784,14 +784,15 @@ pub(crate) fn get_pkg_scope(
 
 pub(crate) fn build_word_index_for_file_paths(
     paths: &[String],
+    prune_comments: bool,
 ) -> anyhow::Result<HashMap<String, Vec<Location>>> {
     let mut index: HashMap<String, Vec<Location>> = HashMap::new();
     for p in paths {
         // str path to url
         if let Ok(url) = Url::from_file_path(p) {
             // read file content and save the word to word index
             let text = read_file(p)?;
-            for (key, values) in build_word_index_for_file_content(text, &url) {
+            for (key, values) in build_word_index_for_file_content(text, &url, prune_comments) {
                 index.entry(key).or_insert_with(Vec::new).extend(values);
             }
         }
@@ -800,21 +801,38 @@ pub(crate) fn build_word_index_for_file_paths(
 }
 
 /// scan and build a word -> Locations index map
-pub(crate) fn build_word_index(path: String) -> anyhow::Result<HashMap<String, Vec<Location>>> {
+pub(crate) fn build_word_index(
+    path: String,
+    prune_comments: bool,
+) -> anyhow::Result<HashMap<String, Vec<Location>>> {
     if let Ok(files) = get_kcl_files(path.clone(), true) {
-        return build_word_index_for_file_paths(&files);
+        return build_word_index_for_file_paths(&files, prune_comments);
     }
     Ok(HashMap::new())
 }
 
 pub(crate) fn build_word_index_for_file_content(
     content: String,
     url: &Url,
+    prune_comments: bool,
 ) -> HashMap<String, Vec<Location>> {
     let mut index: HashMap<String, Vec<Location>> = HashMap::new();
     let lines: Vec<&str> = content.lines().collect();
+    let mut in_docstring = false;
     for (li, line) in lines.into_iter().enumerate() {
-        let words = line_to_words(line.to_string());
+        if prune_comments && !in_docstring {
+            if line.trim_start().starts_with("\"\"\"") {
+                in_docstring = true;
+                continue;
+            }
+        }
+        if prune_comments && in_docstring {
+            if line.trim_end().ends_with("\"\"\"") {
+                in_docstring = false;
+            }
+            continue;
+        }
+        let words = line_to_words(line.to_string(), prune_comments);
         for (key, values) in words {
             index
                 .entry(key)
@@ -878,7 +896,7 @@ fn read_file(path: &String) -> anyhow::Result<String> {
 }
 
 // Split one line into identifier words.
-fn line_to_words(text: String) -> HashMap<String, Vec<Word>> {
+fn line_to_words(text: String, prune_comments: bool) -> HashMap<String, Vec<Word>> {
     let mut result = HashMap::new();
     let mut chars: Vec<char> = text.chars().collect();
     chars.push('\n');
@@ -887,6 +905,9 @@ fn line_to_words(text: String) -> HashMap<String, Vec<Word>> {
     let mut prev_word = false;
     let mut words: Vec<Word> = vec![];
     for (i, ch) in chars.iter().enumerate() {
+        if prune_comments && *ch == '#' {
+            break;
+        }
         let is_id_start = rustc_lexer::is_id_start(*ch);
         let is_id_continue = rustc_lexer::is_id_continue(*ch);
         // If the character is valid identfier start and the previous character is not valid identifier continue, mark the start position.
@@ -921,7 +942,10 @@ fn line_to_words(text: String) -> HashMap<String, Vec<Word>> {
 
 #[cfg(test)]
 mod tests {
-    use super::{build_word_index, line_to_words, word_index_add, word_index_subtract, Word};
+    use super::{
+        build_word_index, build_word_index_for_file_content, line_to_words, word_index_add,
+        word_index_subtract, Word,
+    };
     use lsp_types::{Location, Position, Range, Url};
     use std::collections::HashMap;
     use std::path::PathBuf;
@@ -1132,7 +1156,7 @@ mod tests {
         ]
         .into_iter()
         .collect();
-        match build_word_index(path.to_string()) {
+        match build_word_index(path.to_string(), true) {
             Ok(actual) => {
                 assert_eq!(expect, actual)
             }
@@ -1192,7 +1216,13 @@ mod tests {
 
     #[test]
     fn test_line_to_words() {
-        let lines = ["schema Person:", "name. name again", "some_word word !word"];
+        let lines = [
+            "schema Person:",
+            "name. name again",
+            "some_word word !word",
+            "# this line is a single-line comment",
+            "name # end of line comment",
+        ];
 
         let expects: Vec<HashMap<String, Vec<Word>>> = vec![
             vec![
@@ -1269,10 +1299,106 @@ mod tests {
             ]
             .into_iter()
             .collect(),
+            HashMap::new(),
+            vec![(
+                "name".to_string(),
+                vec![Word {
+                    start_col: 0,
+                    end_col: 4,
+                    word: "name".to_string(),
+                }],
+            )]
+            .into_iter()
+            .collect(),
         ];
         for i in 0..lines.len() {
-            let got = line_to_words(lines[i].to_string());
+            let got = line_to_words(lines[i].to_string(), true);
             assert_eq!(expects[i], got)
         }
     }
+
+    #[test]
+    fn test_build_word_index_for_file_content() {
+        let content = r#"schema Person:
+    """
+    This is a docstring.
+    Person is a schema which defines a person's name and age.
+    """
+    name: str # name must not be empty
+    # age is a positive integer
+    age: int
+"#;
+        let mock_url = Url::parse("file:///path/to/file.k").unwrap();
+        let expects: HashMap<String, Vec<Location>> = vec![
+            (
+                "schema".to_string(),
+                vec![Location {
+                    uri: mock_url.clone(),
+                    range: Range {
+                        start: Position::new(0, 0),
+                        end: Position::new(0, 6),
+                    },
+                }],
+            ),
+            (
+                "Person".to_string(),
+                vec![Location {
+                    uri: mock_url.clone(),
+                    range: Range {
+                        start: Position::new(0, 7),
+                        end: Position::new(0, 13),
+                    },
+                }],
+            ),
+            (
+                "int".to_string(),
+                vec![Location {
+                    uri: mock_url.clone(),
+                    range: Range {
+                        start: Position::new(7, 9),
+                        end: Position::new(7, 12),
+                    },
+                }],
+            ),
+            (
+                "str".to_string(),
+                vec![Location {
+                    uri: mock_url.clone(),
+                    range: Range {
+                        start: Position::new(5, 10),
+                        end: Position::new(5, 13),
+                    },
+                }],
+            ),
+            (
+                "name".to_string(),
+                vec![Location {
+                    uri: mock_url.clone(),
+                    range: Range {
+                        start: Position::new(5, 4),
+                        end: Position::new(5, 8),
+                    },
+                }],
+            ),
+            (
+                "age".to_string(),
+                vec![Location {
+                    uri: mock_url.clone(),
+                    range: Range {
+                        start: Position::new(7, 4),
+                        end: Position::new(7, 7),
+                    },
+                }],
+            ),
+        ]
+        .into_iter()
+        .collect();
+
+        let got = build_word_index_for_file_content(
+            content.to_string(),
+            &mock_url.clone(),
+            true,
+        );
+        assert_eq!(expects, got)
+    }
 }