Remove globals table from Scope (#4686)

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

@@ -29,6 +29,7 @@ use ruff_python_semantic::binding::{
 };
 use ruff_python_semantic::context::ExecutionContext;
 use ruff_python_semantic::definition::{ContextualizedDefinition, Module, ModuleKind};
+use ruff_python_semantic::globals::Globals;
 use ruff_python_semantic::model::{ResolvedReference, SemanticModel, SemanticModelFlags};
 use ruff_python_semantic::node::NodeId;
 use ruff_python_semantic::scope::{Scope, ScopeId, ScopeKind};
@@ -1905,32 +1906,9 @@ where
 
                 self.deferred.functions.push(self.semantic_model.snapshot());
 
-                // If any global bindings don't already exist in the global scope, add them.
+                // Extract any global bindings from the function body.
-                let globals = helpers::extract_globals(body);
+                if let Some(globals) = Globals::from_body(body) {
-                for (name, range) in globals {
+                    self.semantic_model.set_globals(globals);
-                    if self
-                        .semantic_model
-                        .global_scope()
-                        .get(name)
-                        .map_or(true, |binding_id| {
-                            self.semantic_model.bindings[binding_id]
-                                .kind
-                                .is_annotation()
-                        })
-                    {
-                        let id = self.semantic_model.bindings.push(Binding {
-                            kind: BindingKind::Assignment,
-                            range,
-                            references: Vec::new(),
-                            source: self.semantic_model.stmt_id,
-                            context: self.semantic_model.execution_context(),
-                            exceptions: self.semantic_model.exceptions(),
-                            flags: BindingFlags::empty(),
-                        });
-                        self.semantic_model.global_scope_mut().add(name, id);
-                    }
-
-                    self.semantic_model.scope_mut().add_global(name, range);
                 }
             }
             Stmt::ClassDef(
@@ -1962,30 +1940,9 @@ where
 
                 self.semantic_model.push_scope(ScopeKind::Class(class_def));
 
-                // If any global bindings don't already exist in the global scope, add them.
+                // Extract any global bindings from the class body.
-                for (name, range) in helpers::extract_globals(body) {
+                if let Some(globals) = Globals::from_body(body) {
-                    if self
+                    self.semantic_model.set_globals(globals);
-                        .semantic_model
-                        .global_scope()
-                        .get(name)
-                        .map_or(true, |binding_id| {
-                            self.semantic_model.bindings[binding_id]
-                                .kind
-                                .is_annotation()
-                        })
-                    {
-                        let id = self.semantic_model.bindings.push(Binding {
-                            kind: BindingKind::Assignment,
-                            range,
-                            references: Vec::new(),
-                            source: self.semantic_model.stmt_id,
-                            context: self.semantic_model.execution_context(),
-                            exceptions: self.semantic_model.exceptions(),
-                            flags: BindingFlags::empty(),
-                        });
-                        self.semantic_model.global_scope_mut().add(name, id);
-                    }
-                    self.semantic_model.scope_mut().add_global(name, range);
                 }
 
                 self.visit_body(body);
@@ -4256,7 +4213,7 @@ impl<'a> Checker<'a> {
     }
 
     /// Visit an [`Expr`], and treat it as a type definition.
-    pub(crate) fn visit_type_definition(&mut self, expr: &'a Expr) {
+    fn visit_type_definition(&mut self, expr: &'a Expr) {
         let snapshot = self.semantic_model.flags;
         self.semantic_model.flags |= SemanticModelFlags::TYPE_DEFINITION;
         self.visit_expr(expr);
@@ -4264,7 +4221,7 @@ impl<'a> Checker<'a> {
     }
 
     /// Visit an [`Expr`], and treat it as _not_ a type definition.
-    pub(crate) fn visit_non_type_definition(&mut self, expr: &'a Expr) {
+    fn visit_non_type_definition(&mut self, expr: &'a Expr) {
         let snapshot = self.semantic_model.flags;
         self.semantic_model.flags -= SemanticModelFlags::TYPE_DEFINITION;
         self.visit_expr(expr);
@@ -4274,7 +4231,7 @@ impl<'a> Checker<'a> {
     /// Visit an [`Expr`], and treat it as a boolean test. This is useful for detecting whether an
     /// expressions return value is significant, or whether the calling context only relies on
     /// its truthiness.
-    pub(crate) fn visit_boolean_test(&mut self, expr: &'a Expr) {
+    fn visit_boolean_test(&mut self, expr: &'a Expr) {
         let snapshot = self.semantic_model.flags;
         self.semantic_model.flags |= SemanticModelFlags::BOOLEAN_TEST;
         self.visit_expr(expr);

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

@@ -54,8 +54,7 @@ impl Violation for LoadBeforeGlobalDeclaration {
 }
 /// PLE0118
 pub(crate) fn load_before_global_declaration(checker: &mut Checker, name: &str, expr: &Expr) {
-    let scope = checker.semantic_model().scope();
+    if let Some(stmt) = checker.semantic_model().global(name) {
-    if let Some(stmt) = scope.get_global(name) {
         if expr.start() < stmt.start() {
             #[allow(deprecated)]
             let location = checker.locator.compute_source_location(stmt.start());

crates/ruff_python_ast/src/helpers.rs

@@ -959,34 +959,6 @@ where
     }
 }
 
-#[derive(Default)]
-struct GlobalStatementVisitor<'a> {
-    globals: FxHashMap<&'a str, TextRange>,
-}
-
-impl<'a> StatementVisitor<'a> for GlobalStatementVisitor<'a> {
-    fn visit_stmt(&mut self, stmt: &'a Stmt) {
-        match stmt {
-            Stmt::Global(ast::StmtGlobal { names, range }) => {
-                for name in names {
-                    self.globals.insert(name.as_str(), *range);
-                }
-            }
-            Stmt::FunctionDef(_) | Stmt::AsyncFunctionDef(_) | Stmt::ClassDef(_) => {
-                // Don't recurse.
-            }
-            _ => walk_stmt(self, stmt),
-        }
-    }
-}
-
-/// Extract a map from global name to its last-defining [`Stmt`].
-pub fn extract_globals(body: &[Stmt]) -> FxHashMap<&str, TextRange> {
-    let mut visitor = GlobalStatementVisitor::default();
-    visitor.visit_body(body);
-    visitor.globals
-}
-
 /// Return `true` if a [`Ranged`] has leading content.
 pub fn has_leading_content<T>(located: &T, locator: &Locator) -> bool
 where

crates/ruff_python_semantic/src/globals.rs

@@ -0,0 +1,89 @@
+//! When building a semantic model, we often need to know which names in a given scope are declared
+//! as `global`. This module provides data structures for storing and querying the set of `global`
+//! names in a given scope.
+
+use std::ops::Index;
+
+use ruff_text_size::TextRange;
+use rustc_hash::FxHashMap;
+use rustpython_parser::ast;
+use rustpython_parser::ast::Stmt;
+
+use ruff_index::{newtype_index, IndexVec};
+use ruff_python_ast::statement_visitor::{walk_stmt, StatementVisitor};
+
+/// Id uniquely identifying the set of global names for a given scope.
+#[newtype_index]
+pub struct GlobalsId;
+
+#[derive(Debug, Default)]
+pub(crate) struct GlobalsArena<'a>(IndexVec<GlobalsId, Globals<'a>>);
+
+impl<'a> GlobalsArena<'a> {
+    /// Inserts a new set of global names into the global names arena and returns its unique id.
+    pub(crate) fn push(&mut self, globals: Globals<'a>) -> GlobalsId {
+        self.0.push(globals)
+    }
+}
+
+impl<'a> Index<GlobalsId> for GlobalsArena<'a> {
+    type Output = Globals<'a>;
+
+    #[inline]
+    fn index(&self, index: GlobalsId) -> &Self::Output {
+        &self.0[index]
+    }
+}
+
+/// The set of global names for a given scope, represented as a map from the name of the global to
+/// the range of the declaration in the source code.
+#[derive(Debug)]
+pub struct Globals<'a>(FxHashMap<&'a str, TextRange>);
+
+impl<'a> Globals<'a> {
+    /// Extracts the set of global names from a given scope, or return `None` if the scope does not
+    /// contain any `global` declarations.
+    pub fn from_body(body: &'a [Stmt]) -> Option<Self> {
+        let mut builder = GlobalsVisitor::new();
+        builder.visit_body(body);
+        builder.finish()
+    }
+
+    pub fn get(&self, name: &str) -> Option<&TextRange> {
+        self.0.get(name)
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (&&'a str, &TextRange)> + '_ {
+        self.0.iter()
+    }
+}
+
+/// Extracts the set of global names from a given scope.
+#[derive(Debug)]
+struct GlobalsVisitor<'a>(FxHashMap<&'a str, TextRange>);
+
+impl<'a> GlobalsVisitor<'a> {
+    fn new() -> Self {
+        Self(FxHashMap::default())
+    }
+
+    fn finish(self) -> Option<Globals<'a>> {
+        (!self.0.is_empty()).then_some(Globals(self.0))
+    }
+}
+
+impl<'a> StatementVisitor<'a> for GlobalsVisitor<'a> {
+    fn visit_stmt(&mut self, stmt: &'a Stmt) {
+        match stmt {
+            Stmt::Global(ast::StmtGlobal { names, range }) => {
+                for name in names {
+                    self.0.insert(name.as_str(), *range);
+                }
+            }
+            Stmt::FunctionDef(_) | Stmt::AsyncFunctionDef(_) | Stmt::ClassDef(_) => {
+                // Don't recurse.
+            }
+            _ => walk_stmt(self, stmt),
+        }
+    }
+}

crates/ruff_python_semantic/src/lib.rs

@@ -2,6 +2,7 @@ pub mod analyze;
 pub mod binding;
 pub mod context;
 pub mod definition;
+pub mod globals;
 pub mod model;
 pub mod node;
 pub mod reference;

crates/ruff_python_semantic/src/model.rs

@@ -13,11 +13,12 @@ use ruff_python_stdlib::path::is_python_stub_file;
 use ruff_python_stdlib::typing::TYPING_EXTENSIONS;
 
 use crate::binding::{
-    Binding, BindingId, BindingKind, Bindings, Exceptions, FromImportation, Importation,
+    Binding, BindingFlags, BindingId, BindingKind, Bindings, Exceptions, FromImportation,
-    SubmoduleImportation,
+    Importation, SubmoduleImportation,
 };
 use crate::context::ExecutionContext;
 use crate::definition::{Definition, DefinitionId, Definitions, Member, Module};
+use crate::globals::{Globals, GlobalsArena};
 use crate::node::{NodeId, Nodes};
 use crate::reference::References;
 use crate::scope::{Scope, ScopeId, ScopeKind, Scopes};
@@ -43,6 +44,8 @@ pub struct SemanticModel<'a> {
     pub bindings: Bindings<'a>,
     // Stack of all references created in any scope, at any point in execution.
     pub references: References,
+    // Arena of global bindings.
+    globals: GlobalsArena<'a>,
     // Map from binding index to indexes of bindings that shadow it in other scopes.
     pub shadowed_bindings: HashMap<BindingId, Vec<BindingId>, BuildNoHashHasher<BindingId>>,
     // Body iteration; used to peek at siblings.
@@ -68,6 +71,7 @@ impl<'a> SemanticModel<'a> {
             definition_id: DefinitionId::module(),
             bindings: Bindings::default(),
             references: References::default(),
+            globals: GlobalsArena::default(),
             shadowed_bindings: IntMap::default(),
             body: &[],
             body_index: 0,
@@ -546,6 +550,35 @@ impl<'a> SemanticModel<'a> {
         self.stmts.ancestor_ids(node_id).map(|id| self.stmts[id])
     }
 
+    /// Set the [`Globals`] for the current [`Scope`].
+    pub fn set_globals(&mut self, globals: Globals<'a>) {
+        // If any global bindings don't already exist in the global scope, add them.
+        for (name, range) in globals.iter() {
+            if self.global_scope().get(name).map_or(true, |binding_id| {
+                self.bindings[binding_id].kind.is_annotation()
+            }) {
+                let id = self.bindings.push(Binding {
+                    kind: BindingKind::Assignment,
+                    range: *range,
+                    references: Vec::new(),
+                    source: self.stmt_id,
+                    context: self.execution_context(),
+                    exceptions: self.exceptions(),
+                    flags: BindingFlags::empty(),
+                });
+                self.global_scope_mut().add(name, id);
+            }
+        }
+
+        self.scopes[self.scope_id].set_globals_id(self.globals.push(globals));
+    }
+
+    /// Return the [`TextRange`] at which a name is declared as global in the current [`Scope`].
+    pub fn global(&self, name: &str) -> Option<TextRange> {
+        let global_id = self.scopes[self.scope_id].globals_id()?;
+        self.globals[global_id].get(name).copied()
+    }
+
     /// Return `true` if the given [`ScopeId`] matches that of the current scope.
     pub fn is_current_scope(&self, scope_id: ScopeId) -> bool {
         self.scope_id == scope_id

crates/ruff_python_semantic/src/scope.rs

@@ -1,12 +1,12 @@
 use std::ops::{Deref, DerefMut};
 
-use ruff_text_size::TextRange;
 use rustc_hash::FxHashMap;
 use rustpython_parser::ast;
 
 use ruff_index::{newtype_index, Idx, IndexSlice, IndexVec};
 
 use crate::binding::{BindingId, StarImportation};
+use crate::globals::GlobalsId;
 
 #[derive(Debug)]
 pub struct Scope<'a> {
@@ -21,8 +21,8 @@ pub struct Scope<'a> {
     bindings: FxHashMap<&'a str, BindingId>,
     /// A map from bound name to binding index, for bindings that were shadowed later in the scope.
     shadowed_bindings: FxHashMap<&'a str, Vec<BindingId>>,
-    /// A map from global name to the range that declares it.
+    /// Index into the globals arena, if the scope contains any globally-declared symbols.
-    globals: FxHashMap<&'a str, TextRange>,
+    globals_id: Option<GlobalsId>,
 }
 
 impl<'a> Scope<'a> {
@@ -34,7 +34,7 @@ impl<'a> Scope<'a> {
             star_imports: Vec::default(),
             bindings: FxHashMap::default(),
             shadowed_bindings: FxHashMap::default(),
-            globals: FxHashMap::default(),
+            globals_id: None,
         }
     }
 
@@ -46,7 +46,7 @@ impl<'a> Scope<'a> {
             star_imports: Vec::default(),
             bindings: FxHashMap::default(),
             shadowed_bindings: FxHashMap::default(),
-            globals: FxHashMap::default(),
+            globals_id: None,
         }
     }
 
@@ -110,14 +110,14 @@ impl<'a> Scope<'a> {
         self.star_imports.iter()
     }
 
-    /// Add a global name to this scope.
+    /// Set the globals pointer for this scope.
-    pub fn add_global(&mut self, name: &'a str, range: TextRange) {
+    pub fn set_globals_id(&mut self, globals: GlobalsId) {
-        self.globals.insert(name, range);
+        self.globals_id = Some(globals);
     }
 
-    /// Returns the range of the global name with the given name.
+    /// Returns the globals pointer for this scope.
-    pub fn get_global(&self, name: &str) -> Option<TextRange> {
+    pub fn globals_id(&self) -> Option<GlobalsId> {
-        self.globals.get(name).copied()
+        self.globals_id
     }
 }