Remove ScopeStack in favor of child-parent ScopeId pointers (#4138)

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

@@ -1,14 +1,14 @@
-use ruff_python_semantic::scope::ScopeStack;
 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::scope::ScopeId;
 
 use crate::checkers::ast::AnnotationContext;
 use crate::docstrings::definition::Definition;
 
-type Context<'a> = (ScopeStack, Vec<RefEquality<'a, Stmt>>);
+type Context<'a> = (ScopeId, Vec<RefEquality<'a, Stmt>>);
 
 /// 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/rules/pyflakes/rules/return_outside_function.rs

@@ -17,14 +17,12 @@ impl Violation for ReturnOutsideFunction {
 }
 
 pub fn return_outside_function(checker: &mut Checker, stmt: &Stmt) {
-    if let Some(index) = checker.ctx.scope_stack.top() {
-        if matches!(
-            checker.ctx.scopes[index].kind,
-            ScopeKind::Class(_) | ScopeKind::Module
-        ) {
-            checker
-                .diagnostics
-                .push(Diagnostic::new(ReturnOutsideFunction, stmt.range()));
-        }
+    if matches!(
+        checker.ctx.scope().kind,
+        ScopeKind::Class(_) | ScopeKind::Module
+    ) {
+        checker
+            .diagnostics
+            .push(Diagnostic::new(ReturnOutsideFunction, stmt.range()));
     }
 }

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

@@ -2,8 +2,8 @@ use std::string::ToString;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_semantic::binding::Bindings;
-use ruff_python_semantic::scope::{Scope, ScopeKind};
+
+use crate::checkers::ast::Checker;
 
 #[violation]
 pub struct UndefinedLocal {
@@ -18,26 +18,39 @@ impl Violation for UndefinedLocal {
     }
 }
 
-/// F821
-pub fn undefined_local(name: &str, scopes: &[&Scope], bindings: &Bindings) -> Option<Diagnostic> {
-    let current = &scopes.last().expect("No current scope found");
-    if matches!(current.kind, ScopeKind::Function(_)) && !current.defines(name) {
-        for scope in scopes.iter().rev().skip(1) {
-            if matches!(scope.kind, ScopeKind::Function(_) | ScopeKind::Module) {
-                if let Some(binding) = scope.get(name).map(|index| &bindings[*index]) {
-                    if let Some((scope_id, location)) = binding.runtime_usage {
-                        if scope_id == current.id {
-                            return Some(Diagnostic::new(
-                                UndefinedLocal {
-                                    name: name.to_string(),
-                                },
-                                location,
-                            ));
-                        }
-                    }
+/// F823
+pub fn undefined_local(checker: &mut Checker, name: &str) {
+    // If the name hasn't already been defined in the current scope...
+    let current = checker.ctx.scope();
+    if !current.kind.is_function() || current.defines(name) {
+        return;
+    }
+
+    let Some(parent) = current.parent else {
+        return;
+    };
+
+    // For every function and module scope above us...
+    for scope in checker.ctx.scopes.ancestors(parent) {
+        if !(scope.kind.is_function() || scope.kind.is_module()) {
+            continue;
+        }
+
+        // If the name was defined in that scope...
+        if let Some(binding) = scope.get(name).map(|index| &checker.ctx.bindings[*index]) {
+            // And has already been accessed in the current scope...
+            if let Some((scope_id, location)) = binding.runtime_usage {
+                if scope_id == checker.ctx.scope_id {
+                    // Then it's probably an error.
+                    checker.diagnostics.push(Diagnostic::new(
+                        UndefinedLocal {
+                            name: name.to_string(),
+                        },
+                        location,
+                    ));
+                    return;
                 }
             }
         }
     }
-    None
 }

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

@@ -16,7 +16,7 @@ pub fn in_dunder_init(checker: &Checker) -> bool {
     if name != "__init__" {
         return false;
     }
-    let Some(parent) = checker.ctx.parent_scope() else {
+    let Some(parent) = scope.parent.map(|scope_id| &checker.ctx.scopes[scope_id]) else {
         return false;
     };
 

crates/ruff_python_semantic/src/context.rs

@@ -1,23 +1,23 @@
 use std::path::Path;
 
 use nohash_hasher::{BuildNoHashHasher, IntMap};
-use ruff_python_ast::call_path::{collect_call_path, from_unqualified_name, CallPath};
-use ruff_python_ast::helpers::from_relative_import;
-use ruff_python_ast::types::RefEquality;
-use ruff_python_ast::typing::AnnotationKind;
 use rustc_hash::FxHashMap;
 use rustpython_parser::ast::{Expr, Stmt};
 use smallvec::smallvec;
 
-use crate::analyze::visibility::{module_visibility, Modifier, VisibleScope};
+use ruff_python_ast::call_path::{collect_call_path, from_unqualified_name, CallPath};
+use ruff_python_ast::helpers::from_relative_import;
+use ruff_python_ast::types::RefEquality;
+use ruff_python_ast::typing::AnnotationKind;
 use ruff_python_stdlib::path::is_python_stub_file;
 use ruff_python_stdlib::typing::TYPING_EXTENSIONS;
 
+use crate::analyze::visibility::{module_visibility, Modifier, VisibleScope};
 use crate::binding::{
     Binding, BindingId, BindingKind, Bindings, Exceptions, ExecutionContext, FromImportation,
     Importation, SubmoduleImportation,
 };
-use crate::scope::{Scope, ScopeId, ScopeKind, ScopeStack, Scopes};
+use crate::scope::{Scope, ScopeId, ScopeKind, Scopes};
 
 #[allow(clippy::struct_excessive_bools)]
 pub struct Context<'a> {
@@ -35,8 +35,8 @@ pub struct Context<'a> {
         std::collections::HashMap<BindingId, Vec<BindingId>, BuildNoHashHasher<BindingId>>,
     pub exprs: Vec<RefEquality<'a, Expr>>,
     pub scopes: Scopes<'a>,
-    pub scope_stack: ScopeStack,
-    pub dead_scopes: Vec<(ScopeId, ScopeStack)>,
+    pub scope_id: ScopeId,
+    pub dead_scopes: Vec<ScopeId>,
     // Body iteration; used to peek at siblings.
     pub body: &'a [Stmt],
     pub body_index: usize,
@@ -75,7 +75,7 @@ impl<'a> Context<'a> {
             shadowed_bindings: IntMap::default(),
             exprs: Vec::default(),
             scopes: Scopes::default(),
-            scope_stack: ScopeStack::default(),
+            scope_id: ScopeId::global(),
             dead_scopes: Vec::default(),
             body: &[],
             body_index: 0,
@@ -330,19 +330,18 @@ impl<'a> Context<'a> {
             .expect("Attempted to pop without expression");
     }
 
-    pub fn push_scope(&mut self, kind: ScopeKind<'a>) -> ScopeId {
-        let id = self.scopes.push_scope(kind);
-        self.scope_stack.push(id);
-        id
+    /// Push a [`Scope`] with the given [`ScopeKind`] onto the stack.
+    pub fn push_scope(&mut self, kind: ScopeKind<'a>) {
+        let id = self.scopes.push_scope(kind, self.scope_id);
+        self.scope_id = id;
     }
 
+    /// Pop the current [`Scope`] off the stack.
     pub fn pop_scope(&mut self) {
-        self.dead_scopes.push((
-            self.scope_stack
-                .pop()
-                .expect("Attempted to pop without scope"),
-            self.scope_stack.clone(),
-        ));
+        self.dead_scopes.push(self.scope_id);
+        self.scope_id = self.scopes[self.scope_id]
+            .parent
+            .expect("Attempted to pop without scope");
     }
 
     /// Return the current `Stmt`.
@@ -387,31 +386,20 @@ impl<'a> Context<'a> {
 
     /// Returns the current top most scope.
     pub fn scope(&self) -> &Scope<'a> {
-        &self.scopes[self.scope_stack.top().expect("No current scope found")]
-    }
-
-    /// Returns the id of the top-most scope
-    pub fn scope_id(&self) -> ScopeId {
-        self.scope_stack.top().expect("No current scope found")
+        &self.scopes[self.scope_id]
     }
 
     /// Returns a mutable reference to the current top most scope.
     pub fn scope_mut(&mut self) -> &mut Scope<'a> {
-        let top_id = self.scope_stack.top().expect("No current scope found");
-        &mut self.scopes[top_id]
-    }
-
-    pub fn parent_scope(&self) -> Option<&Scope> {
-        self.scope_stack
-            .iter()
-            .nth(1)
-            .map(|index| &self.scopes[*index])
+        &mut self.scopes[self.scope_id]
     }
 
+    /// Returns an iterator over all scopes, starting from the current scope.
     pub fn scopes(&self) -> impl Iterator<Item = &Scope> {
-        self.scope_stack.iter().map(|index| &self.scopes[*index])
+        self.scopes.ancestors(self.scope_id)
     }
 
+    /// Returns `true` if the context is in an exception handler.
     pub const fn in_exception_handler(&self) -> bool {
         self.in_exception_handler
     }

crates/ruff_python_semantic/src/scope.rs

@@ -8,8 +8,9 @@ use crate::binding::{BindingId, StarImportation};
 
 #[derive(Debug)]
 pub struct Scope<'a> {
-    pub id: ScopeId,
     pub kind: ScopeKind<'a>,
+    pub parent: Option<ScopeId>,
+    /// Whether this scope uses the `locals()` builtin.
     pub uses_locals: bool,
     /// A list of star imports in this scope. These represent _module_ imports (e.g., `sys` in
     /// `from sys import *`), rather than individual bindings (e.g., individual members in `sys`).
@@ -22,13 +23,20 @@ pub struct Scope<'a> {
 
 impl<'a> Scope<'a> {
     pub fn global() -> Self {
-        Scope::local(ScopeId::global(), ScopeKind::Module)
+        Scope {
+            kind: ScopeKind::Module,
+            parent: None,
+            uses_locals: false,
+            star_imports: Vec::default(),
+            bindings: FxHashMap::default(),
+            shadowed_bindings: FxHashMap::default(),
+        }
     }
 
-    pub fn local(id: ScopeId, kind: ScopeKind<'a>) -> Self {
+    pub fn local(kind: ScopeKind<'a>, parent: ScopeId) -> Self {
         Scope {
-            id,
             kind,
+            parent: Some(parent),
             uses_locals: false,
             star_imports: Vec::default(),
             bindings: FxHashMap::default(),
@@ -189,11 +197,23 @@ impl<'a> Scopes<'a> {
     }
 
     /// Pushes a new scope and returns its unique id
-    pub fn push_scope(&mut self, kind: ScopeKind<'a>) -> ScopeId {
+    pub fn push_scope(&mut self, kind: ScopeKind<'a>, parent: ScopeId) -> ScopeId {
         let next_id = ScopeId::try_from(self.0.len()).unwrap();
-        self.0.push(Scope::local(next_id, kind));
+        self.0.push(Scope::local(kind, parent));
         next_id
     }
+
+    /// Returns an iterator over all [`ScopeId`] ancestors, starting from the given [`ScopeId`].
+    pub fn ancestor_ids(&self, scope_id: ScopeId) -> impl Iterator<Item = ScopeId> + '_ {
+        std::iter::successors(Some(scope_id), |&scope_id| self[scope_id].parent)
+    }
+
+    /// Returns an iterator over all [`Scope`] ancestors, starting from the given [`ScopeId`].
+    pub fn ancestors(&self, scope_id: ScopeId) -> impl Iterator<Item = &Scope> + '_ {
+        std::iter::successors(Some(&self[scope_id]), |&scope| {
+            scope.parent.map(|scope_id| &self[scope_id])
+        })
+    }
 }
 
 impl Default for Scopes<'_> {
@@ -222,45 +242,3 @@ impl<'a> Deref for Scopes<'a> {
         &self.0
     }
 }
-
-#[derive(Debug, Clone)]
-pub struct ScopeStack(Vec<ScopeId>);
-
-impl ScopeStack {
-    /// Pushes a new scope on the stack
-    pub fn push(&mut self, id: ScopeId) {
-        self.0.push(id);
-    }
-
-    /// Pops the top most scope
-    pub fn pop(&mut self) -> Option<ScopeId> {
-        self.0.pop()
-    }
-
-    /// Returns the id of the top-most
-    pub fn top(&self) -> Option<ScopeId> {
-        self.0.last().copied()
-    }
-
-    /// Returns an iterator from the current scope to the top scope (reverse iterator)
-    pub fn iter(&self) -> std::iter::Rev<std::slice::Iter<ScopeId>> {
-        self.0.iter().rev()
-    }
-
-    pub fn snapshot(&self) -> ScopeStackSnapshot {
-        ScopeStackSnapshot(self.0.len())
-    }
-
-    #[allow(clippy::needless_pass_by_value)]
-    pub fn restore(&mut self, snapshot: ScopeStackSnapshot) {
-        self.0.truncate(snapshot.0);
-    }
-}
-
-pub struct ScopeStackSnapshot(usize);
-
-impl Default for ScopeStack {
-    fn default() -> Self {
-        Self(vec![ScopeId::global()])
-    }
-}