[red-knot] unify LoopState and saved_break_states (#16406)

We currently keep two separate pieces of state regarding the current
loop on `SemanticIndexBuilder`. One is an enum simply reflecting whether
we are currently inside a loop, and the other is the saved flow states
for `break` statements found in the current loop.

For adding loopy control flow, I'll need to add some additional loop
state (`continue` states, for example). Prepare for this by
consolidating our existing loop state into a single struct and
simplifying the API for pushing and popping a loop.

This is purely a refactor, so tests are not changed.

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
This commit is contained in:
Carl Meyer 2025-02-26 14:31:13 -08:00 committed by GitHub
parent 671494a620
commit fb778ee38d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -41,22 +41,22 @@ use crate::Db;
mod except_handlers; mod except_handlers;
/// Are we in a state where a `break` statement is allowed? #[derive(Clone, Debug, Default)]
#[derive(Clone, Copy, Debug)] struct Loop {
enum LoopState { /// Flow states at each `break` in the current loop.
InLoop, break_states: Vec<FlowSnapshot>,
NotInLoop,
} }
impl LoopState { impl Loop {
fn is_inside(self) -> bool { fn push_break(&mut self, state: FlowSnapshot) {
matches!(self, LoopState::InLoop) self.break_states.push(state);
} }
} }
struct ScopeInfo { struct ScopeInfo {
file_scope_id: FileScopeId, file_scope_id: FileScopeId,
loop_state: LoopState, /// Current loop state; None if we are not currently visiting a loop
current_loop: Option<Loop>,
} }
pub(super) struct SemanticIndexBuilder<'db> { pub(super) struct SemanticIndexBuilder<'db> {
@ -73,8 +73,6 @@ pub(super) struct SemanticIndexBuilder<'db> {
/// The name of the first function parameter of the innermost function that we're currently visiting. /// The name of the first function parameter of the innermost function that we're currently visiting.
current_first_parameter_name: Option<&'db str>, current_first_parameter_name: Option<&'db str>,
/// Flow states at each `break` in the current loop.
loop_break_states: Vec<FlowSnapshot>,
/// Per-scope contexts regarding nested `try`/`except` statements /// Per-scope contexts regarding nested `try`/`except` statements
try_node_context_stack_manager: TryNodeContextStackManager, try_node_context_stack_manager: TryNodeContextStackManager,
@ -106,7 +104,6 @@ impl<'db> SemanticIndexBuilder<'db> {
current_assignments: vec![], current_assignments: vec![],
current_match_case: None, current_match_case: None,
current_first_parameter_name: None, current_first_parameter_name: None,
loop_break_states: vec![],
try_node_context_stack_manager: TryNodeContextStackManager::default(), try_node_context_stack_manager: TryNodeContextStackManager::default(),
has_future_annotations: false, has_future_annotations: false,
@ -134,19 +131,20 @@ impl<'db> SemanticIndexBuilder<'db> {
builder builder
} }
fn current_scope(&self) -> FileScopeId { fn current_scope_info(&self) -> &ScopeInfo {
*self
.scope_stack
.last()
.map(|ScopeInfo { file_scope_id, .. }| file_scope_id)
.expect("SemanticIndexBuilder should have created a root scope")
}
fn loop_state(&self) -> LoopState {
self.scope_stack self.scope_stack
.last() .last()
.expect("SemanticIndexBuilder should have created a root scope") .expect("SemanticIndexBuilder should have created a root scope")
.loop_state }
fn current_scope_info_mut(&mut self) -> &mut ScopeInfo {
self.scope_stack
.last_mut()
.expect("SemanticIndexBuilder should have created a root scope")
}
fn current_scope(&self) -> FileScopeId {
self.current_scope_info().file_scope_id
} }
/// Returns the scope ID of the surrounding class body scope if the current scope /// Returns the scope ID of the surrounding class body scope if the current scope
@ -167,11 +165,21 @@ impl<'db> SemanticIndexBuilder<'db> {
} }
} }
fn set_inside_loop(&mut self, state: LoopState) { /// Push a new loop, returning the outer loop, if any.
self.scope_stack fn push_loop(&mut self) -> Option<Loop> {
.last_mut() self.current_scope_info_mut()
.expect("Always to have a root scope") .current_loop
.loop_state = state; .replace(Loop::default())
}
/// Pop a loop, replacing with the previous saved outer loop, if any.
fn pop_loop(&mut self, outer_loop: Option<Loop>) -> Loop {
std::mem::replace(&mut self.current_scope_info_mut().current_loop, outer_loop)
.expect("pop_loop() should not be called without a prior push_loop()")
}
fn current_loop_mut(&mut self) -> Option<&mut Loop> {
self.current_scope_info_mut().current_loop.as_mut()
} }
fn push_scope(&mut self, node: NodeWithScopeRef) { fn push_scope(&mut self, node: NodeWithScopeRef) {
@ -204,7 +212,7 @@ impl<'db> SemanticIndexBuilder<'db> {
self.scope_stack.push(ScopeInfo { self.scope_stack.push(ScopeInfo {
file_scope_id, file_scope_id,
loop_state: LoopState::NotInLoop, current_loop: None,
}); });
} }
@ -1208,15 +1216,9 @@ where
.current_visibility_constraints_mut() .current_visibility_constraints_mut()
.add_atom(later_predicate_id); .add_atom(later_predicate_id);
// Save aside any break states from an outer loop let outer_loop = self.push_loop();
let saved_break_states = std::mem::take(&mut self.loop_break_states);
// 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.visit_body(body);
self.set_inside_loop(outer_loop_state); let this_loop = self.pop_loop(outer_loop);
// If the body is executed, we know that we've evaluated the condition at least // If the body is executed, we know that we've evaluated the condition at least
// once, and that the first evaluation was True. We might not have evaluated the // once, and that the first evaluation was True. We might not have evaluated the
@ -1225,11 +1227,6 @@ where
let body_vis_constraint_id = first_vis_constraint_id; let body_vis_constraint_id = first_vis_constraint_id;
self.record_visibility_constraint_id(body_vis_constraint_id); self.record_visibility_constraint_id(body_vis_constraint_id);
// Get the break states from the body of this loop, and restore the saved outer
// ones.
let break_states =
std::mem::replace(&mut self.loop_break_states, saved_break_states);
// We execute the `else` once the condition evaluates to false. This could happen // We execute the `else` once the condition evaluates to false. This could happen
// without ever executing the body, if the condition is false the first time it's // without ever executing the body, if the condition is false the first time it's
// tested. So the starting flow state of the `else` clause is the union of: // tested. So the starting flow state of the `else` clause is the union of:
@ -1250,7 +1247,7 @@ where
// Breaking out of a while loop bypasses the `else` clause, so merge in the break // Breaking out of a while loop bypasses the `else` clause, so merge in the break
// states after visiting `else`. // states after visiting `else`.
for break_state in break_states { for break_state in this_loop.break_states {
let snapshot = self.flow_snapshot(); let snapshot = self.flow_snapshot();
self.flow_restore(break_state); self.flow_restore(break_state);
self.record_visibility_constraint_id(body_vis_constraint_id); self.record_visibility_constraint_id(body_vis_constraint_id);
@ -1298,7 +1295,6 @@ where
self.record_ambiguous_visibility(); self.record_ambiguous_visibility();
let pre_loop = self.flow_snapshot(); let pre_loop = self.flow_snapshot();
let saved_break_states = std::mem::take(&mut self.loop_break_states);
let current_assignment = match &**target { let current_assignment = match &**target {
ast::Expr::List(_) | ast::Expr::Tuple(_) => Some(CurrentAssignment::For { ast::Expr::List(_) | ast::Expr::Tuple(_) => Some(CurrentAssignment::For {
@ -1346,16 +1342,9 @@ where
self.pop_assignment(); self.pop_assignment();
} }
// TODO: Definitions created by loop variables let outer_loop = self.push_loop();
// (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.visit_body(body);
self.set_inside_loop(outer_loop_state); let this_loop = self.pop_loop(outer_loop);
let break_states =
std::mem::replace(&mut self.loop_break_states, saved_break_states);
// We may execute the `else` clause without ever executing the body, so merge in // We may execute the `else` clause without ever executing the body, so merge in
// the pre-loop state before visiting `else`. // the pre-loop state before visiting `else`.
@ -1364,7 +1353,7 @@ where
// Breaking out of a `for` loop bypasses the `else` clause, so merge in the break // Breaking out of a `for` loop bypasses the `else` clause, so merge in the break
// states after visiting `else`. // states after visiting `else`.
for break_state in break_states { for break_state in this_loop.break_states {
self.flow_merge(break_state); self.flow_merge(break_state);
} }
} }
@ -1547,8 +1536,9 @@ where
} }
ast::Stmt::Break(_) => { ast::Stmt::Break(_) => {
if self.loop_state().is_inside() { let snapshot = self.flow_snapshot();
self.loop_break_states.push(self.flow_snapshot()); if let Some(current_loop) = self.current_loop_mut() {
current_loop.push_break(snapshot);
} }
// Everything in the current block after a terminal statement is unreachable. // Everything in the current block after a terminal statement is unreachable.
self.mark_unreachable(); self.mark_unreachable();