fix: respect cast flow in assign type mismatch diagnostics

Author: xuhuanzy

crates/emmylua_code_analysis/src/diagnostic/checker/assign_type_mismatch.rs

@@ -7,9 +7,9 @@ use emmylua_parser::{
 use rowan::{NodeOrToken, TextRange};
 
 use crate::{
-    DiagnosticCode, LuaBuiltinAttributeKind, LuaDeclExtra, LuaDeclId, LuaMemberKey,
-    LuaSemanticDeclId, LuaType, SemanticDeclLevel, SemanticModel, TypeCheckFailReason,
-    TypeCheckResult, VariadicType, infer_index_expr,
+    DiagnosticCode, LuaBuiltinAttributeKind, LuaDeclId, LuaMemberKey, LuaType, SemanticDeclLevel,
+    SemanticModel, TypeCheckFailReason, TypeCheckResult, VariadicType,
+    semantic::infer_expr_without_condition_flow,
 };
 
 use super::{Checker, DiagnosticContext, humanize_lint_type};
@@ -74,36 +74,13 @@ fn check_name_expr(
     expr: Option<LuaExpr>,
     value_type: LuaType,
 ) -> Option<()> {
-    let semantic_decl = semantic_model.find_decl(
-        NodeOrToken::Node(name_expr.syntax().clone()),
-        SemanticDeclLevel::default(),
-    )?;
-    let source_type = match semantic_decl.clone() {
-        LuaSemanticDeclId::LuaDecl(decl_id) => {
-            let decl = semantic_model
-                .get_db()
-                .get_decl_index()
-                .get_decl(&decl_id)?;
-            match decl.extra {
-                LuaDeclExtra::Param {
-                    idx, signature_id, ..
-                } => {
-                    let signature = semantic_model
-                        .get_db()
-                        .get_signature_index()
-                        .get(&signature_id)?;
-                    let param_type = signature.get_param_info_by_id(idx)?;
-                    Some(param_type.type_ref.clone())
-                }
-                _ => semantic_model
-                    .get_db()
-                    .get_type_index()
-                    .get_type_cache(&decl_id.into())
-                    .map(|cache| cache.as_type().clone()),
-            }
-        }
-        _ => None,
-    };
+    let source_type = infer_expr_without_condition_flow(
+        semantic_model.get_db(),
+        &mut semantic_model.get_cache().borrow_mut(),
+        LuaExpr::NameExpr(name_expr.clone()),
+    )
+    .ok();
+
     check_assign_type_mismatch(
         context,
         semantic_model,
@@ -132,11 +109,10 @@ fn check_index_expr(
     expr: Option<LuaExpr>,
     value_type: LuaType,
 ) -> Option<()> {
-    let source_type = infer_index_expr(
+    let source_type = infer_expr_without_condition_flow(
         semantic_model.get_db(),
         &mut semantic_model.get_cache().borrow_mut(),
-        index_expr.clone(),
-        false,
+        LuaExpr::IndexExpr(index_expr.clone()),
     )
     .ok();
 

crates/emmylua_code_analysis/src/diagnostic/test/assign_type_mismatch_test.rs

@@ -837,6 +837,106 @@ return t
         ));
     }
 
+    #[test]
+    fn test_cast_add_type_flow_assignment() {
+        let mut ws = VirtualWorkspace::new();
+        assert!(ws.has_no_diagnostic(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+                ---@type string | number
+                local val
+
+                ---@cast val + boolean
+                val = true
+        "#,
+        ));
+
+        assert!(ws.has_no_diagnostic(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+                ---@type string | number
+                local val
+
+                ---@cast val boolean
+                val = true
+        "#,
+        ));
+    }
+
+    #[test]
+    fn test_cast_remove_type_flow_assignment() {
+        let mut ws = VirtualWorkspace::new();
+        assert!(!ws.has_no_diagnostic(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+                ---@type string | number | boolean
+                local val
+
+                ---@cast val - boolean
+                val = true
+        "#,
+        ));
+
+        assert!(!ws.has_no_diagnostic(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+                ---@type string | number
+                local val
+
+                ---@cast val boolean
+                val = 1
+        "#,
+        ));
+    }
+
+    #[test]
+    fn test_cast_type_flow_index_assignment() {
+        let mut ws = VirtualWorkspace::new();
+        assert!(ws.has_no_diagnostic(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+                ---@class CastFlowBox
+                ---@field val string | number
+
+                ---@type CastFlowBox
+                local box
+
+                ---@cast box.val + boolean
+                box.val = true
+        "#,
+        ));
+
+        assert!(!ws.has_no_diagnostic(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+                ---@class CastFlowBox
+                ---@field val string | number | boolean
+
+                ---@type CastFlowBox
+                local box
+
+                ---@cast box.val - boolean
+                box.val = true
+        "#,
+        ));
+    }
+
+    #[test]
+    fn test_condition_flow_does_not_narrow_assignment_slot() {
+        let mut ws = VirtualWorkspace::new();
+        assert!(ws.has_no_diagnostic(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+                ---@type string | number
+                local val
+
+                if type(val) == "string" then
+                    val = 1
+                end
+        "#,
+        ));
+    }
+
     #[test]
     fn test_flow_1() {
         let mut ws = VirtualWorkspace::new();

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

@@ -28,6 +28,7 @@ use infer_table::infer_table_expr;
 pub use infer_table::{infer_table_field_value_should_be, infer_table_should_be};
 use infer_unary::infer_unary_expr;
 pub use narrow::VarRefId;
+pub(crate) use narrow::infer_expr_without_condition_flow;
 pub(in crate::semantic) use narrow::{ConditionFlowAction, InferConditionFlow};
 
 use rowan::TextRange;

crates/emmylua_code_analysis/src/semantic/infer/narrow/get_type_at_flow.rs

@@ -421,10 +421,12 @@ struct FlowTypeEngine<'a> {
 }
 
 impl<'a> FlowTypeEngine<'a> {
-    fn run(&mut self, var_ref_id: &VarRefId, flow_id: FlowId) -> InferResult {
+    fn run(&mut self, var_ref_id: &VarRefId, flow_id: FlowId, mode: FlowMode) -> InferResult {
         let mut stack = Vec::new();
+        let mut query = FlowQuery::new(self.cache, var_ref_id, flow_id);
+        query.mode = mode;
         let mut step = SchedulerStep::StartQuery {
-            query: FlowQuery::new(self.cache, var_ref_id, flow_id),
+            query,
             continuation: None,
         };
 
@@ -1641,7 +1643,24 @@ pub(super) fn get_type_at_flow(
         cache,
         root,
     }
-    .run(var_ref_id, flow_id)
+    .run(var_ref_id, flow_id, FlowMode::WithConditions)
+}
+
+pub(super) fn get_type_at_flow_without_conditions(
+    db: &DbIndex,
+    tree: &FlowTree,
+    cache: &mut LuaInferCache,
+    root: &LuaChunk,
+    var_ref_id: &VarRefId,
+    flow_id: FlowId,
+) -> InferResult {
+    FlowTypeEngine {
+        db,
+        tree,
+        cache,
+        root,
+    }
+    .run(var_ref_id, flow_id, FlowMode::WithoutConditions)
 }
 
 fn get_flow_cache_var_ref_id(cache: &mut LuaInferCache, var_ref_id: &VarRefId) -> u32 {

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

@@ -37,6 +37,45 @@ pub fn infer_expr_narrow_type(
     get_type_at_flow::get_type_at_flow(db, flow_tree, cache, &root, &var_ref_id, flow_id)
 }
 
+pub(crate) fn infer_expr_without_condition_flow(
+    db: &DbIndex,
+    cache: &mut LuaInferCache,
+    expr: LuaExpr,
+) -> InferResult {
+    let Some(var_ref_id) = get_var_expr_var_ref_id(db, cache, expr.clone()) else {
+        return Err(InferFailReason::None);
+    };
+
+    if let LuaExpr::IndexExpr(index_expr) = &expr
+        && !cache.index_ref_origin_type_cache.contains_key(&var_ref_id)
+        && let Ok(origin_type) =
+            super::infer_index::infer_index_expr(db, cache, index_expr.clone(), false)
+    {
+        cache
+            .index_ref_origin_type_cache
+            .insert(var_ref_id.clone(), CacheEntry::Cache(origin_type));
+    }
+
+    let file_id = cache.get_file_id();
+    let Some(flow_tree) = db.get_flow_index().get_flow_tree(&file_id) else {
+        return get_var_ref_type(db, cache, &var_ref_id);
+    };
+
+    let Some(flow_id) = flow_tree.get_flow_id(expr.get_syntax_id()) else {
+        return get_var_ref_type(db, cache, &var_ref_id);
+    };
+
+    let root = LuaChunk::cast(expr.get_root()).ok_or(InferFailReason::None)?;
+    get_type_at_flow::get_type_at_flow_without_conditions(
+        db,
+        flow_tree,
+        cache,
+        &root,
+        &var_ref_id,
+        flow_id,
+    )
+}
+
 pub(in crate::semantic) fn get_var_ref_type(
     db: &DbIndex,
     cache: &mut LuaInferCache,

crates/emmylua_code_analysis/src/semantic/mod.rs

@@ -55,6 +55,7 @@ pub use generic::*;
 pub use guard::{InferGuard, InferGuardRef};
 pub use infer::InferFailReason;
 pub use infer::infer_call_expr_func;
+pub(crate) use infer::infer_expr_without_condition_flow;
 pub use infer::infer_param;
 pub(crate) use infer::try_infer_expr_for_index;
 pub(crate) use infer::{infer_expr, try_infer_expr_no_flow};