Remove globals table from Scope (#4686)

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,
-    SubmoduleImportation,
+    Binding, BindingFlags, BindingId, BindingKind, Bindings, Exceptions, FromImportation,
+    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.
-    globals: FxHashMap<&'a str, TextRange>,
+    /// Index into the globals arena, if the scope contains any globally-declared symbols.
+    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.
-    pub fn add_global(&mut self, name: &'a str, range: TextRange) {
-        self.globals.insert(name, range);
+    /// Set the globals pointer for this scope.
+    pub fn set_globals_id(&mut self, globals: GlobalsId) {
+        self.globals_id = Some(globals);
     }
 
-    /// Returns the range of the global name with the given name.
-    pub fn get_global(&self, name: &str) -> Option<TextRange> {
-        self.globals.get(name).copied()
+    /// Returns the globals pointer for this scope.
+    pub fn globals_id(&self) -> Option<GlobalsId> {
+        self.globals_id
     }
 }