Fix setmetatable

Author: CppCXY

crates/emmylua_code_analysis/src/compilation/analyzer/lua/metatable.rs

@@ -14,14 +14,20 @@ pub fn analyze_setmetatable(analyzer: &mut LuaAnalyzer, call_expr: LuaCallExpr)
         return Some(());
     }
 
+    let table = args[0].clone();
     let metatable = args[1].clone();
-    let LuaExpr::TableExpr(table_expr) = metatable else {
+    let LuaExpr::TableExpr(metatable) = metatable else {
         return Some(());
     };
 
     let file_id = analyzer.file_id;
-    let operator_owner = LuaOperatorOwner::Table(InFiled::new(file_id, table_expr.get_range()));
-    for field in table_expr.get_fields() {
+    analyzer.db.get_metatable_index_mut().add(
+        InFiled::new(file_id, table.get_range()),
+        InFiled::new(file_id, metatable.get_range()),
+    );
+
+    let operator_owner = LuaOperatorOwner::Table(InFiled::new(file_id, metatable.get_range()));
+    for field in metatable.get_fields() {
         analyze_metable_field(analyzer, &field, &operator_owner);
     }
 

crates/emmylua_code_analysis/src/db_index/metatable/mod.rs

@@ -0,0 +1,38 @@
+use std::collections::HashMap;
+
+use rowan::TextRange;
+
+use crate::{FileId, InFiled};
+
+use super::LuaIndex;
+
+#[derive(Debug)]
+pub struct LuaMetatableIndex {
+    pub metatables: HashMap<InFiled<TextRange>, InFiled<TextRange>>,
+}
+
+impl LuaMetatableIndex {
+    pub fn new() -> Self {
+        Self {
+            metatables: HashMap::new(),
+        }
+    }
+
+    pub fn add(&mut self, table: InFiled<TextRange>, metatable: InFiled<TextRange>) {
+        self.metatables.insert(table, metatable);
+    }
+
+    pub fn get(&self, table: &InFiled<TextRange>) -> Option<&InFiled<TextRange>> {
+        self.metatables.get(table)
+    }
+}
+
+impl LuaIndex for LuaMetatableIndex {
+    fn remove(&mut self, file_id: FileId) {
+        self.metatables.retain(|key, _| key.file_id != file_id);
+    }
+
+    fn clear(&mut self) {
+        self.metatables.clear();
+    }
+}

crates/emmylua_code_analysis/src/db_index/mod.rs

@@ -3,6 +3,7 @@ mod dependency;
 mod diagnostic;
 mod flow;
 mod member;
+mod metatable;
 mod module;
 mod operators;
 mod property;
@@ -23,6 +24,7 @@ pub use member::{
     LuaMember, LuaMemberFeature, LuaMemberId, LuaMemberIndex, LuaMemberIndexItem, LuaMemberKey,
     LuaMemberOwner,
 };
+use metatable::LuaMetatableIndex;
 use module::LuaModuleIndex;
 pub use module::{ModuleInfo, WorkspaceId};
 pub use operators::{
@@ -50,6 +52,7 @@ pub struct DbIndex {
     flow_index: LuaFlowIndex,
     vfs: Vfs,
     file_dependencies_index: LuaDenpendencyIndex,
+    metatable_index: LuaMetatableIndex,
     emmyrc: Arc<Emmyrc>,
 }
 
@@ -69,6 +72,7 @@ impl DbIndex {
             flow_index: LuaFlowIndex::new(),
             vfs: Vfs::new(),
             file_dependencies_index: LuaDenpendencyIndex::new(),
+            metatable_index: LuaMetatableIndex::new(),
             emmyrc: Arc::new(Emmyrc::default()),
         }
     }
@@ -79,6 +83,14 @@ impl DbIndex {
         }
     }
 
+    pub fn get_metatable_index_mut(&mut self) -> &mut LuaMetatableIndex {
+        &mut self.metatable_index
+    }
+
+    pub fn get_metatable_index(&self) -> &LuaMetatableIndex {
+        &self.metatable_index
+    }
+
     pub fn get_decl_index_mut(&mut self) -> &mut LuaDeclIndex {
         &mut self.decl_index
     }
@@ -199,6 +211,7 @@ impl LuaIndex for DbIndex {
         self.operator_index.remove(file_id);
         self.flow_index.remove(file_id);
         self.file_dependencies_index.remove(file_id);
+        self.metatable_index.remove(file_id);
     }
 
     fn clear(&mut self) {
@@ -213,5 +226,6 @@ impl LuaIndex for DbIndex {
         self.operator_index.clear();
         self.flow_index.clear();
         self.file_dependencies_index.clear();
+        self.metatable_index.clear();
     }
 }

crates/emmylua_code_analysis/src/semantic/infer/infer_call/infer_setmetatable.rs

@@ -2,7 +2,7 @@ use emmylua_parser::{LuaAstNode, LuaCallExpr, LuaExpr, LuaIndexKey};
 
 use crate::{
     infer_expr, semantic::infer::InferResult, DbIndex, InFiled, InferFailReason, LuaInferCache,
-    LuaInstanceType, LuaType,
+    LuaType,
 };
 
 pub fn infer_setmetatable_call(
@@ -20,19 +20,17 @@ pub fn infer_setmetatable_call(
     let basic_table = args[0].clone();
     let metatable = args[1].clone();
 
-    let meta_type = infer_metatable_type(db, cache, metatable)?;
+    let (meta_type, is_index) = infer_metatable_index_type(db, cache, metatable)?;
     match &basic_table {
         LuaExpr::TableExpr(table_expr) => {
-            if table_expr.is_empty() {
+            if table_expr.is_empty() && is_index {
                 return Ok(meta_type);
-            } else {
-                let file_id = cache.get_file_id();
-                let inst = LuaType::Instance(
-                    LuaInstanceType::new(meta_type, InFiled::new(file_id, basic_table.get_range()))
-                        .into(),
-                );
-                return Ok(inst);
             }
+
+            return Ok(LuaType::TableConst(InFiled::new(
+                cache.get_file_id(),
+                table_expr.get_range(),
+            )));
         }
         _ => {
             if meta_type.is_unknown() {
@@ -44,11 +42,11 @@ pub fn infer_setmetatable_call(
     }
 }
 
-fn infer_metatable_type(
+fn infer_metatable_index_type(
     db: &DbIndex,
     cache: &mut LuaInferCache,
     metatable: LuaExpr,
-) -> InferResult {
+) -> Result<(LuaType, bool /*__index type*/), InferFailReason> {
     match &metatable {
         LuaExpr::TableExpr(table) => {
             let fields = table.get_fields();
@@ -71,13 +69,15 @@ fn infer_metatable_type(
                             | LuaExpr::IndexExpr(_)
                             | LuaExpr::NameExpr(_)
                     ) {
-                        return infer_expr(db, cache, field_value);
+                        let meta_type = infer_expr(db, cache, field_value)?;
+                        return Ok((meta_type, true));
                     }
                 }
             }
         }
         _ => {}
     };
 
-    return infer_expr(db, cache, metatable);
+    let meta_type = infer_expr(db, cache, metatable)?;
+    Ok((meta_type, false))
 }

crates/emmylua_code_analysis/src/semantic/infer/infer_call/mod.rs

@@ -138,6 +138,19 @@ fn collect_function_type(
             infer_guard,
             funcs,
         )?,
+        LuaType::TableConst(inst) => {
+            collect_func_by_table(db, cache, &inst, call_expr.clone(), infer_guard, funcs)?;
+        }
+        LuaType::Instance(inst) => {
+            collect_function_type(
+                db,
+                cache,
+                call_expr,
+                inst.get_base().clone(),
+                infer_guard,
+                funcs,
+            )?;
+        }
         LuaType::Union(union_types) => {
             for sub_type in union_types.get_types() {
                 collect_function_type(
@@ -455,3 +468,48 @@ fn is_last_call_expr(call_expr: &LuaCallExpr) -> bool {
 
     false
 }
+
+fn collect_func_by_table(
+    db: &DbIndex,
+    _: &mut LuaInferCache,
+    table: &InFiled<TextRange>,
+    _: LuaCallExpr,
+    _: &mut InferGuard,
+    funcs: &mut Vec<Arc<LuaFunctionType>>,
+) -> Result<(), InferFailReason> {
+    let metatable = db
+        .get_metatable_index()
+        .get(table)
+        .ok_or(InferFailReason::None)?;
+
+    let operator_index = db.get_operator_index();
+    let operator_ids = operator_index
+        .get_operators(&metatable.clone().into(), LuaOperatorMetaMethod::Call)
+        .ok_or(InferFailReason::None)?;
+
+    for overload_id in operator_ids {
+        let operator = operator_index
+            .get_operator(overload_id)
+            .ok_or(InferFailReason::None)?;
+        let func = operator.get_operator_func();
+        match func {
+            LuaType::DocFunction(f) => {
+                funcs.push(f.clone());
+            }
+            LuaType::Signature(signature_id) => {
+                let signature = db
+                    .get_signature_index()
+                    .get(&signature_id)
+                    .ok_or(InferFailReason::None)?;
+                if !signature.is_resolve_return() {
+                    return Err(InferFailReason::UnResolveSignatureReturn(signature_id));
+                }
+
+                funcs.push(signature.to_call_operator_func_type());
+            }
+            _ => {}
+        }
+    }
+
+    Ok(())
+}

crates/emmylua_code_analysis/src/semantic/infer/infer_call_func.rs

@@ -73,7 +73,7 @@ pub fn infer_call_expr_func(
         LuaType::Instance(inst) => {
             infer_instance_type_doc_function(db, cache, &inst, call_expr, infer_guard, args_count)
         }
-        LuaType::TableConst(meta_table) => infer_metatable_type_doc_function(db, meta_table),
+        LuaType::TableConst(meta_table) => infer_table_type_doc_function(db, meta_table),
         LuaType::Union(union) => {
             // 此时我们将其视为泛型实例化联合体
             if union.get_types().len() > 1
@@ -325,7 +325,7 @@ fn infer_instance_type_doc_function(
     args_count: Option<usize>,
 ) -> InferCallFuncResult {
     let base = instance.get_base();
-    let metatable = match &base {
+    let base_table = match &base {
         LuaType::TableConst(meta_table) => meta_table.clone(),
         LuaType::Instance(inst) => {
             return infer_instance_type_doc_function(
@@ -340,14 +340,16 @@ fn infer_instance_type_doc_function(
         _ => return Err(InferFailReason::None),
     };
 
-    infer_metatable_type_doc_function(db, metatable)
+    infer_table_type_doc_function(db, base_table)
 }
 
-fn infer_metatable_type_doc_function(
-    db: &DbIndex,
-    meta_table: InFiled<TextRange>,
-) -> InferCallFuncResult {
-    let meta_table_owner = LuaOperatorOwner::Table(meta_table);
+fn infer_table_type_doc_function(db: &DbIndex, table: InFiled<TextRange>) -> InferCallFuncResult {
+    let meta_table = db
+        .get_metatable_index()
+        .get(&table)
+        .ok_or(InferFailReason::None)?;
+
+    let meta_table_owner = LuaOperatorOwner::Table(meta_table.clone());
 
     let call_operators = db
         .get_operator_index()

crates/emmylua_code_analysis/src/semantic/infer/infer_index.rs

@@ -406,7 +406,12 @@ fn infer_member_by_index_table(
     table_range: &InFiled<TextRange>,
     index_expr: LuaIndexMemberExpr,
 ) -> InferResult {
-    let meta_owner = LuaOperatorOwner::Table(table_range.clone());
+    let metatable = db
+        .get_metatable_index()
+        .get(table_range)
+        .ok_or(InferFailReason::FieldDotFound)?;
+
+    let meta_owner = LuaOperatorOwner::Table(metatable.clone());
     let operator_ids = db
         .get_operator_index()
         .get_operators(&meta_owner, LuaOperatorMetaMethod::Index)

crates/emmylua_ls/src/handlers/signature_helper/build_signature_helper.rs

@@ -329,7 +329,9 @@ fn build_table_call_signature_help(
     colon_call: bool,
     current_idx: usize,
 ) -> Option<SignatureHelp> {
-    let operator_owner = LuaOperatorOwner::Table(meta);
+    let metatable = semantic_model.get_db().get_metatable_index().get(&meta)?;
+
+    let operator_owner = LuaOperatorOwner::Table(metatable.clone());
     let db = semantic_model.get_db();
     let operator_ids = db
         .get_operator_index()