feature: impl implementation

optimize searcher member implementation

Author: xuhuanzy

crates/emmylua_ls/src/handlers/implementation/implementation_searcher.rs

@@ -0,0 +1,216 @@
+use std::collections::HashMap;
+
+use emmylua_code_analysis::{
+    LuaCompilation, LuaDeclId, LuaMemberId, LuaSemanticDeclId, LuaType, LuaTypeDeclId,
+    SemanticDeclLevel, SemanticModel,
+};
+use emmylua_parser::{
+    LuaAstNode, LuaDocTagField, LuaIndexExpr, LuaStat, LuaSyntaxNode, LuaSyntaxToken,
+};
+use lsp_types::Location;
+
+pub fn search_implementations(
+    semantic_model: &SemanticModel,
+    compilation: &LuaCompilation,
+    token: LuaSyntaxToken,
+) -> Option<Vec<Location>> {
+    let mut result = Vec::new();
+    if let Some(semantic_decl) =
+        semantic_model.find_decl(token.clone().into(), SemanticDeclLevel::NoTrace)
+    {
+        match semantic_decl {
+            LuaSemanticDeclId::TypeDecl(type_decl_id) => {
+                search_type_implementations(semantic_model, compilation, type_decl_id, &mut result);
+            }
+            LuaSemanticDeclId::Member(member_id) => {
+                search_member_implementations(semantic_model, compilation, member_id, &mut result);
+            }
+            LuaSemanticDeclId::LuaDecl(decl_id) => {
+                search_decl_implementations(semantic_model, compilation, decl_id, &mut result);
+            }
+            _ => {}
+        }
+    }
+
+    Some(result)
+}
+
+pub fn search_member_implementations(
+    semantic_model: &SemanticModel,
+    compilation: &LuaCompilation,
+    member_id: LuaMemberId,
+    result: &mut Vec<Location>,
+) -> Option<()> {
+    let member = semantic_model
+        .get_db()
+        .get_member_index()
+        .get_member(&member_id)?;
+    let key = member.get_key();
+    let index_references = semantic_model
+        .get_db()
+        .get_reference_index()
+        .get_index_references(&key)?;
+
+    let mut semantic_cache = HashMap::new();
+
+    let property_owner = LuaSemanticDeclId::Member(member_id);
+    for in_filed_syntax_id in index_references {
+        let semantic_model =
+            if let Some(semantic_model) = semantic_cache.get_mut(&in_filed_syntax_id.file_id) {
+                semantic_model
+            } else {
+                let semantic_model = compilation.get_semantic_model(in_filed_syntax_id.file_id)?;
+                semantic_cache.insert(in_filed_syntax_id.file_id, semantic_model);
+                semantic_cache.get_mut(&in_filed_syntax_id.file_id)?
+            };
+        let root = semantic_model.get_root();
+        let node = in_filed_syntax_id.value.to_node_from_root(root.syntax())?;
+
+        if check_member_reference(&semantic_model, node.clone()).is_none() {
+            continue;
+        }
+
+        if !semantic_model.is_reference_to(
+            node,
+            property_owner.clone(),
+            SemanticDeclLevel::default(),
+        ) {
+            continue;
+        }
+
+        let document = semantic_model.get_document();
+        let range = in_filed_syntax_id.value.get_range();
+        let location = document.to_lsp_location(range)?;
+        result.push(location);
+    }
+    Some(())
+}
+
+/// 检查成员引用是否符合实现
+fn check_member_reference(semantic_model: &SemanticModel, node: LuaSyntaxNode) -> Option<()> {
+    match &node {
+        expr_node if LuaIndexExpr::can_cast(expr_node.kind().into()) => {
+            let expr = LuaIndexExpr::cast(expr_node.clone())?;
+            let prefix_type = semantic_model
+                .infer_expr(expr.get_prefix_expr()?.into())
+                .ok()?;
+            match prefix_type {
+                LuaType::Ref(_) => {
+                    return None;
+                }
+                _ => {}
+            };
+            // 往上寻找 stat 节点
+            let stat = expr.ancestors::<LuaStat>().next()?;
+            match stat {
+                LuaStat::FuncStat(_) => {
+                    return Some(());
+                }
+                LuaStat::AssignStat(assign_stat) => {
+                    // 判断是否在左侧
+                    let (vars, _) = assign_stat.get_var_and_expr_list();
+                    for var in vars {
+                        if var
+                            .syntax()
+                            .text_range()
+                            .contains(node.text_range().start())
+                        {
+                            return Some(());
+                        }
+                    }
+                    return None;
+                }
+                _ => {
+                    return None;
+                }
+            }
+        }
+        tag_field_node if LuaDocTagField::can_cast(tag_field_node.kind().into()) => {
+            return Some(());
+        }
+        _ => {}
+    }
+
+    Some(())
+}
+pub fn search_type_implementations(
+    semantic_model: &SemanticModel,
+    compilation: &LuaCompilation,
+    type_decl_id: LuaTypeDeclId,
+    result: &mut Vec<Location>,
+) -> Option<()> {
+    let db = semantic_model.get_db();
+    let type_index = db.get_type_index();
+    let type_decl = type_index.get_type_decl(&type_decl_id)?;
+    let locations = type_decl.get_locations();
+    let mut semantic_cache = HashMap::new();
+    for location in locations {
+        let semantic_model = if let Some(semantic_model) = semantic_cache.get_mut(&location.file_id)
+        {
+            semantic_model
+        } else {
+            let semantic_model = compilation.get_semantic_model(location.file_id)?;
+            semantic_cache.insert(location.file_id, semantic_model);
+            semantic_cache.get_mut(&location.file_id)?
+        };
+        let document = semantic_model.get_document();
+        let range = location.range;
+        let location = document.to_lsp_location(range)?;
+        result.push(location);
+    }
+
+    Some(())
+}
+
+pub fn search_decl_implementations(
+    semantic_model: &SemanticModel,
+    compilation: &LuaCompilation,
+    decl_id: LuaDeclId,
+    result: &mut Vec<Location>,
+) -> Option<()> {
+    let decl = semantic_model
+        .get_db()
+        .get_decl_index()
+        .get_decl(&decl_id)?;
+
+    if decl.is_local() {
+        let document = semantic_model.get_document();
+        let range = decl.get_range();
+        let location = document.to_lsp_location(range)?;
+        result.push(location);
+        return Some(());
+    } else {
+        let name = decl.get_name();
+        let global_decl_ids = semantic_model
+            .get_db()
+            .get_global_index()
+            .get_global_decl_ids(name)?;
+
+        let mut semantic_cache = HashMap::new();
+
+        for global_decl_id in global_decl_ids {
+            let semantic_model =
+                if let Some(semantic_model) = semantic_cache.get_mut(&global_decl_id.file_id) {
+                    semantic_model
+                } else {
+                    let semantic_model = compilation.get_semantic_model(global_decl_id.file_id)?;
+                    semantic_cache.insert(global_decl_id.file_id, semantic_model);
+                    semantic_cache.get_mut(&global_decl_id.file_id)?
+                };
+            let Some(decl) = semantic_model
+                .get_db()
+                .get_decl_index()
+                .get_decl(&global_decl_id)
+            else {
+                continue;
+            };
+
+            let document = semantic_model.get_document();
+            let range = decl.get_range();
+            let location = document.to_lsp_location(range)?;
+            result.push(location);
+        }
+    }
+
+    Some(())
+}

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

