Support complex assert narrow

Fix #583

Author: CppCXY

crates/emmylua_code_analysis/src/compilation/analyzer/flow/bind_analyze/stats.rs

@@ -1,7 +1,7 @@
 use emmylua_parser::{
-    LuaAssignStat, LuaAst, LuaAstNode, LuaBlock, LuaBreakStat, LuaCallExprStat, LuaDoStat,
-    LuaForRangeStat, LuaForStat, LuaFuncStat, LuaGotoStat, LuaIfStat, LuaLabelStat, LuaLocalStat,
-    LuaRepeatStat, LuaReturnStat, LuaWhileStat,
+    LuaAssignStat, LuaAst, LuaAstNode, LuaBlock, LuaBreakStat, LuaCallArgList, LuaCallExprStat,
+    LuaDoStat, LuaForRangeStat, LuaForStat, LuaFuncStat, LuaGotoStat, LuaIfStat, LuaLabelStat,
+    LuaLocalStat, LuaRepeatStat, LuaReturnStat, LuaWhileStat,
 };
 
 use crate::{
@@ -78,23 +78,37 @@ pub fn bind_call_expr_stat(
         None => return current, // If there's no call expression, just return the current flow
     };
 
-    if let Some(ast) = LuaAst::cast(call_expr.syntax().clone()) {
-        bind_each_child(binder, ast, current);
-    }
-
     if call_expr.is_assert() {
-        let assert_flow_id = binder.create_node(FlowNodeKind::AssertCall(call_expr.to_ptr()));
-        binder.add_antecedent(assert_flow_id, current);
-        assert_flow_id
+        let Some(arg_list) = call_expr.get_args_list() else {
+            return current; // If there's no argument list, just return the current flow
+        };
+
+        bind_assert_stat(binder, arg_list, current)
     } else if call_expr.is_error() {
         let return_flow_id = binder.create_return();
         binder.add_antecedent(return_flow_id, current);
         return_flow_id
     } else {
+        if let Some(ast) = LuaAst::cast(call_expr.syntax().clone()) {
+            bind_each_child(binder, ast, current);
+        }
         current
     }
 }
 
+fn bind_assert_stat(binder: &mut FlowBinder, arg_list: LuaCallArgList, current: FlowId) -> FlowId {
+    let false_target = binder.unreachable;
+
+    let mut pre_arg = current;
+    for arg in arg_list.get_args() {
+        let pre_next_arg = binder.create_branch_label();
+        bind_condition_expr(binder, arg, pre_arg, pre_next_arg, false_target);
+        pre_arg = finish_flow_label(binder, pre_next_arg, pre_arg);
+    }
+
+    pre_arg
+}
+
 pub fn bind_label_stat(
     binder: &mut FlowBinder,
     label_stat: LuaLabelStat,

crates/emmylua_code_analysis/src/compilation/test/flow.rs

@@ -850,4 +850,23 @@ end
         let b_expected = ws.ty("B");
         assert_eq!(b, b_expected);
     }
+
+    #[test]
+    fn test_issue_583() {
+        let mut ws = VirtualWorkspace::new_with_init_std_lib();
+
+        ws.check_code_for(
+            DiagnosticCode::AssignTypeMismatch,
+            r#"
+            --- @param sha string
+            local function get_hash_color(sha)
+            local r, g, b = sha:match('(%x)%x(%x)%x(%x)')
+            assert(r and g and b, 'Invalid hash color')
+            local _ = r --- @type string
+            local _ = g --- @type string
+            local _ = b --- @type string
+            end
+            "#,
+        );
+    }
 }

crates/emmylua_code_analysis/src/db_index/flow/flow_node.rs

@@ -1,6 +1,6 @@
 use emmylua_parser::{
-    LuaAssignStat, LuaAstNode, LuaAstPtr, LuaCallExpr, LuaChunk, LuaClosureExpr, LuaDocTagCast,
-    LuaExpr, LuaForStat, LuaSyntaxKind, LuaSyntaxNode,
+    LuaAssignStat, LuaAstNode, LuaAstPtr, LuaChunk, LuaClosureExpr, LuaDocTagCast, LuaExpr,
+    LuaForStat, LuaSyntaxKind, LuaSyntaxNode,
 };
 use internment::ArcIntern;
 use rowan::{TextRange, TextSize};
@@ -52,8 +52,6 @@ pub enum FlowNodeKind {
     ForIStat(LuaAstPtr<LuaForStat>),
     /// Tag cast comment
     TagCast(LuaAstPtr<LuaDocTagCast>),
-    /// Assert call
-    AssertCall(LuaAstPtr<LuaCallExpr>),
     /// Break statement
     Break,
     /// Return statement

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

@@ -1,4 +1,4 @@
-use emmylua_parser::{LuaAssignStat, LuaAstNode, LuaCallExpr, LuaChunk, LuaVarExpr};
+use emmylua_parser::{LuaAssignStat, LuaAstNode, LuaChunk, LuaVarExpr};
 
 use crate::{
     infer_expr,
@@ -8,7 +8,7 @@ use crate::{
             get_multi_antecedents, get_single_antecedent,
             get_type_at_cast_flow::get_type_at_cast_flow,
             get_var_ref_type,
-            narrow_type::{narrow_down_type, remove_false_or_nil},
+            narrow_type::narrow_down_type,
             var_ref_id::get_var_expr_var_ref_id,
             ResultTypeOrContinue,
         },
@@ -138,25 +138,6 @@ pub fn get_type_at_flow(
                     antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
                 }
             }
-            FlowNodeKind::AssertCall(lua_ast_ptr) => {
-                let assert_call = lua_ast_ptr.to_node(root).ok_or(InferFailReason::None)?;
-                let result_or_continue = get_type_at_assert_call(
-                    db,
-                    tree,
-                    cache,
-                    root,
-                    var_ref_id,
-                    flow_node,
-                    assert_call,
-                )?;
-
-                if let ResultTypeOrContinue::Result(assert_type) = result_or_continue {
-                    result_type = assert_type;
-                    break;
-                } else {
-                    antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
-                }
-            }
         }
     }
 
@@ -256,33 +237,3 @@ fn get_type_at_assign_stat(
 
     Ok(ResultTypeOrContinue::Continue)
 }
-
-fn get_type_at_assert_call(
-    db: &DbIndex,
-    tree: &FlowTree,
-    cache: &mut LuaInferCache,
-    root: &LuaChunk,
-    var_ref_id: &VarRefId,
-    flow_node: &FlowNode,
-    assert_call: LuaCallExpr,
-) -> Result<ResultTypeOrContinue, InferFailReason> {
-    let call_arg_list = match assert_call.get_args_list() {
-        Some(args) => args,
-        None => return Ok(ResultTypeOrContinue::Continue),
-    };
-
-    for arg in call_arg_list.get_args() {
-        if let Some(ref_decl_id) = get_var_expr_var_ref_id(db, cache, arg.clone()) {
-            if ref_decl_id == *var_ref_id {
-                let antecedent_flow_id = get_single_antecedent(tree, flow_node)?;
-                let antecedent_type =
-                    get_type_at_flow(db, tree, cache, root, var_ref_id, antecedent_flow_id)?;
-                let result_type = remove_false_or_nil(antecedent_type);
-
-                return Ok(ResultTypeOrContinue::Result(result_type));
-            }
-        }
-    }
-
-    Ok(ResultTypeOrContinue::Continue)
-}