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
     }