is_reference_to 优化, 允许从 table 查找到 class

xuhuanzy · xuhuanzy · commit 6adfcf7ca78f · 2025-05-30T04:34:55.000+08:00
diff --git a/crates/emmylua_code_analysis/src/semantic/reference/mod.rs b/crates/emmylua_code_analysis/src/semantic/reference/mod.rs
@@ -1,8 +1,15 @@
-use emmylua_parser::LuaSyntaxNode;
+use std::collections::HashSet;
 
-use crate::{DbIndex, LuaMemberId, LuaSemanticDeclId};
+use emmylua_parser::{
+    LuaAssignStat, LuaAstNode, LuaSyntaxKind, LuaSyntaxNode, LuaTableExpr, LuaTableField,
+};
 
-use super::{semantic_info::infer_node_semantic_decl, LuaInferCache, SemanticDeclLevel};
+use crate::{DbIndex, LuaMemberId, LuaMemberKey, LuaSemanticDeclId, LuaType};
+
+use super::{
+    infer_table_should_be, member::find_members, semantic_info::infer_node_semantic_decl,
+    LuaInferCache, SemanticDeclLevel,
+};
 
 pub fn is_reference_to(
     db: &DbIndex,
@@ -16,9 +23,13 @@ pub fn is_reference_to(
         return Some(true);
     }
 
-    match (node_semantic_decl_id, semantic_decl) {
+    match (node_semantic_decl_id, &semantic_decl) {
         (LuaSemanticDeclId::Member(node_member_id), LuaSemanticDeclId::Member(member_id)) => {
-            is_member_reference_to(db, node_member_id, member_id)
+            if let Some(true) = is_member_reference_to(db, node_member_id, *member_id) {
+                return Some(true);
+            }
+
+            is_member_origin_reference_to(db, infer_config, node_member_id, semantic_decl)
         }
         _ => Some(false),
     }
@@ -34,3 +45,126 @@ fn is_member_reference_to(
 
     Some(node_owner == owner)
 }
+
+fn is_member_origin_reference_to(
+    db: &DbIndex,
+    infer_config: &mut LuaInferCache,
+    node_member_id: LuaMemberId,
+    semantic_decl: LuaSemanticDeclId,
+) -> Option<bool> {
+    let node_origin = find_member_origin_owner(db, infer_config, node_member_id)
+        .unwrap_or(LuaSemanticDeclId::Member(node_member_id));
+
+    match (node_origin, semantic_decl) {
+        (LuaSemanticDeclId::Member(node_owner), LuaSemanticDeclId::Member(member_owner)) => {
+            is_member_reference_to(db, node_owner, member_owner)
+        }
+        (node_origin, member_origin) => Some(node_origin == member_origin),
+    }
+}
+
+pub fn find_member_origin_owner(
+    db: &DbIndex,
+    infer_config: &mut LuaInferCache,
+    member_id: LuaMemberId,
+) -> Option<LuaSemanticDeclId> {
+    const MAX_ITERATIONS: usize = 50;
+    let mut visited_members = HashSet::new();
+
+    let mut current_owner = resolve_member_owner(db, infer_config, &member_id);
+    let mut final_owner = current_owner.clone();
+    let mut iteration_count = 0;
+
+    while let Some(LuaSemanticDeclId::Member(current_member_id)) = &current_owner {
+        if visited_members.contains(current_member_id) || iteration_count >= MAX_ITERATIONS {
+            break;
+        }
+
+        visited_members.insert(current_member_id.clone());
+        iteration_count += 1;
+
+        match resolve_member_owner(db, infer_config, current_member_id) {
+            Some(next_owner) => {
+                final_owner = Some(next_owner.clone());
+                current_owner = Some(next_owner);
+            }
+            None => break,
+        }
+    }
+
+    final_owner
+}
+
+fn resolve_member_owner(
+    db: &DbIndex,
+    infer_config: &mut LuaInferCache,
+    member_id: &LuaMemberId,
+) -> Option<LuaSemanticDeclId> {
+    let root = db
+        .get_vfs()
+        .get_syntax_tree(&member_id.file_id)?
+        .get_red_root();
+    let current_node = member_id.get_syntax_id().to_node_from_root(&root)?;
+    match member_id.get_syntax_id().get_kind() {
+        LuaSyntaxKind::TableFieldAssign => {
+            if LuaTableField::can_cast(current_node.kind().into()) {
+                let table_field = LuaTableField::cast(current_node.clone())?;
+                // 如果表是类, 那么通过类型推断获取 owner
+                if let Some(owner_id) =
+                    resolve_table_field_through_type_inference(db, infer_config, &table_field)
+                {
+                    return Some(owner_id);
+                }
+                // 非类, 那么通过右值推断
+                let value_expr = table_field.get_value_expr()?;
+                let value_node = value_expr.get_syntax_id().to_node_from_root(&root)?;
+                infer_node_semantic_decl(db, infer_config, value_node, SemanticDeclLevel::default())
+            } else {
+                None
+            }
+        }
+        LuaSyntaxKind::IndexExpr => {
+            let assign_node = current_node.parent()?;
+            let assign_stat = LuaAssignStat::cast(assign_node)?;
+            let (vars, exprs) = assign_stat.get_var_and_expr_list();
+
+            for (var, expr) in vars.iter().zip(exprs.iter()) {
+                if var.syntax().text_range() == current_node.text_range() {
+                    let expr_node = expr.get_syntax_id().to_node_from_root(&root)?;
+                    return infer_node_semantic_decl(
+                        db,
+                        infer_config,
+                        expr_node,
+                        SemanticDeclLevel::default(),
+                    );
+                }
+            }
+            None
+        }
+        _ => None,
+    }
+}
+
+fn resolve_table_field_through_type_inference(
+    db: &DbIndex,
+    infer_config: &mut LuaInferCache,
+    table_field: &LuaTableField,
+) -> Option<LuaSemanticDeclId> {
+    let parent = table_field.syntax().parent()?;
+    let table_expr = LuaTableExpr::cast(parent)?;
+    let table_type = infer_table_should_be(db, infer_config, table_expr).ok()?;
+
+    if !matches!(table_type, LuaType::Ref(_) | LuaType::Def(_)) {
+        return None;
+    }
+
+    let field_key = table_field.get_field_key()?;
+    let key = LuaMemberKey::from_index_key(db, infer_config, &field_key).ok()?;
+    let member_infos = find_members(db, &table_type)?;
+
+    member_infos
+        .iter()
+        .find(|m| m.key == key)?
+        .property_owner_id
+        .clone()
+}
diff --git a/crates/emmylua_ls/src/handlers/definition/goto_def_definition.rs b/crates/emmylua_ls/src/handlers/definition/goto_def_definition.rs
@@ -179,7 +179,7 @@ pub fn goto_str_tpl_ref_definition(
     None
 }
 
-fn find_table_member_definition(
+pub fn find_table_member_definition(
     semantic_model: &SemanticModel,
     trigger_token: &LuaSyntaxToken,
     member_key: &LuaMemberKey,
diff --git a/crates/emmylua_ls/src/handlers/implementation/implementation_searcher.rs b/crates/emmylua_ls/src/handlers/implementation/implementation_searcher.rs
@@ -5,10 +5,12 @@ use emmylua_code_analysis::{
     SemanticDeclLevel, SemanticModel,
 };
 use emmylua_parser::{
-    LuaAstNode, LuaDocTagField, LuaIndexExpr, LuaStat, LuaSyntaxNode, LuaSyntaxToken,
+    LuaAstNode, LuaDocTagField, LuaIndexExpr, LuaStat, LuaSyntaxNode, LuaSyntaxToken, LuaTableField,
 };
 use lsp_types::Location;
 
+use crate::handlers::hover::find_member_origin_owner;
+
 pub fn search_implementations(
     semantic_model: &SemanticModel,
     compilation: &LuaCompilation,
@@ -45,15 +47,17 @@ pub fn search_member_implementations(
         .get_db()
         .get_member_index()
         .get_member(&member_id)?;
-    let key = member.get_key();
+    let member_key = member.get_key();
+
     let index_references = semantic_model
         .get_db()
         .get_reference_index()
-        .get_index_references(&key)?;
+        .get_index_references(&member_key)?;
 
     let mut semantic_cache = HashMap::new();
 
-    let property_owner = LuaSemanticDeclId::Member(member_id);
+    let property_owner = find_member_origin_owner(semantic_model, member_id)
+        .unwrap_or(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) {
@@ -94,6 +98,21 @@ fn check_member_reference(semantic_model: &SemanticModel, node: LuaSyntaxNode) -
             let prefix_type = semantic_model
                 .infer_expr(expr.get_prefix_expr()?.into())
                 .ok()?;
+            // TODO: 需要实现更复杂的逻辑, 即当为`Ref`时, 针对指定的实例定义到其实现
+            /*
+               ---@class A
+               ---@field a number -- 这里寻找实现只匹配到`A.a`, 不能穿透到`a.a`与`b.a`
+               local A = {}
+               A.a = 1
+
+               ---@type A
+               local a = {}
+               a.a = 1 -- 这里寻找实现不能匹配到`b.a`
+
+               ---@type A
+               local b = a
+               b.a = 2 -- 这里寻找实现不能匹配到`a.a`
+            */
             match prefix_type {
                 LuaType::Ref(_) => {
                     return None;
@@ -128,6 +147,14 @@ fn check_member_reference(semantic_model: &SemanticModel, node: LuaSyntaxNode) -
         tag_field_node if LuaDocTagField::can_cast(tag_field_node.kind().into()) => {
             return Some(());
         }
+        table_field_node if LuaTableField::can_cast(table_field_node.kind().into()) => {
+            let table_field = LuaTableField::cast(table_field_node.clone())?;
+            if table_field.is_assign_field() {
+                return Some(());
+            } else {
+                return None;
+            }
+        }
         _ => {}
     }
 
diff --git a/crates/emmylua_ls/src/handlers/test/implementation_test.rs b/crates/emmylua_ls/src/handlers/test/implementation_test.rs
@@ -20,15 +20,15 @@ mod tests {
                delete()
             "#,
         );
-
-        ws.check_implementation(
+        assert!(ws.check_implementation(
             r#"
                 local M = {}
                 function M.de<??>lete(a)
                 end
                 return M
             "#,
-        );
+            1,
+        ));
     }
 
     #[test]
@@ -57,11 +57,12 @@ mod tests {
                 local a = test.a
             "#,
         );
-        ws.check_implementation(
+        assert!(ws.check_implementation(
             r#"
                 t<??>est
             "#,
-        );
+            3,
+        ));
     }
 
     #[test]
@@ -82,10 +83,49 @@ mod tests {
 
             "#,
         );
-        ws.check_implementation(
+        assert!(ws.check_implementation(
             r#"
                 yyy.<??>a = 2
             "#,
-        );
+            3,
+        ));
+    }
+
+    #[test]
+    fn test_table_field_definition_1() {
+        let mut ws = ProviderVirtualWorkspace::new();
+        assert!(ws.check_implementation(
+            r#"
+                ---@class T
+                ---@field func fun(self: T) 注释注释
+
+                ---@type T
+                local t = {
+                    func = function(self)
+                    end,
+                }
+
+                t:fun<??>c()
+            "#,
+            2,
+        ));
+    }
+
+    #[test]
+    fn test_table_field_definition_2() {
+        let mut ws = ProviderVirtualWorkspace::new();
+        assert!(ws.check_implementation(
+            r#"
+                ---@class T
+                ---@field func fun(self: T) 注释注释
+
+                ---@type T
+                local t = {
+                    f<??>unc = function(self)
+                    end,
+                }
+            "#,
+            2,
+        ));
     }
 }
diff --git a/crates/emmylua_ls/src/handlers/test_lib/mod.rs b/crates/emmylua_ls/src/handlers/test_lib/mod.rs
@@ -236,7 +236,7 @@ impl ProviderVirtualWorkspace {
         true
     }
 
-    pub fn check_implementation(&mut self, block_str: &str) -> bool {
+    pub fn check_implementation(&mut self, block_str: &str, len: usize) -> bool {
         let content = Self::handle_file_content(block_str);
         let Some((content, position)) = content else {
             return false;
@@ -251,7 +251,10 @@ impl ProviderVirtualWorkspace {
             return false;
         };
         dbg!(&implementations.len());
-        true
+        if implementations.len() == len {
+            return true;
+        }
+        false
     }
 
     pub fn check_definition(&mut self, block_str: &str) -> bool {

Original file line number	Diff line number	Diff line change
`@@ -179,7 +179,7 @@ pub fn goto_str_tpl_ref_definition(`
`179`	`179`	`None`
`180`	`180`	`}`
`181`	`181`
`182`		`-fn find_table_member_definition(`
	`182`	`+pub fn find_table_member_definition(`
`183`	`183`	`semantic_model: &SemanticModel,`
`184`	`184`	`trigger_token: &LuaSyntaxToken,`
`185`	`185`	`member_key: &LuaMemberKey,`