feat: 重构条件类型推断

Author: xuhuanzy

crates/emmylua_code_analysis/src/compilation/analyzer/doc/file_generic_index.rs

@@ -1,8 +1,10 @@
 use hashbrown::HashMap;
 
 use rowan::{TextRange, TextSize};
+use smol_str::SmolStr;
+use std::sync::Arc;
 
-use crate::{GenericParam, GenericTplId, LuaType};
+use crate::{GenericParam, GenericTpl, GenericTplId, LuaType};
 
 pub trait GenericIndex: std::fmt::Debug {
     fn add_generic_scope(&mut self, ranges: Vec<TextRange>, is_func: bool) -> GenericScopeId;
@@ -15,10 +17,6 @@ pub trait GenericIndex: std::fmt::Debug {
         }
     }
 
-    fn append_pending_type_param(&mut self, _param: GenericParam) {}
-
-    fn clear_pending_type_params(&mut self) {}
-
     fn find_generic(
         &self,
         position: TextSize,
@@ -29,42 +27,11 @@ pub trait GenericIndex: std::fmt::Debug {
 #[derive(Debug, Clone)]
 pub struct FileGenericIndex {
     scopes: Vec<FileGenericScope>,
-    pending_type_params: Vec<GenericParam>,
 }
 
 impl FileGenericIndex {
     pub fn new() -> Self {
-        Self {
-            scopes: Vec::new(),
-            pending_type_params: Vec::new(),
-        }
-    }
-
-    pub fn add_generic_scope(&mut self, ranges: Vec<TextRange>, is_func: bool) -> GenericScopeId {
-        let scope_id = GenericScopeId::new(self.scopes.len());
-        let next_tpl_id = self.next_tpl_id(&ranges, is_func);
-        self.scopes.push(FileGenericScope::new(ranges, next_tpl_id));
-        scope_id
-    }
-
-    pub fn append_generic_param(&mut self, scope_id: GenericScopeId, param: GenericParam) {
-        if let Some(scope) = self.scopes.get_mut(scope_id.id) {
-            scope.insert_param(param);
-        }
-    }
-
-    pub fn append_generic_params(&mut self, scope_id: GenericScopeId, params: Vec<GenericParam>) {
-        for param in params {
-            self.append_generic_param(scope_id, param);
-        }
-    }
-
-    pub fn append_pending_type_param(&mut self, param: GenericParam) {
-        self.pending_type_params.push(param);
-    }
-
-    pub fn clear_pending_type_params(&mut self) {
-        self.pending_type_params.clear();
+        Self { scopes: Vec::new() }
     }
 
     fn next_tpl_id(&self, ranges: &[TextRange], is_func: bool) -> GenericTplId {
@@ -86,10 +53,31 @@ impl FileGenericIndex {
                 .sum(),
         )
     }
+}
+
+impl GenericIndex for FileGenericIndex {
+    fn add_generic_scope(&mut self, ranges: Vec<TextRange>, is_func: bool) -> GenericScopeId {
+        let scope_id = GenericScopeId::new(self.scopes.len());
+        let next_tpl_id = self.next_tpl_id(&ranges, is_func);
+        self.scopes.push(FileGenericScope::new(ranges, next_tpl_id));
+        scope_id
+    }
+
+    fn append_generic_param(&mut self, scope_id: GenericScopeId, param: GenericParam) {
+        if let Some(scope) = self.scopes.get_mut(scope_id.id) {
+            scope.insert_param(param);
+        }
+    }
+
+    fn append_generic_params(&mut self, scope_id: GenericScopeId, params: Vec<GenericParam>) {
+        for param in params {
+            self.append_generic_param(scope_id, param);
+        }
+    }
 
     /// Find generic parameter by position and name.
     /// return (GenericTplId, constraint, default)
-    pub fn find_generic(
+    fn find_generic(
         &self,
         position: TextSize,
         name: &str,
@@ -108,49 +96,7 @@ impl FileGenericIndex {
             }
         }
 
-        // 搜索前置类型参数, 例如 ---@alias Pick<T, K extends keyof T>
-        self.pending_type_params
-            .iter()
-            .enumerate()
-            .rev()
-            .find(|(_, param)| param.name == name)
-            .map(|(idx, param)| {
-                (
-                    GenericTplId::Type(idx as u32),
-                    param.type_constraint.clone(),
-                    param.default_type.clone(),
-                )
-            })
-    }
-}
-
-impl GenericIndex for FileGenericIndex {
-    fn add_generic_scope(&mut self, ranges: Vec<TextRange>, is_func: bool) -> GenericScopeId {
-        FileGenericIndex::add_generic_scope(self, ranges, is_func)
-    }
-
-    fn append_generic_param(&mut self, scope_id: GenericScopeId, param: GenericParam) {
-        FileGenericIndex::append_generic_param(self, scope_id, param);
-    }
-
-    fn append_generic_params(&mut self, scope_id: GenericScopeId, params: Vec<GenericParam>) {
-        FileGenericIndex::append_generic_params(self, scope_id, params);
-    }
-
-    fn append_pending_type_param(&mut self, param: GenericParam) {
-        FileGenericIndex::append_pending_type_param(self, param);
-    }
-
-    fn clear_pending_type_params(&mut self) {
-        FileGenericIndex::clear_pending_type_params(self);
-    }
-
-    fn find_generic(
-        &self,
-        position: TextSize,
-        name: &str,
-    ) -> Option<(GenericTplId, Option<LuaType>, Option<LuaType>)> {
-        FileGenericIndex::find_generic(self, position, name)
+        None
     }
 }
 
@@ -195,3 +141,86 @@ impl FileGenericScope {
         self.ranges.iter().any(|range| range.contains(position))
     }
 }
+
+#[derive(Debug, Clone, Default)]
+pub struct ConditionalInferIndex {
+    scopes: Vec<ConditionalInferScope>,
+    next_infer_id: u32,
+}
+
+impl ConditionalInferIndex {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn enter_scope(&mut self) {
+        self.scopes.push(ConditionalInferScope::new());
+    }
+
+    pub fn leave_scope(&mut self) -> Option<ConditionalInferScope> {
+        self.scopes.pop()
+    }
+
+    pub fn set_current_refs_visible(&mut self, visible: bool) {
+        if let Some(scope) = self.scopes.last_mut() {
+            scope.refs_visible = visible;
+        }
+    }
+
+    pub fn declare(&mut self, name: &str) -> Option<Arc<GenericTpl>> {
+        let scope_idx = self.scopes.len().checked_sub(1)?;
+        if let Some(tpl) = self.scopes[scope_idx].bindings.get(name) {
+            return Some(tpl.clone());
+        }
+
+        let tpl_id = GenericTplId::ConditionalInfer(self.next_infer_id);
+        self.next_infer_id += 1;
+        let tpl = Arc::new(GenericTpl::new(
+            tpl_id,
+            SmolStr::new(name).into(),
+            None,
+            None,
+        ));
+
+        let scope = &mut self.scopes[scope_idx];
+        scope.bindings.insert(name.to_string(), tpl.clone());
+        scope
+            .params
+            .push(GenericParam::new(SmolStr::new(name), None, None, None));
+        Some(tpl)
+    }
+
+    pub fn find_ref(&self, name: &str) -> Option<Arc<GenericTpl>> {
+        self.scopes
+            .iter()
+            .rev()
+            .filter(|scope| scope.refs_visible)
+            .find_map(|scope| scope.bindings.get(name).cloned())
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ConditionalInferScope {
+    /// 是否允许在当前阶段把普通名字解析为 conditional infer 引用.
+    /// condition 阶段只声明 `infer P`, true 分支阶段才允许引用 `P`.
+    refs_visible: bool,
+    /// 当前 conditional 作用域内的 `infer` 名字到实际模板的绑定.
+    /// 同名 `infer P` 会复用同一个 `GenericTplId::ConditionalInfer`.
+    bindings: HashMap<String, Arc<GenericTpl>>,
+    /// 当前 conditional 声明过的 infer 参数元数据, 保留给 `LuaConditionalType`.
+    params: Vec<GenericParam>,
+}
+
+impl ConditionalInferScope {
+    fn new() -> Self {
+        Self {
+            refs_visible: false,
+            bindings: HashMap::new(),
+            params: Vec::new(),
+        }
+    }
+
+    pub fn into_params(self) -> Vec<GenericParam> {
+        self.params
+    }
+}

crates/emmylua_code_analysis/src/compilation/analyzer/doc/infer_type.rs

@@ -3,12 +3,11 @@ use std::sync::Arc;
 use emmylua_parser::{
     LuaAst, LuaAstNode, LuaComment, LuaDocAttributeType, LuaDocBinaryType, LuaDocConditionalType,
     LuaDocDescriptionOwner, LuaDocFuncType, LuaDocGenericDecl, LuaDocGenericDeclList,
-    LuaDocGenericType, LuaDocIndexAccessType, LuaDocInferType, LuaDocMappedType,
-    LuaDocMultiLineUnionType, LuaDocObjectFieldKey, LuaDocObjectType, LuaDocStrTplType, LuaDocType,
-    LuaDocUnaryType, LuaDocVariadicType, LuaLiteralToken, LuaSyntaxKind, LuaTypeBinaryOperator,
+    LuaDocGenericType, LuaDocIndexAccessType, LuaDocMappedType, LuaDocMultiLineUnionType,
+    LuaDocObjectFieldKey, LuaDocObjectType, LuaDocStrTplType, LuaDocType, LuaDocUnaryType,
+    LuaDocVariadicType, LuaLiteralToken, LuaSyntaxKind, LuaTypeBinaryOperator,
     LuaTypeUnaryOperator, LuaVarExpr, NumberResult,
 };
-use internment::ArcIntern;
 use rowan::TextRange;
 use smol_str::SmolStr;
 
@@ -23,7 +22,10 @@ use crate::{
     },
 };
 
-use super::{file_generic_index::GenericIndex, preprocess_description};
+use super::{
+    file_generic_index::{ConditionalInferIndex, GenericIndex},
+    preprocess_description,
+};
 
 #[derive(Debug)]
 pub struct DocTypeAnalyzeContext<'a> {
@@ -33,6 +35,7 @@ pub struct DocTypeAnalyzeContext<'a> {
     pub workspace_id: WorkspaceId,
     comment: Option<LuaComment>,
     options: DocTypeAnalyzeOptions,
+    conditional_infer_index: ConditionalInferIndex,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -77,6 +80,7 @@ impl<'a> DocTypeAnalyzeContext<'a> {
             workspace_id,
             comment: None,
             options: DocTypeAnalyzeOptions::default(),
+            conditional_infer_index: ConditionalInferIndex::new(),
         }
     }
 
@@ -203,7 +207,9 @@ pub fn infer_type(analyzer: &mut DocTypeAnalyzeContext<'_>, node: LuaDocType) ->
         }
         LuaDocType::Infer(infer_type) => {
             if let Some(name) = infer_type.get_generic_decl_name_text() {
-                return LuaType::ConditionalInfer(ArcIntern::new(SmolStr::new(&name)));
+                if let Some(tpl) = analyzer.conditional_infer_index.declare(&name) {
+                    return LuaType::TplRef(tpl);
+                }
             }
         }
         LuaDocType::Mapped(mapped_type) => {
@@ -246,6 +252,10 @@ fn infer_buildin_or_ref_type(
             LuaType::Table
         }
         _ => {
+            if let Some(tpl) = analyzer.conditional_infer_index.find_ref(name) {
+                return LuaType::TplRef(tpl);
+            }
+
             if let Some((tpl_id, constraint, default_type)) =
                 analyzer.generic_index.find_generic(position, name)
             {
@@ -899,26 +909,37 @@ fn infer_conditional_type(
     cond_type: &LuaDocConditionalType,
 ) -> LuaType {
     if let Some((condition, when_true, when_false)) = cond_type.get_types() {
-        // 收集条件中的所有 infer 声明
-        let infer_params = collect_cond_infer_params(&condition);
-        if !infer_params.is_empty() {
-            // 条件表达式中 infer 声明的类型参数只允许在`true`分支中使用
-            let true_range = when_true.get_range();
-            let scope_id = analyzer
-                .generic_index
-                .add_generic_scope(vec![true_range], false);
-            analyzer
-                .generic_index
-                .append_generic_params(scope_id, infer_params.clone());
-        }
-
-        // 处理条件和分支类型
+        analyzer.conditional_infer_index.enter_scope();
+
         let condition_type = infer_type(analyzer, condition);
+        let LuaType::Call(alias_call) = condition_type else {
+            analyzer.conditional_infer_index.leave_scope();
+            return LuaType::Unknown;
+        };
+        if alias_call.get_call_kind() != LuaAliasCallKind::Extends
+            || alias_call.get_operands().len() != 2
+        {
+            analyzer.conditional_infer_index.leave_scope();
+            return LuaType::Unknown;
+        }
+        let operands = alias_call.get_operands();
+        let checked_type = operands[0].clone();
+        let extends_type = operands[1].clone();
+
+        analyzer
+            .conditional_infer_index
+            .set_current_refs_visible(true);
         let true_type = infer_type(analyzer, when_true);
+        let infer_params = analyzer
+            .conditional_infer_index
+            .leave_scope() // 退出当前作用域
+            .map(|scope| scope.into_params())
+            .unwrap_or_default();
         let false_type = infer_type(analyzer, when_false);
 
         return LuaConditionalType::new(
-            condition_type,
+            checked_type,
+            extends_type,
             true_type,
             false_type,
             infer_params,
@@ -930,18 +951,6 @@ fn infer_conditional_type(
     LuaType::Unknown
 }
 
-/// 收集条件类型中的条件表达式中所有 infer 声明
-fn collect_cond_infer_params(doc_type: &LuaDocType) -> Vec<GenericParam> {
-    let mut params = Vec::new();
-    let doc_infer_types = doc_type.descendants::<LuaDocInferType>();
-    for infer_type in doc_infer_types {
-        if let Some(name) = infer_type.get_generic_decl_name_text() {
-            params.push(GenericParam::new(SmolStr::new(&name), None, None, None));
-        }
-    }
-    params
-}
-
 fn infer_mapped_type(
     analyzer: &mut DocTypeAnalyzeContext<'_>,
     mapped_type: &LuaDocMappedType,