Replace parents statement stack with a Nodes abstraction

crates/ruff/src/checkers/ast/deferred.rs

@@ -1,14 +1,16 @@
 use ruff_text_size::TextRange;
 use rustpython_parser::ast::{Expr, Stmt};
 
-use ruff_python_ast::types::RefEquality;
 use ruff_python_semantic::analyze::visibility::{Visibility, VisibleScope};
+use ruff_python_semantic::node::NodeId;
 use ruff_python_semantic::scope::ScopeId;
 
 use crate::checkers::ast::AnnotationContext;
 use crate::docstrings::definition::Definition;
 
-type Context<'a> = (ScopeId, Vec<RefEquality<'a, Stmt>>);
+/// A snapshot of the current scope and statement, which will be restored when visiting any
+/// deferred definitions.
+type Context<'a> = (ScopeId, Option<NodeId>);
 
 /// A collection of AST nodes that are deferred for later analysis.
 /// Used to, e.g., store functions, whose bodies shouldn't be analyzed until all

crates/ruff/src/codes.rs

@@ -1,6 +1,7 @@
-use crate::registry::{Linter, Rule};
 use std::fmt::Formatter;
 
+use crate::registry::{Linter, Rule};
+
 #[derive(PartialEq, Eq, PartialOrd, Ord)]
 pub struct NoqaCode(&'static str, &'static str);
 

crates/ruff/src/rules/flake8_bugbear/rules/unused_loop_control_variable.rs