@@ -0,0 +1,70 @@
+mod implementation_searcher;
+mod test;
+
+use crate::context::ServerContextSnapshot;
+use emmylua_code_analysis::{EmmyLuaAnalysis, FileId};
+use emmylua_parser::LuaAstNode;
+use implementation_searcher::search_implementations;
+use lsp_types::{
+    request::GotoImplementationParams, ClientCapabilities, GotoDefinitionResponse,
+    ImplementationProviderCapability, Position, ServerCapabilities,
+};
+use rowan::TokenAtOffset;
+use tokio_util::sync::CancellationToken;
+
+use super::RegisterCapabilities;
+
+pub async fn on_implementation_handler(
+    context: ServerContextSnapshot,
+    params: GotoImplementationParams,
+    _: CancellationToken,
+) -> Option<GotoDefinitionResponse> {
+    let uri = params.text_document_position_params.text_document.uri;
+    let analysis = context.analysis.read().await;
+    let file_id = analysis.get_file_id(&uri)?;
+    let position = params.text_document_position_params.position;
+
+    implementation(&analysis, file_id, position)
+}
+
+pub fn implementation(
+    analysis: &EmmyLuaAnalysis,
+    file_id: FileId,
+    position: Position,
+) -> Option<GotoDefinitionResponse> {
+    let mut semantic_model = analysis.compilation.get_semantic_model(file_id)?;
+
+    let root = semantic_model.get_root();
+    let position_offset = {
+        let document = semantic_model.get_document();
+        document.get_offset(position.line as usize, position.character as usize)?
+    };
+
+    if position_offset > root.syntax().text_range().end() {
+        return None;
+    }
+
+    let token = match root.syntax().token_at_offset(position_offset) {
+        TokenAtOffset::None => return None,
+        TokenAtOffset::Single(token) => token,
+        TokenAtOffset::Between(token, _) => token,
+    };
+
+    let implementations =
+        search_implementations(&mut semantic_model, &analysis.compilation, token)?;
+
+    if implementations.is_empty() {
+        return None;
+    }
+
+    Some(GotoDefinitionResponse::Array(implementations))
+}
+
+pub struct ImplementationCapabilities;
+
+impl RegisterCapabilities for ImplementationCapabilities {
+    fn register_capabilities(server_capabilities: &mut ServerCapabilities, _: &ClientCapabilities) {
+        server_capabilities.implementation_provider =
+            Some(ImplementationProviderCapability::Simple(true));
+    }
+}

