Replace parents statement stack with a Nodes abstraction (#4233)

crates/ruff_python_semantic/src/context.rs

@@ -1,7 +1,6 @@
 use std::path::Path;
 
 use nohash_hasher::{BuildNoHashHasher, IntMap};
-use rustc_hash::FxHashMap;
 use rustpython_parser::ast::{Expr, Stmt};
 use smallvec::smallvec;
 
@@ -17,26 +16,26 @@ use crate::binding::{
     Binding, BindingId, BindingKind, Bindings, Exceptions, ExecutionContext, FromImportation,
     Importation, SubmoduleImportation,
 };
+use crate::node::{NodeId, Nodes};
 use crate::scope::{Scope, ScopeId, ScopeKind, Scopes};
 
 #[allow(clippy::struct_excessive_bools)]
 pub struct Context<'a> {
     pub typing_modules: &'a [String],
     pub module_path: Option<Vec<String>>,
-    // Retain all scopes and parent nodes, along with a stack of indices to track which are active
-    // at various points in time.
-    pub parents: Vec<RefEquality<'a, Stmt>>,
-    pub depths: FxHashMap<RefEquality<'a, Stmt>, usize>,
-    pub child_to_parent: FxHashMap<RefEquality<'a, Stmt>, RefEquality<'a, Stmt>>,
+    // Stack of all visited statements, along with the identifier of the current statement.
+    pub stmts: Nodes<'a>,
+    pub stmt_id: Option<NodeId>,
+    // Stack of all scopes, along with the identifier of the current scope.
+    pub scopes: Scopes<'a>,
+    pub scope_id: ScopeId,
+    pub dead_scopes: Vec<ScopeId>,
     // A stack of all bindings created in any scope, at any point in execution.
     pub bindings: Bindings<'a>,
     // Map from binding index to indexes of bindings that shadow it in other scopes.
     pub shadowed_bindings:
         std::collections::HashMap<BindingId, Vec<BindingId>, BuildNoHashHasher<BindingId>>,
     pub exprs: Vec<RefEquality<'a, Expr>>,
-    pub scopes: Scopes<'a>,
-    pub scope_id: ScopeId,
-    pub dead_scopes: Vec<ScopeId>,
     // Body iteration; used to peek at siblings.
     pub body: &'a [Stmt],
     pub body_index: usize,
@@ -68,15 +67,14 @@ impl<'a> Context<'a> {
         Self {
             typing_modules,
             module_path,
-            parents: Vec::default(),
-            depths: FxHashMap::default(),
-            child_to_parent: FxHashMap::default(),
-            bindings: Bindings::default(),
-            shadowed_bindings: IntMap::default(),
-            exprs: Vec::default(),
+            stmts: Nodes::default(),
+            stmt_id: None,
             scopes: Scopes::default(),
             scope_id: ScopeId::global(),
             dead_scopes: Vec::default(),
+            bindings: Bindings::default(),
+            shadowed_bindings: IntMap::default(),
+            exprs: Vec::default(),
             body: &[],
             body_index: 0,
             visible_scope: VisibleScope {
@@ -254,10 +252,7 @@ impl<'a> Context<'a> {
                                 .take(scope_index)
                                 .all(|scope| scope.get(name).is_none())
                             {
-                                return Some((
-                                    binding.source.as_ref().unwrap().into(),
-                                    format!("{name}.{member}"),
-                                ));
+                                return Some((binding.source.unwrap(), format!("{name}.{member}")));
                             }
                         }
                     }
@@ -273,10 +268,7 @@ impl<'a> Context<'a> {
                                     .take(scope_index)
                                     .all(|scope| scope.get(name).is_none())
                                 {
-                                    return Some((
-                                        binding.source.as_ref().unwrap().into(),
-                                        (*name).to_string(),
-                                    ));
+                                    return Some((binding.source.unwrap(), (*name).to_string()));
                                 }
                             }
                         }
@@ -291,10 +283,7 @@ impl<'a> Context<'a> {
                                 .take(scope_index)
                                 .all(|scope| scope.get(name).is_none())
                             {
-                                return Some((
-                                    binding.source.as_ref().unwrap().into(),
-                                    format!("{name}.{member}"),
-                                ));
+                                return Some((binding.source.unwrap(), format!("{name}.{member}")));
                             }
                         }
                     }
@@ -306,18 +295,15 @@ impl<'a> Context<'a> {
         })
     }
 
-    pub fn push_parent(&mut self, parent: &'a Stmt) {
-        let num_existing = self.parents.len();
-        self.parents.push(RefEquality(parent));
-        self.depths.insert(self.parents[num_existing], num_existing);
-        if num_existing > 0 {
-            self.child_to_parent
-                .insert(self.parents[num_existing], self.parents[num_existing - 1]);
-        }
+    /// Push a [`Stmt`] onto the stack.
+    pub fn push_stmt(&mut self, stmt: &'a Stmt) {
+        self.stmt_id = Some(self.stmts.insert(stmt, self.stmt_id));
     }
 
-    pub fn pop_parent(&mut self) {
-        self.parents.pop().expect("Attempted to pop without parent");
+    /// Pop the current [`Stmt`] off the stack.
+    pub fn pop_stmt(&mut self) {
+        let node_id = self.stmt_id.expect("Attempted to pop without statement");
+        self.stmt_id = self.stmts.parent_id(node_id);
     }
 
     pub fn push_expr(&mut self, expr: &'a Expr) {
@@ -345,13 +331,16 @@ impl<'a> Context<'a> {
     }
 
     /// Return the current `Stmt`.
-    pub fn current_stmt(&self) -> &RefEquality<'a, Stmt> {
-        self.parents.iter().rev().next().expect("No parent found")
+    pub fn current_stmt(&self) -> &'a Stmt {
+        let node_id = self.stmt_id.expect("No current statement");
+        self.stmts[node_id]
     }
 
     /// Return the parent `Stmt` of the current `Stmt`, if any.
-    pub fn current_stmt_parent(&self) -> Option<&RefEquality<'a, Stmt>> {
-        self.parents.iter().rev().nth(1)
+    pub fn current_stmt_parent(&self) -> Option<&'a Stmt> {
+        let node_id = self.stmt_id.expect("No current statement");
+        let parent_id = self.stmts.parent_id(node_id)?;
+        Some(self.stmts[parent_id])
     }
 
     /// Return the parent `Expr` of the current `Expr`.
@@ -399,6 +388,11 @@ impl<'a> Context<'a> {
         self.scopes.ancestors(self.scope_id)
     }
 
+    pub fn parents(&self) -> impl Iterator<Item = &Stmt> + '_ {
+        let node_id = self.stmt_id.expect("No current statement");
+        self.stmts.ancestor_ids(node_id).map(|id| self.stmts[id])
+    }
+
     /// Returns `true` if the context is in an exception handler.
     pub const fn in_exception_handler(&self) -> bool {
         self.in_exception_handler