Fix #3587

1 files changed, 28 insertions, 1 deletions

diff --git a/src/check_stmt.cpp b/src/check_stmt.cpp
index ee55ff0d7..875503874 100644
--- a/src/check_stmt.cpp
+++ b/src/check_stmt.cpp
@@ -169,9 +169,16 @@ gb_internal bool check_has_break_list(Slice<Ast *> const &stmts, String const &l
 	return false;
 }
 
+gb_internal bool check_has_break_expr(Ast * expr, String const &label) {
+	if (expr && expr->viral_state_flags & ViralStateFlag_ContainsOrBreak) {
+		return true;
+	}
+	return false;
+}
+
 gb_internal bool check_has_break_expr_list(Slice<Ast *> const &exprs, String const &label) {
 	for (Ast *expr : exprs) {
-		if (expr && expr->viral_state_flags & ViralStateFlag_ContainsOrBreak) {
+		if (check_has_break_expr(expr, label)) {
 			return true;
 		}
 	}
@@ -196,6 +203,13 @@ gb_internal bool check_has_break(Ast *stmt, String const &label, bool implicit)
 		return check_has_break_list(stmt->BlockStmt.stmts, label, implicit);
 
 	case Ast_IfStmt:
+		if (stmt->IfStmt.init && check_has_break(stmt->IfStmt.init, label, implicit)) {
+			return true;
+		}
+		if (stmt->IfStmt.cond && check_has_break_expr(stmt->IfStmt.cond, label)) {
+			return true;
+		}
+
 		if (check_has_break(stmt->IfStmt.body, label, implicit) ||
 		    (stmt->IfStmt.else_stmt != nullptr && check_has_break(stmt->IfStmt.else_stmt, label, implicit))) {
 			return true;
@@ -206,6 +220,9 @@ gb_internal bool check_has_break(Ast *stmt, String const &label, bool implicit)
 		return check_has_break_list(stmt->CaseClause.stmts, label, implicit);
 
 	case Ast_SwitchStmt:
+		if (stmt->SwitchStmt.init && check_has_break_expr(stmt->SwitchStmt.init, label)) {
+			return true;
+		}
 		if (label != "" && check_has_break(stmt->SwitchStmt.body, label, false)) {
 			return true;
 		}
@@ -218,6 +235,16 @@ gb_internal bool check_has_break(Ast *stmt, String const &label, bool implicit)
 		break;
 
 	case Ast_ForStmt:
+		if (stmt->ForStmt.init && check_has_break(stmt->ForStmt.init, label, implicit)) {
+			return true;
+		}
+		if (stmt->ForStmt.cond && check_has_break_expr(stmt->ForStmt.cond, label)) {
+			return true;
+		}
+		if (stmt->ForStmt.post && check_has_break(stmt->ForStmt.post, label, implicit)) {
+			return true;
+		}
+
 		if (label != "" && check_has_break(stmt->ForStmt.body, label, false)) {
 			return true;
 		}