Track all read references in semantic model (#4610)

crates/ruff_python_semantic/src/binding.rs

@@ -1,14 +1,15 @@
 use std::ops::{Deref, DerefMut};
 
-use crate::model::SemanticModel;
 use bitflags::bitflags;
+use ruff_text_size::TextRange;
+
 use ruff_index::{newtype_index, IndexSlice, IndexVec};
 use ruff_python_ast::helpers;
 use ruff_python_ast::source_code::Locator;
-use ruff_text_size::TextRange;
 
+use crate::model::SemanticModel;
 use crate::node::NodeId;
-use crate::scope::ScopeId;
+use crate::reference::ReferenceId;
 
 #[derive(Debug, Clone)]
 pub struct Binding<'a> {
@@ -18,33 +19,20 @@ pub struct Binding<'a> {
     pub context: ExecutionContext,
     /// The statement in which the [`Binding`] was defined.
     pub source: Option<NodeId>,
-    /// Tuple of (scope index, range) indicating the scope and range at which
-    /// the binding was last used in a runtime context.
-    pub runtime_usage: Option<(ScopeId, TextRange)>,
-    /// Tuple of (scope index, range) indicating the scope and range at which
-    /// the binding was last used in a typing-time context.
-    pub typing_usage: Option<(ScopeId, TextRange)>,
-    /// Tuple of (scope index, range) indicating the scope and range at which
-    /// the binding was last used in a synthetic context. This is used for
-    /// (e.g.) `__future__` imports, explicit re-exports, and other bindings
-    /// that should be considered used even if they're never referenced.
-    pub synthetic_usage: Option<(ScopeId, TextRange)>,
+    /// The references to the binding.
+    pub references: Vec<ReferenceId>,
     /// The exceptions that were handled when the binding was defined.
     pub exceptions: Exceptions,
 }
 
 impl<'a> Binding<'a> {
-    pub fn mark_used(&mut self, scope: ScopeId, range: TextRange, context: ExecutionContext) {
-        match context {
-            ExecutionContext::Runtime => self.runtime_usage = Some((scope, range)),
-            ExecutionContext::Typing => self.typing_usage = Some((scope, range)),
-        }
+    pub fn is_used(&self) -> bool {
+        !self.references.is_empty()
     }
 
-    pub const fn used(&self) -> bool {
-        self.runtime_usage.is_some()
-            || self.synthetic_usage.is_some()
-            || self.typing_usage.is_some()
+    /// Returns an iterator over all references for the current [`Binding`].
+    pub fn references(&self) -> impl Iterator<Item = ReferenceId> + '_ {
+        self.references.iter().copied()
     }
 
     pub const fn is_definition(&self) -> bool {
@@ -266,7 +254,7 @@ bitflags! {
     }
 }
 
-#[derive(Copy, Debug, Clone)]
+#[derive(Copy, Debug, Clone, is_macro::Is)]
 pub enum ExecutionContext {
     Runtime,
     Typing,

crates/ruff_python_semantic/src/lib.rs

@@ -3,4 +3,5 @@ pub mod binding;
 pub mod definition;
 pub mod model;
 pub mod node;
+pub mod reference;
 pub mod scope;

crates/ruff_python_semantic/src/model.rs

@@ -18,6 +18,7 @@ use crate::binding::{
 };
 use crate::definition::{Definition, DefinitionId, Definitions, Member, Module};
 use crate::node::{NodeId, Nodes};
+use crate::reference::{ReferenceContext, References};
 use crate::scope::{Scope, ScopeId, ScopeKind, Scopes};
 
 /// A semantic model for a Python module, to enable querying the module's semantic information.
@@ -39,6 +40,8 @@ pub struct SemanticModel<'a> {
     pub definition_id: DefinitionId,
     // A stack of all bindings created in any scope, at any point in execution.
     pub bindings: Bindings<'a>,
+    // Stack of all references created in any scope, at any point in execution.
+    pub references: References,
     // 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.
@@ -63,6 +66,7 @@ impl<'a> SemanticModel<'a> {
             definitions: Definitions::for_module(module),
             definition_id: DefinitionId::module(),
             bindings: Bindings::default(),
+            references: References::default(),
             shadowed_bindings: IntMap::default(),
             body: &[],
             body_index: 0,
@@ -119,17 +123,33 @@ impl<'a> SemanticModel<'a> {
         // PEP 563 indicates that if a forward reference can be resolved in the module scope, we
         // should prefer it over local resolutions.
         if self.in_deferred_type_definition() {
-            if let Some(binding_id) = self.scopes.global().get(symbol) {
+            if let Some(binding_id) = self.scopes.global().get(symbol).copied() {
                 // Mark the binding as used.
                 let context = self.execution_context();
-                self.bindings[*binding_id].mark_used(ScopeId::global(), range, context);
+                let reference_id = self.references.push(
+                    ScopeId::global(),
+                    range,
+                    match context {
+                        ExecutionContext::Runtime => ReferenceContext::Runtime,
+                        ExecutionContext::Typing => ReferenceContext::Typing,
+                    },
+                );
+                self.bindings[binding_id].references.push(reference_id);
 
                 // Mark any submodule aliases as used.
-                if let Some(binding_id) = self.resolve_submodule(ScopeId::global(), *binding_id) {
-                    self.bindings[binding_id].mark_used(ScopeId::global(), range, context);
+                if let Some(binding_id) = self.resolve_submodule(ScopeId::global(), binding_id) {
+                    let reference_id = self.references.push(
+                        ScopeId::global(),
+                        range,
+                        match context {
+                            ExecutionContext::Runtime => ReferenceContext::Runtime,
+                            ExecutionContext::Typing => ReferenceContext::Typing,
+                        },
+                    );
+                    self.bindings[binding_id].references.push(reference_id);
                 }
 
-                return ResolvedReference::Resolved(*binding_id);
+                return ResolvedReference::Resolved(binding_id);
             }
         }
 
@@ -153,14 +173,30 @@ impl<'a> SemanticModel<'a> {
                 }
             }
 
-            if let Some(binding_id) = scope.get(symbol) {
+            if let Some(binding_id) = scope.get(symbol).copied() {
                 // Mark the binding as used.
                 let context = self.execution_context();
-                self.bindings[*binding_id].mark_used(self.scope_id, range, context);
+                let reference_id = self.references.push(
+                    self.scope_id,
+                    range,
+                    match context {
+                        ExecutionContext::Runtime => ReferenceContext::Runtime,
+                        ExecutionContext::Typing => ReferenceContext::Typing,
+                    },
+                );
+                self.bindings[binding_id].references.push(reference_id);
 
                 // Mark any submodule aliases as used.
-                if let Some(binding_id) = self.resolve_submodule(scope_id, *binding_id) {
-                    self.bindings[binding_id].mark_used(ScopeId::global(), range, context);
+                if let Some(binding_id) = self.resolve_submodule(scope_id, binding_id) {
+                    let reference_id = self.references.push(
+                        self.scope_id,
+                        range,
+                        match context {
+                            ExecutionContext::Runtime => ReferenceContext::Runtime,
+                            ExecutionContext::Typing => ReferenceContext::Typing,
+                        },
+                    );
+                    self.bindings[binding_id].references.push(reference_id);
                 }
 
                 // But if it's a type annotation, don't treat it as resolved, unless we're in a
@@ -174,12 +210,12 @@ impl<'a> SemanticModel<'a> {
                 // The `name` in `print(name)` should be treated as unresolved, but the `name` in
                 // `name: str` should be treated as used.
                 if !self.in_deferred_type_definition()
-                    && self.bindings[*binding_id].kind.is_annotation()
+                    && self.bindings[binding_id].kind.is_annotation()
                 {
                     continue;
                 }
 
-                return ResolvedReference::Resolved(*binding_id);
+                return ResolvedReference::Resolved(binding_id);
             }
 
             // Allow usages of `__module__` and `__qualname__` within class scopes, e.g.:
@@ -343,8 +379,8 @@ impl<'a> SemanticModel<'a> {
         member: &str,
     ) -> Option<(&Stmt, String)> {
         self.scopes().enumerate().find_map(|(scope_index, scope)| {
-            scope.binding_ids().find_map(|binding_index| {
-                let binding = &self.bindings[*binding_index];
+            scope.binding_ids().copied().find_map(|binding_id| {
+                let binding = &self.bindings[binding_id];
                 match &binding.kind {
                     // Ex) Given `module="sys"` and `object="exit"`:
                     // `import sys`         -> `sys.exit`
@@ -519,11 +555,17 @@ impl<'a> SemanticModel<'a> {
         self.scopes.ancestors(self.scope_id)
     }
 
+    /// Returns an iterator over all parent statements.
     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])
     }
 
+    /// 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
+    }
+
     /// Return `true` if the context is at the top level of the module (i.e., in the module scope,
     /// and not nested within any statements).
     pub fn at_top_level(&self) -> bool {
@@ -533,6 +575,33 @@ impl<'a> SemanticModel<'a> {
                 .map_or(true, |stmt_id| self.stmts.parent_id(stmt_id).is_none())
     }
 
+    /// Returns `true` if the given [`BindingId`] is used.
+    pub fn is_used(&self, binding_id: BindingId) -> bool {
+        self.bindings[binding_id].is_used()
+    }
+
+    /// Add a reference to the given [`BindingId`] in the local scope.
+    pub fn add_local_reference(
+        &mut self,
+        binding_id: BindingId,
+        range: TextRange,
+        context: ReferenceContext,
+    ) {
+        let reference_id = self.references.push(self.scope_id, range, context);
+        self.bindings[binding_id].references.push(reference_id);
+    }
+
+    /// Add a reference to the given [`BindingId`] in the global scope.
+    pub fn add_global_reference(
+        &mut self,
+        binding_id: BindingId,
+        range: TextRange,
+        context: ReferenceContext,
+    ) {
+        let reference_id = self.references.push(ScopeId::global(), range, context);
+        self.bindings[binding_id].references.push(reference_id);
+    }
+
     /// Return the [`ExecutionContext`] of the current scope.
     pub const fn execution_context(&self) -> ExecutionContext {
         if self.in_type_checking_block()

crates/ruff_python_semantic/src/reference.rs

@@ -0,0 +1,70 @@
+use ruff_text_size::TextRange;
+
+use ruff_index::{newtype_index, IndexVec};
+
+use crate::scope::ScopeId;
+
+#[derive(Debug, Clone)]
+pub struct Reference {
+    /// The scope in which the reference is defined.
+    scope_id: ScopeId,
+    /// The range of the reference in the source code.
+    range: TextRange,
+    /// The context in which the reference occurs.
+    context: ReferenceContext,
+}
+
+impl Reference {
+    pub const fn scope_id(&self) -> ScopeId {
+        self.scope_id
+    }
+
+    pub const fn range(&self) -> TextRange {
+        self.range
+    }
+
+    pub const fn context(&self) -> &ReferenceContext {
+        &self.context
+    }
+}
+
+#[derive(Debug, Clone, is_macro::Is)]
+pub enum ReferenceContext {
+    /// The reference occurs in a runtime context.
+    Runtime,
+    /// The reference occurs in a typing-only context.
+    Typing,
+    /// The reference occurs in a synthetic context, used for `__future__` imports, explicit
+    /// re-exports, and other bindings that should be considered used even if they're never
+    /// "referenced".
+    Synthetic,
+}
+
+/// Id uniquely identifying a read reference in a program.
+#[newtype_index]
+pub struct ReferenceId;
+
+/// The references of a program indexed by [`ReferenceId`].
+#[derive(Debug, Default)]
+pub struct References(IndexVec<ReferenceId, Reference>);
+
+impl References {
+    /// Pushes a new read reference and returns its unique id.
+    pub fn push(
+        &mut self,
+        scope_id: ScopeId,
+        range: TextRange,
+        context: ReferenceContext,
+    ) -> ReferenceId {
+        self.0.push(Reference {
+            scope_id,
+            range,
+            context,
+        })
+    }
+
+    /// Returns the [`Reference`] with the given id.
+    pub fn resolve(&self, id: ReferenceId) -> &Reference {
+        &self.0[id]
+    }
+}