Fix false diagnostics and fixes for certain loops

crates/cairo-lint-core/src/lints/loops/loop_for_while.rs

@@ -78,25 +78,28 @@ fn check_single_loop_for_while(
     let Expr::Block(block_expr) = &arenas.exprs[loop_expr.body] else {
         return;
     };
-    // Checks if one of the statements is an if expression that only contains a break instruction
-    for statement in &block_expr.statements {
-        if_chain! {
-            if let Statement::Expr(ref expr_statement) = arenas.statements[*statement];
-            if check_if_contains_break(&expr_statement.expr, arenas);
-            then {
-                diagnostics.push(PluginDiagnostic {
-                    stable_ptr: loop_expr.stable_ptr.untyped(),
-                    message: LoopForWhile.diagnostic_message().to_string(),
-                    severity: Severity::Warning,
-                });
-            }
+
+    // Checks if the first statement is an if expression that only contains a break instruction.
+    if_chain! {
+        if let Some(statement) = block_expr.statements.first();
+        if let Statement::Expr(ref expr_statement) = arenas.statements[*statement];
+        if check_if_contains_break_with_no_return_value(&expr_statement.expr, arenas);
+        then {
+            diagnostics.push(PluginDiagnostic {
+                stable_ptr: loop_expr.stable_ptr.untyped(),
+                message: LoopForWhile.diagnostic_message().to_string(),
+                severity: Severity::Warning,
+            });
         }
     }
 
     // Do the same thing if the if is in the tail of the block
     if_chain! {
+        // Loop with single if-else statement will only have a tail expr and the statements will be empty.
+        // We check if the tail if statement is a single one. If it's not, we ignore the loop as a whole.
+        if block_expr.statements.is_empty();
         if let Some(tail_expr) = block_expr.tail;
-        if check_if_contains_break(&tail_expr, arenas);
+        if check_if_contains_break_with_no_return_value(&tail_expr, arenas);
         then {
             diagnostics.push(PluginDiagnostic {
                 stable_ptr: loop_expr.stable_ptr.untyped(),
@@ -107,7 +110,7 @@ fn check_single_loop_for_while(
     }
 }
 
-fn check_if_contains_break(expr: &ExprId, arenas: &Arenas) -> bool {
+fn check_if_contains_break_with_no_return_value(expr: &ExprId, arenas: &Arenas) -> bool {
     if_chain! {
         // Is an if expression
         if let Expr::If(ref if_expr) = arenas.exprs[*expr];
@@ -116,9 +119,10 @@ fn check_if_contains_break(expr: &ExprId, arenas: &Arenas) -> bool {
         // Get the first statement of the if
         if let Some(inner_stmt) = if_block.statements.first();
         // Is it a break statement
-        if matches!(arenas.statements[*inner_stmt], Statement::Break(_));
+        if let Statement::Break(break_expr) = &arenas.statements[*inner_stmt];
         then {
-            return true;
+            // If break also has a return value like `break 1;` then it's not a simple break.
+            return break_expr.expr_option.is_none();
         }
     }
     false

crates/cairo-lint-core/tests/loops/loop_for_while.rs

@@ -144,6 +144,56 @@ fn main() {
 }
 "#;
 
+const ADVANCED_LOOP_WITH_BREAK_IN_THE_MIDDLE: &str = r#"
+fn main() -> u32 {
+    let mut exponent: u32 = 3;
+    let two: u32 = 2;
+    let mut result: u32 = 0;
+    let mut base: u32 = 0;
+    loop {
+        if exponent % two != 0 {
+            let new_result = 10;
+            result = new_result;
+        }
+
+        exponent = exponent / two;
+
+        if exponent == 0 {
+            break result;
+        }
+
+        let new_base = 2;
+
+        base = new_base;
+    };
+    1
+}
+"#;
+
+const SIMPLE_LOOP_WITH_BREAK_AT_THE_END: &str = r#"
+fn main() {
+    let mut x: u16 = 0;
+    loop {
+        x += 1;
+        if x == 10 {
+            break;
+        }
+    }
+}
+"#;
+
+const SIMPLE_LOOP_WITH_BREAK_WITH_RETURN_VALUE: &str = r#"
+fn main() -> u16 {
+    let mut x: u16 = 0;
+    loop {
+        if x == 10 {
+            break x;
+        }
+        x += 1;
+    }
+}
+"#;
+
 #[test]
 fn simple_loop_with_break_diagnostics() {
     test_lint_diagnostics!(SIMPLE_LOOP_WITH_BREAK, @r"
@@ -447,3 +497,77 @@ fn loop_with_condition_depending_on_external_variable_fixer() {
     }
     "#);
 }
+
+#[test]
+fn advanced_loop_with_break_in_the_middle_diagnostics() {
+    test_lint_diagnostics!(ADVANCED_LOOP_WITH_BREAK_IN_THE_MIDDLE, @"");
+}
+
+#[test]
+fn advanced_loop_with_break_in_the_middle_fixer() {
+    test_lint_fixer!(ADVANCED_LOOP_WITH_BREAK_IN_THE_MIDDLE, @r"
+    fn main() -> u32 {
+        let mut exponent: u32 = 3;
+        let two: u32 = 2;
+        let mut result: u32 = 0;
+        let mut base: u32 = 0;
+        loop {
+            if exponent % two != 0 {
+                let new_result = 10;
+                result = new_result;
+            }
+
+            exponent = exponent / two;
+
+            if exponent == 0 {
+                break result;
+            }
+
+            let new_base = 2;
+
+            base = new_base;
+        };
+        1
+    }
+    ");
+}
+
+#[test]
+fn simple_loop_with_break_at_the_end_diagnostics() {
+    test_lint_diagnostics!(SIMPLE_LOOP_WITH_BREAK_AT_THE_END, @r"");
+}
+
+#[test]
+fn simple_loop_with_break_at_the_end_fixer() {
+    test_lint_fixer!(SIMPLE_LOOP_WITH_BREAK_AT_THE_END, @r"
+    fn main() {
+        let mut x: u16 = 0;
+        loop {
+            x += 1;
+            if x == 10 {
+                break;
+            }
+        }
+    }
+    ");
+}
+
+#[test]
+fn simple_loop_with_break_with_return_value_diagnostics() {
+    test_lint_diagnostics!(SIMPLE_LOOP_WITH_BREAK_WITH_RETURN_VALUE, @"");
+}
+
+#[test]
+fn simple_loop_with_break_with_return_value_fixer() {
+    test_lint_fixer!(SIMPLE_LOOP_WITH_BREAK_WITH_RETURN_VALUE, @r"
+    fn main() -> u16 {
+        let mut x: u16 = 0;
+        loop {
+            if x == 10 {
+                break x;
+            }
+            x += 1;
+        }
+    }
+    ");
+}