crates/emmylua_ls/src/handlers/implementation/test/implementation_test.rs

@@ -0,0 +1,91 @@
+#[cfg(test)]
+mod tests {
+
+    use crate::handlers::test_lib::ProviderVirtualWorkspace;
+
+    #[test]
+    fn test_1() {
+        let mut ws = ProviderVirtualWorkspace::new();
+        ws.def_file(
+            "2.lua",
+            r#"
+               delete = require("virtual_0").delete
+               delete()
+            "#,
+        );
+        ws.def_file(
+            "3.lua",
+            r#"
+               delete = require("virtual_0").delete
+               delete()
+            "#,
+        );
+
+        ws.check_implementation(
+            r#"
+                local M = {}
+                function M.de<??>lete(a)
+                end
+                return M
+            "#,
+        );
+    }
+
+    #[test]
+    fn test_2() {
+        let mut ws = ProviderVirtualWorkspace::new();
+        ws.def_file(
+            "1.lua",
+            r#"
+                ---@class (partial) Test
+                test = {}
+
+                test.a = 1
+            "#,
+        );
+        ws.def_file(
+            "2.lua",
+            r#"
+                ---@class (partial) Test
+                test = {}
+                test.a = 1
+            "#,
+        );
+        ws.def_file(
+            "3.lua",
+            r#"
+                local a = test.a
+            "#,
+        );
+        ws.check_implementation(
+            r#"
+                t<??>est
+            "#,
+        );
+    }
+
+    #[test]
+    fn test_3() {
+        let mut ws = ProviderVirtualWorkspace::new();
+        ws.def_file(
+            "1.lua",
+            r#"
+                ---@class YYY
+                ---@field a number
+                yyy = {}
+
+                if false then
+                    yyy.a = 1
+                    if yyy.a then
+                    end
+                end
+
+            "#,
+        );
+        ws.check_implementation(
+            r#"
+                yyy.<??>a = 2
+            "#,
+        );
+    }
+}

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

@@ -0,0 +1 @@
+mod implementation_test;