[ty] Use an interval map for scopes by expression

Author: MichaReiser

crates/ruff_python_ast/src/node_index.rs

@@ -48,6 +48,16 @@ impl NodeIndex {
     pub fn as_usize(self) -> usize {
         self.0 as _
     }
+
+    pub fn as_u32(self) -> u32 {
+        self.0
+    }
+}
+
+impl From<u32> for NodeIndex {
+    fn from(value: u32) -> Self {
+        NodeIndex(value)
+    }
 }
 
 impl From<u32> for AtomicNodeIndex {

crates/ty_python_semantic/src/semantic_index.rs

@@ -5,6 +5,7 @@ use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
 use ruff_index::{IndexSlice, IndexVec};
 
+use ruff_python_ast::NodeIndex;
 use ruff_python_parser::semantic_errors::SemanticSyntaxError;
 use rustc_hash::{FxHashMap, FxHashSet};
 use salsa::Update;
@@ -24,6 +25,7 @@ use crate::semantic_index::place::{
     ScopeKind, ScopedPlaceId,
 };
 use crate::semantic_index::use_def::{EagerSnapshotKey, ScopedEagerSnapshotId, UseDefMap};
+use crate::semantic_model::HasTrackedScope;
 use crate::util::get_size::untracked_arc_size;
 
 pub mod ast_ids;
@@ -179,7 +181,7 @@ pub(crate) struct SemanticIndex<'db> {
     scopes: IndexVec<FileScopeId, Scope>,
 
     /// Map expressions to their corresponding scope.
-    scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
+    scopes_by_expression: ExpressionsScopeMap,
 
     /// Map from a node creating a definition to its definition.
     definitions_by_node: FxHashMap<DefinitionNodeKey, Definitions<'db>>,
@@ -247,25 +249,26 @@ impl<'db> SemanticIndex<'db> {
 
     /// Returns the ID of the `expression`'s enclosing scope.
     #[track_caller]
-    pub(crate) fn expression_scope_id(
-        &self,
-        expression: impl Into<ExpressionNodeKey>,
-    ) -> FileScopeId {
-        self.scopes_by_expression[&expression.into()]
+    pub(crate) fn expression_scope_id<E>(&self, expression: &E) -> FileScopeId
+    where
+        E: HasTrackedScope,
+    {
+        self.try_expression_scope_id(expression)
+            .expect("Expression to be part of a scope if it is from the same module")
     }
 
     /// Returns the ID of the `expression`'s enclosing scope.
-    pub(crate) fn try_expression_scope_id(
-        &self,
-        expression: impl Into<ExpressionNodeKey>,
-    ) -> Option<FileScopeId> {
-        self.scopes_by_expression.get(&expression.into()).copied()
+    pub(crate) fn try_expression_scope_id<E>(&self, expression: &E) -> Option<FileScopeId>
+    where
+        E: HasTrackedScope,
+    {
+        self.scopes_by_expression.try_get(expression)
     }
 
     /// Returns the [`Scope`] of the `expression`'s enclosing scope.
     #[allow(unused)]
     #[track_caller]
-    pub(crate) fn expression_scope(&self, expression: impl Into<ExpressionNodeKey>) -> &Scope {
+    pub(crate) fn expression_scope(&self, expression: &impl HasTrackedScope) -> &Scope {
         &self.scopes[self.expression_scope_id(expression)]
     }
 
@@ -583,6 +586,40 @@ impl<'a> Iterator for ChildrenIter<'a> {
 
 impl FusedIterator for ChildrenIter<'_> {}
 
+/// Interval map that maps a range of expression node ids to their corresponding scopes.
+///
+/// Lookups require `O(log n)` time, where `n` is roughly the number of scopes (roughly
+/// becaues sub-scopes can be interleaved with expressions in the outer scope, e.g. function, some statements, a function).
+#[derive(Eq, PartialEq, Debug, get_size2::GetSize)]
+struct ExpressionsScopeMap {
+    scopes_by_expression: Box<[(std::ops::Range<NodeIndex>, FileScopeId)]>,
+}
+
+impl ExpressionsScopeMap {
+    fn try_get<E>(&self, node: &E) -> Option<FileScopeId>
+    where
+        E: HasTrackedScope,
+    {
+        let node_index = node.node_index().load();
+
+        let entry = self
+            .scopes_by_expression
+            .binary_search_by_key(&node_index, |(range, _)| range.start);
+
+        let index = match entry {
+            Ok(index) => index,
+            Err(index) => index.checked_sub(1)?,
+        };
+
+        let (range, scope_id) = &self.scopes_by_expression[index];
+        if range.contains(&node_index) {
+            Some(*scope_id)
+        } else {
+            None
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use ruff_db::files::{File, system_path_to_file};

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -2,6 +2,7 @@ use std::cell::{OnceCell, RefCell};
 use std::sync::Arc;
 
 use except_handlers::TryNodeContextStackManager;
+use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet};
 
 use ruff_db::files::File;
@@ -10,7 +11,7 @@ use ruff_db::source::{SourceText, source_text};
 use ruff_index::IndexVec;
 use ruff_python_ast::name::Name;
 use ruff_python_ast::visitor::{Visitor, walk_expr, walk_pattern, walk_stmt};
-use ruff_python_ast::{self as ast, PySourceType, PythonVersion};
+use ruff_python_ast::{self as ast, NodeIndex, PySourceType, PythonVersion};
 use ruff_python_parser::semantic_errors::{
     SemanticSyntaxChecker, SemanticSyntaxContext, SemanticSyntaxError, SemanticSyntaxErrorKind,
 };
@@ -45,7 +46,8 @@ use crate::semantic_index::reachability_constraints::{
 use crate::semantic_index::use_def::{
     EagerSnapshotKey, FlowSnapshot, ScopedEagerSnapshotId, UseDefMapBuilder,
 };
-use crate::semantic_index::{ArcUseDefMap, SemanticIndex};
+use crate::semantic_index::{ArcUseDefMap, ExpressionsScopeMap, SemanticIndex};
+use crate::semantic_model::HasTrackedScope;
 use crate::unpack::{Unpack, UnpackKind, UnpackPosition, UnpackValue};
 use crate::{Db, Program};
 
@@ -102,7 +104,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     ast_ids: IndexVec<FileScopeId, AstIdsBuilder>,
     use_def_maps: IndexVec<FileScopeId, UseDefMapBuilder<'db>>,
     scopes_by_node: FxHashMap<NodeWithScopeKey, FileScopeId>,
-    scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
+    scopes_by_expression: ExpressionsScopeMapBuilder,
     globals_by_scope: FxHashMap<FileScopeId, FxHashSet<ScopedPlaceId>>,
     definitions_by_node: FxHashMap<DefinitionNodeKey, Definitions<'db>>,
     expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
@@ -137,7 +139,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             scope_ids_by_scope: IndexVec::new(),
             use_def_maps: IndexVec::new(),
 
-            scopes_by_expression: FxHashMap::default(),
+            scopes_by_expression: ExpressionsScopeMapBuilder::new(),
             scopes_by_node: FxHashMap::default(),
             definitions_by_node: FxHashMap::default(),
             expressions_by_node: FxHashMap::default(),
@@ -1011,7 +1013,6 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         place_tables.shrink_to_fit();
         use_def_maps.shrink_to_fit();
         ast_ids.shrink_to_fit();
-        self.scopes_by_expression.shrink_to_fit();
         self.definitions_by_node.shrink_to_fit();
 
         self.scope_ids_by_scope.shrink_to_fit();
@@ -1028,7 +1029,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             scope_ids_by_scope: self.scope_ids_by_scope,
             globals_by_scope: self.globals_by_scope,
             ast_ids,
-            scopes_by_expression: self.scopes_by_expression,
+            scopes_by_expression: self.scopes_by_expression.build(),
             scopes_by_node: self.scopes_by_node,
             use_def_maps,
             imported_modules: Arc::new(self.imported_modules),
@@ -1901,7 +1902,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
         self.with_semantic_checker(|semantic, context| semantic.visit_expr(expr, context));
 
         self.scopes_by_expression
-            .insert(expr.into(), self.current_scope());
+            .record_expression(expr, self.current_scope());
 
         let node_key = NodeKey::from_node(expr);
 
@@ -2539,3 +2540,79 @@ fn dunder_all_extend_argument(value: &ast::Expr) -> Option<&ast::Expr> {
 
     (attr == "__all__").then_some(value)
 }
+
+/// Builds a map of expressions to their enclosing scopes.
+///
+/// Ideally, we would condense the expression node ids to scoope index directly when
+/// registering the expression. However, there are a few things that make this hard (but maybe still possible?)
+///
+/// * The expression ids are assigned in source order, but we visit the expressions in semantic order.
+///   There's no difference for most expressions but some expressions will be registered out of order.
+/// * The expression ids aren't always consecutive elements (in fact, they normally aren't). That makes it
+///   harder to detect out of order expressions.
+struct ExpressionsScopeMapBuilder {
+    expression_and_scope: Vec<(NodeIndex, FileScopeId)>,
+    out_of_order: Vec<(NodeIndex, FileScopeId)>,
+}
+
+impl ExpressionsScopeMapBuilder {
+    fn new() -> Self {
+        Self {
+            expression_and_scope: vec![],
+            out_of_order: vec![],
+        }
+    }
+
+    // Ugh, the expression ranges can be off because they are in source order but we visit the
+    // expressions in semantic order.
+    fn record_expression(&mut self, expression: &impl HasTrackedScope, scope: FileScopeId) {
+        let expression_index = expression.node_index().load();
+
+        if self
+            .expression_and_scope
+            .last()
+            .is_some_and(|last| last.0 > expression_index)
+        {
+            self.out_of_order.push((expression_index, scope));
+        } else {
+            self.expression_and_scope.push((expression_index, scope));
+        }
+    }
+
+    fn build(mut self) -> ExpressionsScopeMap {
+        self.out_of_order.sort_by_key(|(index, _)| *index);
+
+        let mut condensed = Vec::new();
+
+        let mut iter = self
+            .expression_and_scope
+            .into_iter()
+            .merge_by(self.out_of_order, |left, right| left.0 <= right.0);
+
+        let Some(first) = iter.next() else {
+            return ExpressionsScopeMap {
+                scopes_by_expression: Box::default(),
+            };
+        };
+
+        let mut current_scope = first.1;
+        let mut range = first.0..NodeIndex::from(first.0.as_u32() + 1);
+
+        for (index, scope) in iter {
+            if scope == current_scope {
+                range.end = NodeIndex::from(index.as_u32() + 1);
+                continue;
+            }
+
+            condensed.push((range, current_scope));
+            current_scope = scope;
+            range = index..NodeIndex::from(index.as_u32() + 1);
+        }
+
+        condensed.push((range, current_scope));
+
+        ExpressionsScopeMap {
+            scopes_by_expression: condensed.into_boxed_slice(),
+        }
+    }
+}

crates/ty_python_semantic/src/semantic_model.rs

@@ -1,7 +1,7 @@
 use ruff_db::files::{File, FilePath};
 use ruff_db::source::line_index;
 use ruff_python_ast as ast;
-use ruff_python_ast::{Expr, ExprRef, name::Name};
+use ruff_python_ast::{Expr, ExprRef, HasNodeIndex, name::Name};
 use ruff_source_file::LineIndex;
 
 use crate::Db;
@@ -108,7 +108,7 @@ impl<'db> SemanticModel<'db> {
                 // expression that we're in, then just
                 // fall back to the global scope.
                 None => Some(FileScopeId::global()),
-                Some(expr) => index.try_expression_scope_id(expr),
+                Some(expr) => index.try_expression_scope_id(&expr),
             },
         }) else {
             return vec![];
@@ -155,7 +155,7 @@ pub trait HasType {
 impl HasType for ast::ExprRef<'_> {
     fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
         let index = semantic_index(model.db, model.file);
-        let file_scope = index.expression_scope_id(*self);
+        let file_scope = index.expression_scope_id(self);
         let scope = file_scope.to_scope_id(model.db, model.file);
 
         infer_scope_types(model.db, scope).expression_type(*self)
@@ -277,6 +277,18 @@ impl HasType for ast::Alias {
     }
 }
 
+/// Implemented by types for which the semantic index tracks their scope.
+pub(crate) trait HasTrackedScope: HasNodeIndex {}
+
+impl HasTrackedScope for ast::Expr {}
+
+impl HasTrackedScope for ast::ExprRef<'_> {}
+impl HasTrackedScope for &ast::ExprRef<'_> {}
+
+// See https://github.com/astral-sh/ty/issues/572 why this implementation exists
+// even when we never register identifiers during semantic index building.
+impl HasTrackedScope for ast::Identifier {}
+
 #[cfg(test)]
 mod tests {
     use ruff_db::files::system_path_to_file;