Merge remote-tracking branch 'origin/main' into merge

xuhuanzy · xuhuanzy · commit de59988e300c · 2025-07-11T01:14:48.000+08:00
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/flow/bind_analyze/comment.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/flow/bind_analyze/comment.rs
@@ -1,6 +1,9 @@
 use emmylua_parser::{LuaAstNode, LuaComment, LuaDocTag};
 
-use crate::{compilation::analyzer::flow::binder::FlowBinder, FlowId, FlowNodeKind};
+use crate::{
+    compilation::analyzer::flow::{bind_analyze::exprs::bind_expr, binder::FlowBinder},
+    FlowId, FlowNodeKind,
+};
 
 pub fn bind_comment(binder: &mut FlowBinder, lua_comment: LuaComment, current: FlowId) -> FlowId {
     let cast_tags = lua_comment.get_doc_tags().filter_map(|it| match it {
@@ -11,7 +14,9 @@ pub fn bind_comment(binder: &mut FlowBinder, lua_comment: LuaComment, current: F
     let mut parent = current;
     for cast in cast_tags {
         let expr = cast.get_key_expr();
-        if expr.is_some() {
+        if let Some(expr) = expr {
+            bind_expr(binder, expr, current);
+
             let flow_id = binder.create_node(FlowNodeKind::TagCast(cast.to_ptr()));
             binder.add_antecedent(flow_id, parent);
             parent = flow_id;
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/flow/bind_analyze/stats.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/flow/bind_analyze/stats.rs
@@ -1,7 +1,7 @@
 use emmylua_parser::{
-    LuaAssignStat, LuaAst, LuaAstNode, LuaBlock, LuaBreakStat, LuaCallArgList, LuaCallExprStat,
-    LuaDoStat, LuaForRangeStat, LuaForStat, LuaFuncStat, LuaGotoStat, LuaIfStat, LuaLabelStat,
-    LuaLocalStat, LuaRepeatStat, LuaReturnStat, LuaWhileStat,
+    BinaryOperator, LuaAssignStat, LuaAst, LuaAstNode, LuaBlock, LuaBreakStat, LuaCallArgList,
+    LuaCallExprStat, LuaDoStat, LuaExpr, LuaForRangeStat, LuaForStat, LuaFuncStat, LuaGotoStat,
+    LuaIfStat, LuaLabelStat, LuaLocalStat, LuaRepeatStat, LuaReturnStat, LuaWhileStat,
 };
 
 use crate::{
@@ -28,8 +28,9 @@ pub fn bind_local_stat(
         let name = &local_names[i];
         let value = &values[i];
         let decl_id = LuaDeclId::new(binder.file_id, name.get_position());
-        let flow_id = bind_expr(binder, value.clone(), current);
-        binder.decl_bind_flow_ref.insert(decl_id, flow_id);
+        if check_local_immutable(binder, decl_id) && check_value_expr_is_check_expr(value.clone()) {
+            binder.decl_bind_expr_ref.insert(decl_id, value.to_ptr());
+        }
     }
 
     for value in values {
@@ -42,6 +43,41 @@ pub fn bind_local_stat(
     local_flow_id
 }
 
+fn check_local_immutable(binder: &mut FlowBinder, decl_id: LuaDeclId) -> bool {
+    let Some(refs) = binder
+        .db
+        .get_reference_index()
+        .get_decl_references(&binder.file_id, &decl_id)
+    else {
+        return true;
+    };
+
+    for r in refs {
+        if r.is_write {
+            return false;
+        }
+    }
+
+    true
+}
+
+fn check_value_expr_is_check_expr(value_expr: LuaExpr) -> bool {
+    match value_expr {
+        LuaExpr::BinaryExpr(binary_expr) => {
+            let Some(op) = binary_expr.get_op_token() else {
+                return false;
+            };
+
+            match op.get_op() {
+                BinaryOperator::OpEq | BinaryOperator::OpNe => true,
+                _ => false,
+            }
+        }
+        LuaExpr::CallExpr(_) => true,
+        _ => false, // Other expressions can be checked
+    }
+}
+
 pub fn bind_assign_stat(
     binder: &mut FlowBinder,
     assign_stat: LuaAssignStat,
diff --git a/crates/emmylua_code_analysis/src/compilation/analyzer/flow/binder.rs b/crates/emmylua_code_analysis/src/compilation/analyzer/flow/binder.rs
@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use emmylua_parser::LuaSyntaxId;
+use emmylua_parser::{LuaAstPtr, LuaExpr, LuaSyntaxId};
 use internment::ArcIntern;
 use rowan::TextSize;
 use smol_str::SmolStr;
@@ -14,7 +14,7 @@ use crate::{
 pub struct FlowBinder<'a> {
     pub db: &'a mut DbIndex,
     pub file_id: FileId,
-    pub decl_bind_flow_ref: HashMap<LuaDeclId, FlowId>,
+    pub decl_bind_expr_ref: HashMap<LuaDeclId, LuaAstPtr<LuaExpr>>,
     pub start: FlowId,
     pub unreachable: FlowId,
     pub loop_label: FlowId,
@@ -35,7 +35,7 @@ impl<'a> FlowBinder<'a> {
             file_id,
             flow_nodes: Vec::new(),
             multiple_antecedents: Vec::new(),
-            decl_bind_flow_ref: HashMap::new(),
+            decl_bind_expr_ref: HashMap::new(),
             labels: HashMap::new(),
             start: FlowId::default(),
             unreachable: FlowId::default(),
@@ -171,7 +171,7 @@ impl<'a> FlowBinder<'a> {
 
     pub fn finish(self) -> FlowTree {
         FlowTree::new(
-            self.decl_bind_flow_ref,
+            self.decl_bind_expr_ref,
             self.flow_nodes,
             self.multiple_antecedents,
             // self.labels,
diff --git a/crates/emmylua_code_analysis/src/compilation/test/flow.rs b/crates/emmylua_code_analysis/src/compilation/test/flow.rs
@@ -891,4 +891,58 @@ end
             "#,
         );
     }
+
+    #[test]
+    fn test_feature_inherit_flow_from_const_local() {
+        let mut ws = VirtualWorkspace::new_with_init_std_lib();
+
+        ws.def(
+            r#"
+            local ret --- @type string | nil
+
+            local h = type(ret) == "string"
+            if h then
+                a = ret
+            end
+
+            local e = type(ret)
+            if e == "string" then
+                b = ret
+            end
+            "#,
+        );
+
+        let a = ws.expr_ty("a");
+        let a_expected = ws.ty("string");
+        assert_eq!(a, a_expected);
+        let b = ws.expr_ty("b");
+        let b_expected = ws.ty("string");
+        assert_eq!(b, b_expected);
+    }
+
+    #[test]
+    fn test_feature_generic_type_guard() {
+        let mut ws = VirtualWorkspace::new();
+
+        ws.def(
+            r#"
+            ---@generic T
+            ---@param type `T`
+            ---@return TypeGuard<T>
+            local function instanceOf(inst, type)
+                return true
+            end
+            
+            local ret --- @type string | nil
+            
+            if instanceOf(ret, "string") then
+                a = ret
+            end
+            "#,
+        );
+
+        let a = ws.expr_ty("a");
+        let a_expected = ws.ty("string");
+        assert_eq!(a, a_expected);
+    }
 }
diff --git a/crates/emmylua_code_analysis/src/db_index/flow/flow_tree.rs b/crates/emmylua_code_analysis/src/db_index/flow/flow_tree.rs
@@ -1,13 +1,12 @@
 use std::collections::HashMap;
 
-use emmylua_parser::LuaSyntaxId;
+use emmylua_parser::{LuaAstPtr, LuaExpr, LuaSyntaxId};
 
 use crate::{FlowId, FlowNode, LuaDeclId};
 
 #[derive(Debug)]
 pub struct FlowTree {
-    #[allow(unused)]
-    decl_bind_flow_ref: HashMap<LuaDeclId, FlowId>,
+    decl_bind_expr_ref: HashMap<LuaDeclId, LuaAstPtr<LuaExpr>>,
     flow_nodes: Vec<FlowNode>,
     multiple_antecedents: Vec<Vec<FlowId>>,
     // labels: HashMap<LuaClosureId, HashMap<SmolStr, FlowId>>,
@@ -16,17 +15,16 @@ pub struct FlowTree {
 
 impl FlowTree {
     pub fn new(
-        decl_bind_flow_ref: HashMap<LuaDeclId, FlowId>,
+        decl_bind_expr_ref: HashMap<LuaDeclId, LuaAstPtr<LuaExpr>>,
         flow_nodes: Vec<FlowNode>,
         multiple_antecedents: Vec<Vec<FlowId>>,
         // labels: HashMap<LuaClosureId, HashMap<SmolStr, FlowId>>,
         bindings: HashMap<LuaSyntaxId, FlowId>,
     ) -> Self {
         Self {
-            decl_bind_flow_ref,
+            decl_bind_expr_ref,
             flow_nodes,
             multiple_antecedents,
-            // labels,
             bindings,
         }
     }
@@ -44,4 +42,8 @@ impl FlowTree {
             .get(id as usize)
             .map(|v| v.as_slice())
     }
+
+    pub fn get_decl_ref_expr(&self, decl_id: &LuaDeclId) -> Option<LuaAstPtr<LuaExpr>> {
+        self.decl_bind_expr_ref.get(decl_id).cloned()
+    }
 }
diff --git a/crates/emmylua_code_analysis/src/db_index/type/types.rs b/crates/emmylua_code_analysis/src/db_index/type/types.rs
@@ -291,6 +291,7 @@ impl LuaType {
             LuaType::Nil | LuaType::Boolean | LuaType::Any | LuaType::Unknown => false,
             LuaType::BooleanConst(boolean) | LuaType::DocBooleanConst(boolean) => boolean.clone(),
             LuaType::Union(u) => u.is_always_truthy(),
+            LuaType::TypeGuard(_) => false,
             _ => true,
         }
     }
@@ -299,6 +300,7 @@ impl LuaType {
         match self {
             LuaType::Nil | LuaType::BooleanConst(false) | LuaType::DocBooleanConst(false) => true,
             LuaType::Union(u) => u.is_always_falsy(),
+            LuaType::TypeGuard(_) => false,
             _ => false,
         }
     }
@@ -400,6 +402,7 @@ impl LuaType {
             LuaType::StrTplRef(_) => true,
             LuaType::SelfInfer => true,
             LuaType::MultiLineUnion(inner) => inner.contain_tpl(),
+            LuaType::TypeGuard(inner) => inner.contain_tpl(),
             _ => false,
         }
     }
diff --git a/crates/emmylua_code_analysis/src/semantic/generic/instantiate_type_generic.rs b/crates/emmylua_code_analysis/src/semantic/generic/instantiate_type_generic.rs
@@ -45,6 +45,10 @@ pub fn instantiate_type_generic(
                 LuaType::SelfInfer
             }
         }
+        LuaType::TypeGuard(guard) => {
+            let inner = instantiate_type_generic(db, guard.deref(), substitutor);
+            LuaType::TypeGuard(inner.into())
+        }
         _ => ty.clone(),
     }
 }
diff --git a/crates/emmylua_code_analysis/src/semantic/infer/infer_call/mod.rs b/crates/emmylua_code_analysis/src/semantic/infer/infer_call/mod.rs
@@ -107,15 +107,19 @@ pub fn infer_call_expr_func(
     };
 
     let result = if let Ok(func_ty) = result {
-        unwrapp_return_type(db, cache, func_ty.get_ret().clone(), call_expr).map(|new_ret| {
-            LuaFunctionType::new(
-                func_ty.is_async(),
-                func_ty.is_colon_define(),
-                func_ty.get_params().to_vec(),
-                new_ret,
-            )
-            .into()
-        })
+        let func_ret = func_ty.get_ret();
+        match func_ret {
+            LuaType::TypeGuard(_) => Ok(func_ty),
+            _ => unwrapp_return_type(db, cache, func_ret.clone(), call_expr).map(|new_ret| {
+                LuaFunctionType::new(
+                    func_ty.is_async(),
+                    func_ty.is_colon_define(),
+                    func_ty.get_params().to_vec(),
+                    new_ret,
+                )
+                .into()
+            }),
+        }
     } else {
         result
     };
diff --git a/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/binary_flow.rs b/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/binary_flow.rs
@@ -1,5 +1,5 @@
 use emmylua_parser::{
-    BinaryOperator, LuaBinaryExpr, LuaCallExpr, LuaChunk, LuaExpr, LuaIndexMemberExpr,
+    BinaryOperator, LuaAstNode, LuaBinaryExpr, LuaCallExpr, LuaChunk, LuaExpr, LuaIndexMemberExpr,
     LuaLiteralToken,
 };
 
@@ -128,6 +128,38 @@ fn maybe_type_guard_binary(
                 }
             }
         }
+        // may ref a type value
+    } else if let LuaExpr::NameExpr(name_expr) = left_expr {
+        if let LuaExpr::LiteralExpr(literal_expr) = right_expr {
+            let Some(decl_id) = db
+                .get_reference_index()
+                .get_var_reference_decl(&cache.get_file_id(), name_expr.get_range())
+            else {
+                return Ok(ResultTypeOrContinue::Continue);
+            };
+
+            let Some(expr_ptr) = tree.get_decl_ref_expr(&decl_id) else {
+                return Ok(ResultTypeOrContinue::Continue);
+            };
+
+            let Some(expr) = expr_ptr.to_node(root) else {
+                return Ok(ResultTypeOrContinue::Continue);
+            };
+
+            if let LuaExpr::CallExpr(call_expr) = expr {
+                if call_expr.is_type() {
+                    type_guard_expr = Some(call_expr);
+                    match literal_expr.get_literal() {
+                        Some(LuaLiteralToken::String(s)) => {
+                            literal_string = s.get_value();
+                        }
+                        _ => return Ok(ResultTypeOrContinue::Continue),
+                    }
+                }
+            } else {
+                return Ok(ResultTypeOrContinue::Continue);
+            }
+        }
     }
 
     if type_guard_expr.is_none() || literal_string.is_empty() {
diff --git a/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/call_flow.rs b/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/call_flow.rs
diff --git a/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/mod.rs b/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -291,6 +291,7 @@ impl LuaType {`
`291`	`291`	`LuaType::Nil \| LuaType::Boolean \| LuaType::Any \| LuaType::Unknown => false,`
`292`	`292`	`LuaType::BooleanConst(boolean) \| LuaType::DocBooleanConst(boolean) => boolean.clone(),`
`293`	`293`	`LuaType::Union(u) => u.is_always_truthy(),`
	`294`	`+ LuaType::TypeGuard(_) => false,`
`294`	`295`	`_ => true,`
`295`	`296`	`}`
`296`	`297`	`}`
`@@ -299,6 +300,7 @@ impl LuaType {`
`299`	`300`	`match self {`
`300`	`301`	`LuaType::Nil \| LuaType::BooleanConst(false) \| LuaType::DocBooleanConst(false) => true,`
`301`	`302`	`LuaType::Union(u) => u.is_always_falsy(),`
	`303`	`+ LuaType::TypeGuard(_) => false,`
`302`	`304`	`_ => false,`
`303`	`305`	`}`
`304`	`306`	`}`
`@@ -400,6 +402,7 @@ impl LuaType {`
`400`	`402`	`LuaType::StrTplRef(_) => true,`
`401`	`403`	`LuaType::SelfInfer => true,`
`402`	`404`	`LuaType::MultiLineUnion(inner) => inner.contain_tpl(),`
	`405`	`+ LuaType::TypeGuard(inner) => inner.contain_tpl(),`
`403`	`406`	`_ => false,`
`404`	`407`	`}`
`405`	`408`	`}`
Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,10 @@ pub fn instantiate_type_generic(`
`45`	`45`	`LuaType::SelfInfer`
`46`	`46`	`}`
`47`	`47`	`}`
	`48`	`+ LuaType::TypeGuard(guard) => {`
	`49`	`+ let inner = instantiate_type_generic(db, guard.deref(), substitutor);`
	`50`	`+ LuaType::TypeGuard(inner.into())`
	`51`	`+ }`
`48`	`52`	`_ => ty.clone(),`
`49`	`53`	`}`
`50`	`54`	`}`