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

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()])
-    }
-}