[red-knot] Semantic index: handle invalid breaks (#14522)

## Summary

This fix addresses panics related to invalid syntax like the following
where a `break` statement is used in a nested definition inside a
loop:

```py
while True:

    def b():
        x: int

        break
```

closes #14342

## Test Plan

* New corpus regression tests.
* New unit test to make sure we handle nested while loops correctly.
This test is passing on `main`, but can easily fail if the
`is_inside_loop` state isn't properly saved/restored.

crates/red_knot_python_semantic/resources/mdtest/loops/while_loop.md

@@ -52,3 +52,29 @@ else:
 reveal_type(x)  # revealed: Literal[2, 3]
 reveal_type(y)  # revealed: Literal[1, 2, 4]
 ```
+
+## Nested while loops
+
+```py
+def flag() -> bool:
+    return True
+
+x = 1
+
+while flag():
+    x = 2
+
+    while flag():
+        x = 3
+        if flag():
+            break
+    else:
+        x = 4
+
+    if flag():
+        break
+else:
+    x = 5
+
+reveal_type(x)  # revealed: Literal[3, 4, 5]
+```

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -36,12 +36,25 @@ use super::definition::{
 
 mod except_handlers;
 
+/// Are we in a state where a `break` statement is allowed?
+#[derive(Clone, Copy, Debug)]
+enum LoopState {
+    InLoop,
+    NotInLoop,
+}
+
+impl LoopState {
+    fn is_inside(self) -> bool {
+        matches!(self, LoopState::InLoop)
+    }
+}
+
 pub(super) struct SemanticIndexBuilder<'db> {
     // Builder state
     db: &'db dyn Db,
     file: File,
     module: &'db ParsedModule,
-    scope_stack: Vec<FileScopeId>,
+    scope_stack: Vec<(FileScopeId, LoopState)>,
     /// The assignments we're currently visiting, with
     /// the most recent visit at the end of the Vec
     current_assignments: Vec<CurrentAssignment<'db>>,
@@ -103,9 +116,24 @@ impl<'db> SemanticIndexBuilder<'db> {
         *self
             .scope_stack
             .last()
+            .map(|(scope, _)| scope)
             .expect("Always to have a root scope")
     }
 
+    fn loop_state(&self) -> LoopState {
+        self.scope_stack
+            .last()
+            .expect("Always to have a root scope")
+            .1
+    }
+
+    fn set_inside_loop(&mut self, state: LoopState) {
+        self.scope_stack
+            .last_mut()
+            .expect("Always to have a root scope")
+            .1 = state;
+    }
+
     fn push_scope(&mut self, node: NodeWithScopeRef) {
         let parent = self.current_scope();
         self.push_scope_with_parent(node, Some(parent));
@@ -136,11 +164,11 @@ impl<'db> SemanticIndexBuilder<'db> {
 
         debug_assert_eq!(ast_id_scope, file_scope_id);
 
-        self.scope_stack.push(file_scope_id);
+        self.scope_stack.push((file_scope_id, LoopState::NotInLoop));
     }
 
     fn pop_scope(&mut self) -> FileScopeId {
-        let id = self.scope_stack.pop().expect("Root scope to be present");
+        let (id, _) = self.scope_stack.pop().expect("Root scope to be present");
         let children_end = self.scopes.next_index();
         let scope = &mut self.scopes[id];
         scope.descendents = scope.descendents.start..children_end;
@@ -785,7 +813,10 @@ where
 
                 // TODO: definitions created inside the body should be fully visible
                 // to other statements/expressions inside the body --Alex/Carl
+                let outer_loop_state = self.loop_state();
+                self.set_inside_loop(LoopState::InLoop);
                 self.visit_body(body);
+                self.set_inside_loop(outer_loop_state);
 
                 // Get the break states from the body of this loop, and restore the saved outer
                 // ones.
@@ -824,7 +855,9 @@ where
                 self.visit_body(body);
             }
             ast::Stmt::Break(_) => {
-                self.loop_break_states.push(self.flow_snapshot());
+                if self.loop_state().is_inside() {
+                    self.loop_break_states.push(self.flow_snapshot());
+                }
             }
 
             ast::Stmt::For(
@@ -851,7 +884,10 @@ where
                 // TODO: Definitions created by loop variables
                 // (and definitions created inside the body)
                 // are fully visible to other statements/expressions inside the body --Alex/Carl
+                let outer_loop_state = self.loop_state();
+                self.set_inside_loop(LoopState::InLoop);
                 self.visit_body(body);
+                self.set_inside_loop(outer_loop_state);
 
                 let break_states =
                     std::mem::replace(&mut self.loop_break_states, saved_break_states);

crates/red_knot_workspace/resources/test/corpus/15_while_break_invalid_in_class.py

@@ -0,0 +1,6 @@
+while True:
+
+    class A:
+        x: int
+
+        break

crates/red_knot_workspace/resources/test/corpus/15_while_break_invalid_in_func.py

@@ -0,0 +1,6 @@
+while True:
+
+    def b():
+        x: int
+
+        break

crates/red_knot_workspace/resources/test/corpus/16_for_break_invalid_in_class.py

@@ -0,0 +1,6 @@
+for _ in range(1):
+
+    class A:
+        x: int
+
+        break

crates/red_knot_workspace/resources/test/corpus/16_for_break_invalid_in_func.py

@@ -0,0 +1,6 @@
+for _ in range(1):
+
+    def b():
+        x: int
+
+        break