支持泛型约束来自泛型时补全

Author: xuhuanzy

crates/emmylua_ls/src/handlers/completion/providers/function_provider.rs

@@ -2,7 +2,8 @@ use emmylua_code_analysis::{
     DbIndex, GenericTplId, InferGuard, InferGuardRef, LuaAliasCallKind, LuaAliasCallType,
     LuaDeclLocation, LuaFunctionType, LuaMember, LuaMemberKey, LuaMemberOwner, LuaMultiLineUnion,
     LuaSemanticDeclId, LuaStringTplType, LuaType, LuaTypeCache, LuaTypeDeclId, LuaUnionType,
-    RenderLevel, SemanticDeclLevel, get_real_type,
+    RenderLevel, SemanticDeclLevel, TypeSubstitutor, build_call_constraint_context, get_real_type,
+    instantiate_type_generic, normalize_constraint_type,
 };
 use emmylua_parser::{
     LuaAssignStat, LuaAst, LuaAstNode, LuaAstToken, LuaCallArgList, LuaCallExpr, LuaClosureExpr,
@@ -11,6 +12,7 @@ use emmylua_parser::{
 };
 use itertools::Itertools;
 use lsp_types::{CompletionItem, Documentation};
+use std::sync::Arc;
 
 use crate::handlers::{
     completion::{
@@ -292,29 +294,95 @@ fn infer_call_arg_list(
             param_idx += 1;
         }
     }
+    let constraint_substitutor = build_call_constraint_context(&builder.semantic_model, &call_expr)
+        .map(|(ctx, _)| ctx.substitutor);
+    let substitutor = constraint_substitutor.as_ref();
     let typ = call_expr_func
         .get_params()
         .get(param_idx)?
         .1
         .clone()
         .unwrap_or(LuaType::Unknown);
+    let typ = resolve_param_type(builder, typ, substitutor);
     let mut types = Vec::new();
     types.push(typ);
     push_function_overloads_param(
         builder,
         &call_expr,
         call_expr_func.get_params(),
         param_idx,
+        substitutor,
         &mut types,
     );
     Some(types.into_iter().unique().collect()) // 需要去重
 }
 
+fn resolve_param_type(
+    builder: &CompletionBuilder,
+    mut typ: LuaType,
+    substitutor: Option<&TypeSubstitutor>,
+) -> LuaType {
+    let db = builder.semantic_model.get_db();
+    if let Some(substitutor) = substitutor {
+        typ = apply_substitutor_to_type(db, typ, substitutor);
+    }
+    normalize_constraint_type(db, typ)
+}
+
+fn apply_substitutor_to_type(db: &DbIndex, typ: LuaType, substitutor: &TypeSubstitutor) -> LuaType {
+    if let LuaType::Call(alias_call) = &typ {
+        if alias_call.get_call_kind() == LuaAliasCallKind::KeyOf {
+            let operands = alias_call
+                .get_operands()
+                .iter()
+                .map(|operand| instantiate_type_generic(db, operand, substitutor))
+                .collect::<Vec<_>>();
+            return LuaType::Call(Arc::new(LuaAliasCallType::new(
+                alias_call.get_call_kind(),
+                operands,
+            )));
+        }
+    }
+    if let Some(alias_call) = rebuild_keyof_alias_call(db, &typ, substitutor) {
+        return alias_call;
+    }
+    instantiate_type_generic(db, &typ, substitutor)
+}
+
+fn rebuild_keyof_alias_call(
+    db: &DbIndex,
+    original_type: &LuaType,
+    substitutor: &TypeSubstitutor,
+) -> Option<LuaType> {
+    let tpl = match original_type {
+        LuaType::TplRef(tpl) | LuaType::ConstTplRef(tpl) => tpl,
+        _ => return None,
+    };
+    let constraint = tpl.get_constraint()?;
+    let LuaType::Call(alias_call) = constraint else {
+        return None;
+    };
+    if alias_call.get_call_kind() != LuaAliasCallKind::KeyOf {
+        return None;
+    }
+
+    let operands = alias_call
+        .get_operands()
+        .iter()
+        .map(|operand| instantiate_type_generic(db, operand, substitutor))
+        .collect::<Vec<_>>();
+    Some(LuaType::Call(Arc::new(LuaAliasCallType::new(
+        alias_call.get_call_kind(),
+        operands,
+    ))))
+}
+
 fn push_function_overloads_param(
     builder: &mut CompletionBuilder,
     call_expr: &LuaCallExpr,
     call_params: &[(String, Option<LuaType>)],
     param_idx: usize,
+    substitutor: Option<&TypeSubstitutor>,
     types: &mut Vec<LuaType>,
 ) -> Option<()> {
     let member_index = builder.semantic_model.get_db().get_member_index();
@@ -394,6 +462,7 @@ fn push_function_overloads_param(
 
         // 添加匹配的参数类型
         if let Some(param_type) = overload_params.get(param_idx).and_then(|p| p.1.clone()) {
+            let param_type = resolve_param_type(builder, param_type, substitutor);
             types.push(param_type);
         }
     }

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

@@ -2266,4 +2266,77 @@ mod tests {
         ));
         Ok(())
     }
+
+    #[gtest]
+    fn test_generic_constraint() -> Result<()> {
+        let mut ws = ProviderVirtualWorkspace::new();
+        ws.def(
+            r#"
+            ---@alias std.RawGet<T, K> unknown
+
+            ---@alias std.ConstTpl<T> unknown
+
+            ---@generic T, K extends keyof T
+            ---@param object T
+            ---@param key K
+            ---@return std.RawGet<T, K>
+            function pick(object, key)
+            end
+
+            ---@class Person
+            ---@field age integer
+        "#,
+        );
+
+        check!(ws.check_completion_with_kind(
+            r#"
+            ---@type Person
+            local person
+
+            pick(person, <??>)
+            "#,
+            vec![VirtualCompletionItem {
+                label: "\"age\"".to_string(),
+                kind: CompletionItemKind::VARIABLE,
+                ..Default::default()
+            },],
+            CompletionTriggerKind::TRIGGER_CHARACTER
+        ),);
+        Ok(())
+    }
+
+    #[gtest]
+    fn test_generic_constraint_inline_object_completion() -> Result<()> {
+        let mut ws = ProviderVirtualWorkspace::new();
+        ws.def(
+            r#"
+            ---@generic T, K extends keyof T
+            ---@param object T
+            ---@param key K
+            function pick(object, key)
+            end
+            "#,
+        );
+
+        check!(ws.check_completion_with_kind(
+            r#"
+            pick({ foo = 1, bar = 2 }, <??>)
+            "#,
+            vec![
+                VirtualCompletionItem {
+                    label: "\"bar\"".to_string(),
+                    kind: CompletionItemKind::CONSTANT,
+                    ..Default::default()
+                },
+                VirtualCompletionItem {
+                    label: "\"foo\"".to_string(),
+                    kind: CompletionItemKind::CONSTANT,
+                    ..Default::default()
+                },
+            ],
+            CompletionTriggerKind::TRIGGER_CHARACTER
+        ));
+
+        Ok(())
+    }
 }