support inherit flow infer from local const

CppCXY · CppCXY · commit 5f6351ce1530 · 2025-07-10T16:09:35.000+08:00
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,32 @@ 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);
+    }
 }
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/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/mod.rs b/crates/emmylua_code_analysis/src/semantic/infer/narrow/condition_flow/mod.rs
@@ -2,7 +2,7 @@ mod binary_flow;
 mod call_flow;
 mod index_flow;
 
-use emmylua_parser::{LuaChunk, LuaExpr, LuaNameExpr, LuaUnaryExpr, UnaryOperator};
+use emmylua_parser::{LuaAstNode, LuaChunk, LuaExpr, LuaNameExpr, LuaUnaryExpr, UnaryOperator};
 
 use crate::{
     semantic::infer::{
@@ -146,7 +146,16 @@ fn get_type_at_name_expr(
     };
 
     if name_var_ref_id != *var_ref_id {
-        return Ok(ResultTypeOrContinue::Continue);
+        return get_type_at_name_ref(
+            db,
+            tree,
+            cache,
+            root,
+            var_ref_id,
+            flow_node,
+            name_expr,
+            condition_flow,
+        );
     }
 
     let antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
@@ -160,6 +169,43 @@ fn get_type_at_name_expr(
     Ok(ResultTypeOrContinue::Result(result_type))
 }
 
+fn get_type_at_name_ref(
+    db: &DbIndex,
+    tree: &FlowTree,
+    cache: &mut LuaInferCache,
+    root: &LuaChunk,
+    var_ref_id: &VarRefId,
+    flow_node: &FlowNode,
+    name_expr: LuaNameExpr,
+    condition_flow: InferConditionFlow,
+) -> Result<ResultTypeOrContinue, InferFailReason> {
+    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);
+    };
+
+    get_type_at_condition_flow(
+        db,
+        tree,
+        cache,
+        root,
+        var_ref_id,
+        flow_node,
+        expr,
+        condition_flow,
+    )
+}
+
 fn get_type_at_unary_flow(
     db: &DbIndex,
     tree: &FlowTree,
