update references: 本地函数现在能够搜索导出的引用

Author: xuhuanzy

crates/emmylua_ls/src/handlers/call_hierarchy/build_call_hierarchy.rs

@@ -100,7 +100,7 @@ pub fn build_incoming_hierarchy(
     let mut locations = vec![];
     match semantic_decl {
         LuaSemanticDeclId::LuaDecl(decl_id) => {
-            search_decl_references(semantic_model, decl_id, &mut locations);
+            search_decl_references(semantic_model, compilation, decl_id, &mut locations);
         }
         LuaSemanticDeclId::Member(member_id) => {
             search_member_references(semantic_model, compilation, member_id, &mut locations);

crates/emmylua_ls/src/handlers/code_lens/resolve_code_lens.rs

@@ -44,7 +44,7 @@ pub fn resolve_code_lens(
             let file_id = decl_id.file_id;
             let mut semantic_model = compilation.get_semantic_model(file_id)?;
             let mut results = Vec::new();
-            search_decl_references(&mut semantic_model, decl_id, &mut results);
+            search_decl_references(&mut semantic_model, compilation, decl_id, &mut results);
             let ref_count = results.len();
             let uri = semantic_model.get_document().get_uri();
             let command = make_usage_command(uri, code_lens.range, ref_count, client_id, results);

crates/emmylua_ls/src/handlers/references/mod.rs

@@ -1,8 +1,11 @@
 mod reference_seacher;
 
 use crate::context::ServerContextSnapshot;
+use emmylua_code_analysis::{EmmyLuaAnalysis, FileId};
 use emmylua_parser::{LuaAstNode, LuaTokenKind};
-use lsp_types::{ClientCapabilities, Location, OneOf, ReferenceParams, ServerCapabilities};
+use lsp_types::{
+    ClientCapabilities, Location, OneOf, Position, ReferenceParams, ServerCapabilities,
+};
 use reference_seacher::search_references;
 pub use reference_seacher::{search_decl_references, search_member_references};
 use rowan::TokenAtOffset;
@@ -19,6 +22,15 @@ pub async fn on_references_handler(
     let analysis = context.analysis.read().await;
     let file_id = analysis.get_file_id(&uri)?;
     let position = params.text_document_position.position;
+
+    references(&analysis, file_id, position)
+}
+
+pub fn references(
+    analysis: &EmmyLuaAnalysis,
+    file_id: FileId,
+    position: Position,
+) -> Option<Vec<Location>> {
     let mut semantic_model = analysis.compilation.get_semantic_model(file_id)?;
     if !semantic_model.get_emmyrc().references.enable {
         return None;

crates/emmylua_ls/src/handlers/references/reference_seacher.rs

@@ -1,11 +1,12 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use emmylua_code_analysis::{
-    LuaCompilation, LuaDeclId, LuaMemberId, LuaMemberKey, LuaSemanticDeclId, LuaTypeDeclId,
-    SemanticDeclLevel, SemanticModel,
+    DeclReference, LuaCompilation, LuaDeclId, LuaMemberId, LuaMemberKey, LuaSemanticDeclId,
+    LuaTypeDeclId, SemanticDeclLevel, SemanticModel,
 };
 use emmylua_parser::{
-    LuaAst, LuaAstNode, LuaAstToken, LuaNameToken, LuaStringToken, LuaSyntaxNode, LuaSyntaxToken,
+    LuaAssignStat, LuaAst, LuaAstNode, LuaAstToken, LuaNameToken, LuaStringToken, LuaSyntaxNode,
+    LuaSyntaxToken,
 };
 use lsp_types::Location;
 
@@ -20,7 +21,7 @@ pub fn search_references(
     {
         match semantic_decl {
             LuaSemanticDeclId::LuaDecl(decl_id) => {
-                search_decl_references(semantic_model, decl_id, &mut result);
+                search_decl_references(semantic_model, compilation, decl_id, &mut result);
             }
             LuaSemanticDeclId::Member(member_id) => {
                 search_member_references(semantic_model, compilation, member_id, &mut result);
@@ -36,11 +37,16 @@ pub fn search_references(
         fuzzy_search_references(compilation, token, &mut result);
     }
 
+    // 简单过滤, 同行的多个引用只保留一个
+    // let filtered_result = filter_duplicate_and_covered_locations(result);
+    // Some(filtered_result)
+
     Some(result)
 }
 
 pub fn search_decl_references(
     semantic_model: &SemanticModel,
+    compilation: &LuaCompilation,
     decl_id: LuaDeclId,
     result: &mut Vec<Location>,
 ) -> Option<()> {
@@ -58,9 +64,15 @@ pub fn search_decl_references(
         if let Some(location) = document.to_lsp_location(decl.get_range()) {
             result.push(location);
         }
+        let typ = semantic_model.get_type(decl.get_id().into());
+        let is_signature = typ.is_signature();
+
         for decl_ref in decl_refs {
             let location = document.to_lsp_location(decl_ref.range.clone())?;
             result.push(location);
+            if is_signature {
+                get_signature_decl_member_references(semantic_model, compilation, result, decl_ref);
+            }
         }
 
         return Some(());
@@ -119,7 +131,7 @@ pub fn search_member_references(
             let range = in_filed_syntax_id.value.get_range();
             let location = document.to_lsp_location(range)?;
             result.push(location);
-            search_member_secondary_references(semantic_model, node, result);
+            search_member_secondary_references(semantic_model, compilation, node, result);
         }
     }
 
@@ -128,6 +140,7 @@ pub fn search_member_references(
 
 fn search_member_secondary_references(
     semantic_model: &SemanticModel,
+    compilation: &LuaCompilation,
     node: LuaSyntaxNode,
     result: &mut Vec<Location>,
 ) -> Option<()> {
@@ -141,7 +154,7 @@ fn search_member_secondary_references(
                 .position(|value| value.get_position() == position)?;
             let var = vars.get(idx)?;
             let decl_id = LuaDeclId::new(semantic_model.get_file_id(), var.get_position());
-            search_decl_references(semantic_model, decl_id, result);
+            search_decl_references(semantic_model, compilation, decl_id, result);
             let document = semantic_model.get_document();
             let range = document.to_lsp_location(var.get_range())?;
             result.push(range);
@@ -152,7 +165,7 @@ fn search_member_secondary_references(
             let idx = values.position(|value| value.get_position() == position)?;
             let name = local_names.get(idx)?;
             let decl_id = LuaDeclId::new(semantic_model.get_file_id(), name.get_position());
-            search_decl_references(semantic_model, decl_id, result);
+            search_decl_references(semantic_model, compilation, decl_id, result);
             let document = semantic_model.get_document();
             let range = document.to_lsp_location(name.get_range())?;
             result.push(range);
@@ -241,3 +254,80 @@ fn search_type_decl_references(
 
     Some(())
 }
+
+fn get_signature_decl_member_references(
+    semantic_model: &SemanticModel,
+    compilation: &LuaCompilation,
+    result: &mut Vec<Location>,
+    decl_ref: &DeclReference,
+) -> Option<Vec<Location>> {
+    let root = semantic_model.get_root();
+    let position = decl_ref.range.start();
+    let token = root.syntax().token_at_offset(position).right_biased()?;
+    let parent = token.parent()?;
+
+    match parent.parent()? {
+        assign_stat_node if LuaAssignStat::can_cast(assign_stat_node.kind().into()) => {
+            let assign_stat = LuaAssignStat::cast(assign_stat_node)?;
+            let (vars, values) = assign_stat.get_var_and_expr_list();
+            let idx = values
+                .iter()
+                .position(|value| value.get_position() == position)?;
+            let var = vars.get(idx)?;
+            let decl_id = semantic_model
+                .find_decl(var.syntax().clone().into(), SemanticDeclLevel::default())?;
+            if let LuaSemanticDeclId::Member(member_id) = decl_id {
+                search_member_references(semantic_model, compilation, member_id, result);
+            }
+        }
+
+        _ => {}
+    }
+    None
+}
+
+#[allow(unused)]
+fn filter_duplicate_and_covered_locations(locations: Vec<Location>) -> Vec<Location> {
+    if locations.is_empty() {
+        return locations;
+    }
+    let mut sorted_locations = locations;
+    sorted_locations.sort_by(|a, b| {
+        a.uri
+            .to_string()
+            .cmp(&b.uri.to_string())
+            .then_with(|| a.range.start.line.cmp(&b.range.start.line))
+            .then_with(|| b.range.end.line.cmp(&a.range.end.line))
+    });
+
+    let mut result = Vec::new();
+    let mut seen_lines_by_uri: HashMap<String, HashSet<u32>> = HashMap::new();
+
+    for location in sorted_locations {
+        let uri_str = location.uri.to_string();
+        let seen_lines = seen_lines_by_uri.entry(uri_str).or_default();
+
+        let start_line = location.range.start.line;
+        let end_line = location.range.end.line;
+
+        let is_covered = (start_line..=end_line).any(|line| seen_lines.contains(&line));
+
+        if !is_covered {
+            for line in start_line..=end_line {
+                seen_lines.insert(line);
+            }
+            result.push(location);
+        }
+    }
+
+    // 最终按位置排序
+    result.sort_by(|a, b| {
+        a.uri
+            .to_string()
+            .cmp(&b.uri.to_string())
+            .then_with(|| a.range.start.line.cmp(&b.range.start.line))
+            .then_with(|| a.range.start.character.cmp(&b.range.start.character))
+    });
+
+    result
+}

crates/emmylua_ls/src/handlers/test/mod.rs

@@ -6,6 +6,7 @@ mod hover_function_test;
 mod hover_test;
 mod implementation_test;
 mod inlay_hint_test;
+mod references_test;
 mod rename_test;
 mod semantic_token_test;
 mod signature_helper_test;

crates/emmylua_ls/src/handlers/test/references_test.rs

@@ -0,0 +1,29 @@
+#[cfg(test)]
+mod tests {
+
+    use crate::handlers::test_lib::ProviderVirtualWorkspace;
+
+    #[test]
+    fn test_function_references() {
+        let mut ws = ProviderVirtualWorkspace::new();
+        ws.def_file(
+            "1.lua",
+            r#"
+                local flush = require("virtual_0").flush
+                flush()
+            "#,
+        );
+        let result = ws.check_references(
+            r#"
+                local export = {}
+                local function fl<??>ush()
+                end
+                export.flush = flush
+                return export
+            "#,
+        );
+        assert!(result.is_some());
+        let locations = result.unwrap();
+        assert!(locations.len() >= 4);
+    }
+}

crates/emmylua_ls/src/handlers/test_lib/mod.rs

@@ -20,7 +20,7 @@ use crate::{
     },
 };
 
-use super::{hover::hover, implementation::implementation};
+use super::{hover::hover, implementation::implementation, references::references};
 
 /// A virtual workspace for testing.
 #[allow(unused)]
@@ -375,4 +375,16 @@ impl ProviderVirtualWorkspace {
 
         true
     }
+
+    pub fn check_references(&mut self, block_str: &str) -> Option<Vec<lsp_types::Location>> {
+        let content = Self::handle_file_content(block_str);
+        let Some((content, position)) = content else {
+            return None;
+        };
+        let file_id = self.def(&content);
+        let result = references(&self.analysis, file_id, position);
+        // dbg!(&result);
+        dbg!(&result.as_ref().unwrap().len());
+        result
+    }
 }