Skip to content

Commit

Permalink
Pop assignments even outside of debug builds
Browse files Browse the repository at this point in the history
  • Loading branch information
@rtpg
rtpg committed Oct 15, 2024
1 parent 47a3625 commit 65537da
Showing 1 changed file with 39 additions and 27 deletions.
66 changes: 39 additions & 27 deletions crates/red_knot_python_semantic/src/semantic_index/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@ pub(super) struct SemanticIndexBuilder<'db> {
file: File,
module: &'db ParsedModule,
scope_stack: Vec<FileScopeId>,
/// The assignments we're currently visiting.
/// The assignments we're currently visiting, with
/// the most recent visit at the end of the Vec
current_assignments: Vec<CurrentAssignment<'db>>,
/// The match case we're currently visiting.
current_match_case: Option<CurrentMatchCase<'db>>,
Expand Down Expand Up @@ -238,6 +239,19 @@ impl<'db> SemanticIndexBuilder<'db> {
expression
}

fn push_assignment(&mut self, assignment: CurrentAssignment<'db>) {
self.current_assignments.push(assignment);
}

fn pop_assignment(&mut self) {
let popped_assignment = self.current_assignments.pop();
debug_assert!(popped_assignment.is_some());
}

fn current_assignment(&self) -> Option<&CurrentAssignment<'db>> {
self.current_assignments.last()
}

fn add_pattern_constraint(
&mut self,
subject: &ast::Expr,
Expand Down Expand Up @@ -359,13 +373,12 @@ impl<'db> SemanticIndexBuilder<'db> {
self.visit_expr(&generator.iter);
self.push_scope(scope);

self.current_assignments
.push(CurrentAssignment::Comprehension {
node: generator,
first: true,
});
self.push_assignment(CurrentAssignment::Comprehension {
node: generator,
first: true,
});
self.visit_expr(&generator.target);
debug_assert!(self.current_assignments.pop().is_some());
self.pop_assignment();

for expr in &generator.ifs {
self.visit_expr(expr);
Expand All @@ -375,13 +388,12 @@ impl<'db> SemanticIndexBuilder<'db> {
self.add_standalone_expression(&generator.iter);
self.visit_expr(&generator.iter);

self.current_assignments
.push(CurrentAssignment::Comprehension {
node: generator,
first: false,
});
self.push_assignment(CurrentAssignment::Comprehension {
node: generator,
first: false,
});
self.visit_expr(&generator.target);
debug_assert!(self.current_assignments.pop().is_some());
self.pop_assignment();

for expr in &generator.ifs {
self.visit_expr(expr);
Expand Down Expand Up @@ -568,21 +580,21 @@ where
debug_assert!(self.current_assignments.is_empty());
self.visit_expr(&node.value);
self.add_standalone_expression(&node.value);
self.current_assignments.push(node.into());
self.push_assignment(node.into());
for target in &node.targets {
self.visit_expr(target);
}
debug_assert!(self.current_assignments.pop().is_some());
self.pop_assignment();
}
ast::Stmt::AnnAssign(node) => {
debug_assert!(self.current_assignments.is_empty());
self.visit_expr(&node.annotation);
if let Some(value) = &node.value {
self.visit_expr(value);
}
self.current_assignments.push(node.into());
self.push_assignment(node.into());
self.visit_expr(&node.target);
debug_assert!(self.current_assignments.pop().is_some());
self.pop_assignment();
}
ast::Stmt::AugAssign(
aug_assign @ ast::StmtAugAssign {
Expand All @@ -594,9 +606,9 @@ where
) => {
debug_assert!(self.current_assignments.is_empty());
self.visit_expr(value);
self.current_assignments.push(aug_assign.into());
self.push_assignment(aug_assign.into());
self.visit_expr(target);
debug_assert!(self.current_assignments.pop().is_some());
self.pop_assignment();
}
ast::Stmt::If(node) => {
self.visit_expr(&node.test);
Expand Down Expand Up @@ -664,9 +676,9 @@ where
self.visit_expr(&item.context_expr);
if let Some(optional_vars) = item.optional_vars.as_deref() {
self.add_standalone_expression(&item.context_expr);
self.current_assignments.push(item.into());
self.push_assignment(item.into());
self.visit_expr(optional_vars);
self.current_assignments.pop();
self.pop_assignment();
}
}
self.visit_body(body);
Expand All @@ -692,9 +704,9 @@ where
let saved_break_states = std::mem::take(&mut self.loop_break_states);

debug_assert!(self.current_assignments.is_empty());
self.current_assignments.push(for_stmt.into());
self.push_assignment(for_stmt.into());
self.visit_expr(target);
debug_assert!(self.current_assignments.pop().is_some());
self.pop_assignment();

// TODO: Definitions created by loop variables
// (and definitions created inside the body)
Expand Down Expand Up @@ -804,7 +816,7 @@ where

match expr {
ast::Expr::Name(name_node @ ast::ExprName { id, ctx, .. }) => {
let (is_use, is_definition) = match (ctx, self.current_assignments.last()) {
let (is_use, is_definition) = match (ctx, self.current_assignment()) {
(ast::ExprContext::Store, Some(CurrentAssignment::AugAssign(_))) => {
// For augmented assignment, the target expression is also used.
(true, true)
Expand All @@ -817,7 +829,7 @@ where
let symbol = self.add_symbol(id.clone());

if is_definition {
match self.current_assignments.last().copied() {
match self.current_assignment().copied() {
Some(CurrentAssignment::Assign(assignment)) => {
self.add_definition(
symbol,
Expand Down Expand Up @@ -884,9 +896,9 @@ where
ast::Expr::Named(node) => {
// TODO walrus in comprehensions is implicitly nonlocal
self.visit_expr(&node.value);
self.current_assignments.push(node.into());
self.push_assignment(node.into());
self.visit_expr(&node.target);
debug_assert!(self.current_assignments.pop().is_some());
self.pop_assignment();
}
ast::Expr::Lambda(lambda) => {
if let Some(parameters) = &lambda.parameters {
Expand Down

0 comments on commit 65537da

Please sign in to comment.