diff --git a/crates/red_knot_python_semantic/resources/mdtest/loops/while_loop.md b/crates/red_knot_python_semantic/resources/mdtest/loops/while_loop.md index 51a123edc2..a03c2cf83f 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/loops/while_loop.md +++ b/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] +``` diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index a1c7a5e0ed..dfca6328db 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/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, + 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>, @@ -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); diff --git a/crates/red_knot_workspace/resources/test/corpus/15_while_break_invalid_in_class.py b/crates/red_knot_workspace/resources/test/corpus/15_while_break_invalid_in_class.py new file mode 100644 index 0000000000..934ec4dc17 --- /dev/null +++ b/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 diff --git a/crates/red_knot_workspace/resources/test/corpus/15_while_break_invalid_in_func.py b/crates/red_knot_workspace/resources/test/corpus/15_while_break_invalid_in_func.py new file mode 100644 index 0000000000..83e818f80e --- /dev/null +++ b/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 diff --git a/crates/red_knot_workspace/resources/test/corpus/16_for_break_invalid_in_class.py b/crates/red_knot_workspace/resources/test/corpus/16_for_break_invalid_in_class.py new file mode 100644 index 0000000000..9f6aa02802 --- /dev/null +++ b/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 diff --git a/crates/red_knot_workspace/resources/test/corpus/16_for_break_invalid_in_func.py b/crates/red_knot_workspace/resources/test/corpus/16_for_break_invalid_in_func.py new file mode 100644 index 0000000000..a4cf76e29f --- /dev/null +++ b/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