@@ -168,10 +168,9 @@ pub fn unused_loop_control_variable(
                 let scope = checker.ctx.scope();
                 let binding = scope.bindings_for_name(name).find_map(|index| {
                     let binding = &checker.ctx.bindings[*index];
-                    binding
-                        .source
-                        .as_ref()
-                        .and_then(|source| (source == &RefEquality(stmt)).then_some(binding))
+                    binding.source.and_then(|source| {
+                        (RefEquality(source) == RefEquality(stmt)).then_some(binding)
+                    })
                 });
                 if let Some(binding) = binding {
                     if binding.kind.is_loop_var() {

crates/ruff/src/rules/flake8_simplify/rules/open_file_with_context_handler.rs

@@ -33,7 +33,7 @@ fn match_async_exit_stack(checker: &Checker) -> bool {
     if attr != "enter_async_context" {
         return false;
     }
-    for parent in &checker.ctx.parents {
+    for parent in checker.ctx.parents() {
         if let StmtKind::With { items, .. } = &parent.node {
             for item in items {
                 if let ExprKind::Call { func, .. } = &item.context_expr.node {
@@ -68,7 +68,7 @@ fn match_exit_stack(checker: &Checker) -> bool {
     if attr != "enter_context" {
         return false;
     }
-    for parent in &checker.ctx.parents {
+    for parent in checker.ctx.parents() {
         if let StmtKind::With { items, .. } = &parent.node {
             for item in items {
                 if let ExprKind::Call { func, .. } = &item.context_expr.node {

crates/ruff/src/rules/flake8_type_checking/rules/empty_type_checking_block.rs

@@ -60,11 +60,7 @@ pub fn empty_type_checking_block<'a, 'b>(
 
         // Delete the entire type-checking block.
         if checker.patch(diagnostic.kind.rule()) {
-            let parent = checker
-                .ctx
-                .child_to_parent
-                .get(&RefEquality(stmt))
-                .map(Into::into);
+            let parent = checker.ctx.stmts.parent(stmt);
             let deleted: Vec<&Stmt> = checker.deletions.iter().map(Into::into).collect();
             match delete_stmt(
                 stmt,

crates/ruff/src/rules/pyflakes/rules/unused_variable.rs

@@ -214,11 +214,7 @@ fn remove_unused_variable(
                     ))
                 } else {
                     // If (e.g.) assigning to a constant (`x = 1`), delete the entire statement.
-                    let parent = checker
-                        .ctx
-                        .child_to_parent
-                        .get(&RefEquality(stmt))
-                        .map(Into::into);
+                    let parent = checker.ctx.stmts.parent(stmt);
                     let deleted: Vec<&Stmt> = checker.deletions.iter().map(Into::into).collect();
                     match delete_stmt(
                         stmt,
@@ -259,11 +255,7 @@ fn remove_unused_variable(
                 ))
             } else {
                 // If assigning to a constant (`x = 1`), delete the entire statement.
-                let parent = checker
-                    .ctx
-                    .child_to_parent
-                    .get(&RefEquality(stmt))
-                    .map(Into::into);
+                let parent = checker.ctx.stmts.parent(stmt);
                 let deleted: Vec<&Stmt> = checker.deletions.iter().map(Into::into).collect();
                 match delete_stmt(
                     stmt,
@@ -336,7 +328,7 @@ pub fn unused_variable(checker: &mut Checker, scope: ScopeId) {
                 binding.range,
             );
             if checker.patch(diagnostic.kind.rule()) {
-                if let Some(stmt) = binding.source.as_ref().map(Into::into) {
+                if let Some(stmt) = binding.source {
                     if let Some((kind, fix)) = remove_unused_variable(stmt, binding.range, checker)
                     {
                         if matches!(kind, DeletionKind::Whole) {

crates/ruff/src/rules/pylint/rules/global_statement.rs

@@ -57,9 +57,7 @@ pub fn global_statement(checker: &mut Checker, name: &str) {
         if binding.kind.is_global() {
             let source: &Stmt = binding
                 .source
-                .as_ref()
-                .expect("`global` bindings should always have a `source`")
-                .into();
+                .expect("`global` bindings should always have a `source`");
             let diagnostic = Diagnostic::new(
                 GlobalStatement {
                     name: name.to_string(),

crates/ruff/src/rules/pyupgrade/rules/outdated_version_block.rs

@@ -164,9 +164,9 @@ fn fix_py2_block(
         let defined_by = checker.ctx.current_stmt();
         let defined_in = checker.ctx.current_stmt_parent();
         return match delete_stmt(
-            defined_by.into(),
+            defined_by,
             if block.starter == Tok::If {
-                defined_in.map(Into::into)
+                defined_in
             } else {
                 None
             },
@@ -176,7 +176,7 @@ fn fix_py2_block(
             checker.stylist,
         ) {
             Ok(fix) => {
-                checker.deletions.insert(RefEquality(defined_by.into()));
+                checker.deletions.insert(RefEquality(defined_by));
                 Some(fix)
             }
             Err(err) => {

crates/ruff/src/rules/pyupgrade/rules/super_call_with_parameters.rs

@@ -1,4 +1,4 @@
-use rustpython_parser::ast::{ArgData, Expr, ExprKind, Stmt, StmtKind};
+use rustpython_parser::ast::{ArgData, Expr, ExprKind, StmtKind};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
 use ruff_macros::{derive_message_formats, violation};
@@ -39,14 +39,13 @@ pub fn super_call_with_parameters(checker: &mut Checker, expr: &Expr, func: &Exp
         return;
     }
     let scope = checker.ctx.scope();
-    let parents: Vec<&Stmt> = checker.ctx.parents.iter().map(Into::into).collect();
 
     // Check: are we in a Function scope?
     if !matches!(scope.kind, ScopeKind::Function { .. }) {
         return;
     }
 
-    let mut parents = parents.iter().rev();
+    let mut parents = checker.ctx.parents();
 
     // For a `super` invocation to be unnecessary, the first argument needs to match
     // the enclosing class, and the second argument needs to match the first

crates/ruff/src/rules/pyupgrade/rules/unnecessary_builtin_import.rs

@@ -4,6 +4,7 @@ use rustpython_parser::ast::{Alias, AliasData, Located, Stmt};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::types::RefEquality;
 
 use crate::autofix;
 use crate::checkers::ast::Checker;
@@ -113,16 +114,16 @@ pub fn unnecessary_builtin_import(
             .collect();
         match autofix::actions::remove_unused_imports(
             unused_imports.iter().map(String::as_str),
-            defined_by.into(),
-            defined_in.map(Into::into),
+            defined_by,
+            defined_in,
             &deleted,
             checker.locator,
             checker.indexer,
             checker.stylist,
         ) {
             Ok(fix) => {
                 if fix.is_deletion() || fix.content() == Some("pass") {
-                    checker.deletions.insert(*defined_by);
+                    checker.deletions.insert(RefEquality(defined_by));
                 }
                 diagnostic.set_fix(fix);
             }

crates/ruff/src/rules/pyupgrade/rules/unnecessary_future_import.rs

@@ -4,6 +4,7 @@ use rustpython_parser::ast::{Alias, AliasData, Located, Stmt};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::types::RefEquality;
 
 use crate::autofix;
 use crate::checkers::ast::Checker;
@@ -93,16 +94,16 @@ pub fn unnecessary_future_import(checker: &mut Checker, stmt: &Stmt, names: &[Lo
             .collect();
         match autofix::actions::remove_unused_imports(
             unused_imports.iter().map(String::as_str),
-            defined_by.into(),
-            defined_in.map(Into::into),
+            defined_by,
+            defined_in,
             &deleted,
             checker.locator,
             checker.indexer,
             checker.stylist,
         ) {
             Ok(fix) => {
                 if fix.is_deletion() || fix.content() == Some("pass") {
-                    checker.deletions.insert(*defined_by);
+                    checker.deletions.insert(RefEquality(defined_by));
                 }
                 diagnostic.set_fix(fix);
             }

crates/ruff/src/rules/pyupgrade/rules/useless_metaclass_type.rs

@@ -4,6 +4,7 @@ use rustpython_parser::ast::{Expr, ExprKind, Stmt};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::types::RefEquality;
 
 use crate::autofix::actions;
 use crate::checkers::ast::Checker;
@@ -53,16 +54,16 @@ pub fn useless_metaclass_type(checker: &mut Checker, stmt: &Stmt, value: &Expr,
         let defined_by = checker.ctx.current_stmt();
         let defined_in = checker.ctx.current_stmt_parent();
         match actions::delete_stmt(
-            defined_by.into(),
-            defined_in.map(Into::into),
+            defined_by,
+            defined_in,
             &deleted,
             checker.locator,
             checker.indexer,
             checker.stylist,
         ) {
             Ok(fix) => {
                 if fix.is_deletion() || fix.content() == Some("pass") {
-                    checker.deletions.insert(*defined_by);
+                    checker.deletions.insert(RefEquality(defined_by));
                 }
                 diagnostic.set_fix(fix);
             }

crates/ruff_python_ast/src/lib.rs

@@ -1,5 +1,4 @@
 pub mod all;
-pub mod branch_detection;
 pub mod call_path;
 pub mod cast;
 pub mod comparable;

crates/ruff_python_ast/src/types.rs

@@ -70,3 +70,15 @@ impl<'a> From<&RefEquality<'a, Expr>> for &'a Expr {
         r.0
     }
 }
+
+impl<'a> From<RefEquality<'a, Stmt>> for &'a Stmt {
+    fn from(r: RefEquality<'a, Stmt>) -> Self {
+        r.0
+    }
+}
+
+impl<'a> From<RefEquality<'a, Expr>> for &'a Expr {
+    fn from(r: RefEquality<'a, Expr>) -> Self {
+        r.0
+    }
+}