diff --git a/crates/red_knot_python_semantic/resources/mdtest/comprehensions/basic.md b/crates/red_knot_python_semantic/resources/mdtest/comprehensions/basic.md index caeb5f95b747f..5236b27cce7e4 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/comprehensions/basic.md +++ b/crates/red_knot_python_semantic/resources/mdtest/comprehensions/basic.md @@ -43,8 +43,7 @@ class IntIterable: def __iter__(self) -> IntIterator: return IntIterator() -# TODO: This could be a `tuple[int, int]` if we model that `y` can not be modified in the outer comprehension scope -# revealed: tuple[int, Unknown | int] +# revealed: tuple[int, int] [[reveal_type((x, y)) for x in IntIterable()] for y in IntIterable()] ``` @@ -67,8 +66,7 @@ class IterableOfIterables: def __iter__(self) -> IteratorOfIterables: return IteratorOfIterables() -# TODO: This could be a `tuple[int, int]` (see above) -# revealed: tuple[int, Unknown | IntIterable] +# revealed: tuple[int, IntIterable] [[reveal_type((x, y)) for x in y] for y in IterableOfIterables()] ``` diff --git a/crates/red_knot_python_semantic/resources/mdtest/scopes/eager.md b/crates/red_knot_python_semantic/resources/mdtest/scopes/eager.md new file mode 100644 index 0000000000000..0c0362425a0c4 --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/scopes/eager.md @@ -0,0 +1,382 @@ +# Eager scopes + +Some scopes are executed eagerly: references to variables defined in enclosing scopes are resolved +_immediately_. This is in constrast to (for instance) function scopes, where those references are +resolved when the function is called. + +## Function definitions + +Function definitions are evaluated lazily. + +```py +x = 1 + +def f(): + reveal_type(x) # revealed: Unknown | Literal[2] + +x = 2 +``` + +## Class definitions + +Class definitions are evaluated eagerly. + +```py +def _(): + x = 1 + + class A: + reveal_type(x) # revealed: Literal[1] + + y = x + + x = 2 + + reveal_type(A.y) # revealed: Unknown | Literal[1] +``` + +## List comprehensions + +List comprehensions are evaluated eagerly. + +```py +def _(): + x = 1 + + # revealed: Literal[1] + [reveal_type(x) for a in range(1)] + + x = 2 +``` + +## Set comprehensions + +Set comprehensions are evaluated eagerly. + +```py +def _(): + x = 1 + + # revealed: Literal[1] + {reveal_type(x) for a in range(1)} + + x = 2 +``` + +## Dict comprehensions + +Dict comprehensions are evaluated eagerly. + +```py +def _(): + x = 1 + + # revealed: Literal[1] + {a: reveal_type(x) for a in range(1)} + + x = 2 +``` + +## Generator expressions + +Generator expressions don't necessarily run eagerly, but in practice usually they do, so assuming +they do is the better default. + +```py +def _(): + x = 1 + + # revealed: Literal[1] + list(reveal_type(x) for a in range(1)) + + x = 2 +``` + +But that does lead to incorrect results when the generator expression isn't run immediately: + +```py +def evaluated_later(): + x = 1 + + # revealed: Literal[1] + y = (reveal_type(x) for a in range(1)) + + x = 2 + + # The generator isn't evaluated until here, so at runtime, `x` will evaluate to 2, contradicting + # our inferred type. + print(next(y)) +``` + +Though note that “the iterable expression in the leftmost `for` clause is immediately evaluated” +\[[spec][generators]\]: + +```py +def iterable_evaluated_eagerly(): + x = 1 + + # revealed: Literal[1] + y = (a for a in [reveal_type(x)]) + + x = 2 + + # Even though the generator isn't evaluated until here, the first iterable was evaluated + # immediately, so our inferred type is correct. + print(next(y)) +``` + +## Top-level eager scopes + +All of the above examples behave identically when the eager scopes are directly nested in the global +scope. + +### Class definitions + +```py +x = 1 + +class A: + reveal_type(x) # revealed: Literal[1] + + y = x + +x = 2 + +reveal_type(A.y) # revealed: Unknown | Literal[1] +``` + +### List comprehensions + +```py +x = 1 + +# revealed: Literal[1] +[reveal_type(x) for a in range(1)] + +x = 2 +``` + +### Set comprehensions + +```py +x = 1 + +# revealed: Literal[1] +{reveal_type(x) for a in range(1)} + +x = 2 +``` + +### Dict comprehensions + +```py +x = 1 + +# revealed: Literal[1] +{a: reveal_type(x) for a in range(1)} + +x = 2 +``` + +### Generator expressions + +```py +x = 1 + +# revealed: Literal[1] +list(reveal_type(x) for a in range(1)) + +x = 2 +``` + +`evaluated_later.py`: + +```py +x = 1 + +# revealed: Literal[1] +y = (reveal_type(x) for a in range(1)) + +x = 2 + +# The generator isn't evaluated until here, so at runtime, `x` will evaluate to 2, contradicting +# our inferred type. +print(next(y)) +``` + +`iterable_evaluated_eagerly.py`: + +```py +x = 1 + +# revealed: Literal[1] +y = (a for a in [reveal_type(x)]) + +x = 2 + +# Even though the generator isn't evaluated until here, the first iterable was evaluated +# immediately, so our inferred type is correct. +print(next(y)) +``` + +## Lazy scopes are "sticky" + +As we look through each enclosing scope when resolving a reference, lookups become lazy as soon as +we encounter any lazy scope, even if there are other eager scopes that enclose it. + +### Eager scope within eager scope + +If we don't encounter a lazy scope, lookup remains eager. The resolved binding is not necessarily in +the immediately enclosing scope. Here, the list comprehension and class definition are both eager +scopes, and we immediately resolve the use of `x` to (only) the `x = 1` binding. + +```py +def _(): + x = 1 + + class A: + # revealed: Literal[1] + [reveal_type(x) for a in range(1)] + + x = 2 +``` + +### Class definition bindings are not visible in nested scopes + +Class definitions are eager scopes, but any bindings in them are explicitly not visible to any +nested scopes. (Those nested scopes are typically (lazy) function definitions, but the rule also +applies to nested eager scopes like comprehensions and other class definitions.) + +```py +def _(): + x = 1 + + class A: + x = 4 + + # revealed: Literal[1] + [reveal_type(x) for a in range(1)] + + class B: + # revealed: Literal[1] + [reveal_type(x) for a in range(1)] + + x = 2 +``` + +### Eager scope within a lazy scope + +The list comprehension is an eager scope, and it is enclosed within a function definition, which is +a lazy scope. Because we pass through this lazy scope before encountering any bindings or +definitions, the lookup is lazy. + +```py +def _(): + x = 1 + + def f(): + # revealed: Unknown | Literal[2] + [reveal_type(x) for a in range(1)] + x = 2 +``` + +### Lazy scope within an eager scope + +The function definition is a lazy scope, and it is enclosed within a class definition, which is an +eager scope. Even though we pass through an eager scope before encountering any bindings or +definitions, the lookup remains lazy. + +```py +def _(): + x = 1 + + class A: + def f(): + # revealed: Unknown | Literal[2] + reveal_type(x) + + x = 2 +``` + +### Lazy scope within a lazy scope + +No matter how many lazy scopes we pass through before encountering a binding or definition, the +lookup remains lazy. + +```py +def _(): + x = 1 + + def f(): + def g(): + # revealed: Unknown | Literal[2] + reveal_type(x) + x = 2 +``` + +### Eager scope within a lazy scope within another eager scope + +We have a list comprehension (eager scope), enclosed within a function definition (lazy scope), +enclosed within a class definition (eager scope), all of which we must pass through before +encountering any binding of `x`. Even though the last scope we pass through is eager, the lookup is +lazy, since we encountered a lazy scope on the way. + +```py +def _(): + x = 1 + + class A: + def f(): + # revealed: Unknown | Literal[2] + [reveal_type(x) for a in range(1)] + + x = 2 +``` + +## Annotations + +Type annotations are sometimes deferred. When they are, the types that are referenced in an +annotation are looked up lazily, even if they occur in an eager scope. + +### Eager annotations in a Python file + +```py +x = int + +class C: + var: x + +reveal_type(C.var) # revealed: int + +x = str +``` + +### Deferred annotations in a Python file + +```py +from __future__ import annotations + +x = int + +class C: + var: x + +reveal_type(C.var) # revealed: Unknown | str + +x = str +``` + +### Deferred annotations in a stub file + +```pyi +x = int + +class C: + var: x + +reveal_type(C.var) # revealed: Unknown | str + +x = str +``` + +[generators]: https://docs.python.org/3/reference/expressions.html#generator-expressions diff --git a/crates/red_knot_python_semantic/src/semantic_index.rs b/crates/red_knot_python_semantic/src/semantic_index.rs index f2058e54b1c3c..c3c72cfe2812e 100644 --- a/crates/red_knot_python_semantic/src/semantic_index.rs +++ b/crates/red_knot_python_semantic/src/semantic_index.rs @@ -19,7 +19,7 @@ use crate::semantic_index::expression::Expression; use crate::semantic_index::symbol::{ FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopedSymbolId, SymbolTable, }; -use crate::semantic_index::use_def::UseDefMap; +use crate::semantic_index::use_def::{EagerBindingsKey, ScopedEagerBindingsId, UseDefMap}; use crate::Db; pub mod ast_ids; @@ -165,6 +165,9 @@ pub(crate) struct SemanticIndex<'db> { /// Maps from class body scopes to attribute assignments that were found /// in methods of that class. attribute_assignments: FxHashMap>>, + + /// Map of all of the eager bindings that appear in this file. + eager_bindings: FxHashMap, } impl<'db> SemanticIndex<'db> { @@ -220,7 +223,7 @@ impl<'db> SemanticIndex<'db> { /// Returns the id of the parent scope. pub(crate) fn parent_scope_id(&self, scope_id: FileScopeId) -> Option { let scope = self.scope(scope_id); - scope.parent + scope.parent() } /// Returns the parent scope of `scope_id`. @@ -290,6 +293,23 @@ impl<'db> SemanticIndex<'db> { pub(super) fn has_future_annotations(&self) -> bool { self.has_future_annotations } + + /// Returns an iterator of bindings for a particular nested eager scope reference. + pub(crate) fn eager_bindings( + &self, + enclosing_scope: FileScopeId, + symbol: &str, + nested_scope: FileScopeId, + ) -> Option> { + let symbol_id = self.symbol_tables[enclosing_scope].symbol_id_by_name(symbol)?; + let key = EagerBindingsKey { + enclosing_scope, + enclosing_symbol: symbol_id, + nested_scope, + }; + let id = self.eager_bindings.get(&key)?; + self.use_def_maps[enclosing_scope].eager_bindings(*id) + } } pub struct AncestorsIter<'a> { @@ -312,7 +332,7 @@ impl<'a> Iterator for AncestorsIter<'a> { fn next(&mut self) -> Option { let current_id = self.next_id?; let current = &self.scopes[current_id]; - self.next_id = current.parent; + self.next_id = current.parent(); Some((current_id, current)) } @@ -328,7 +348,7 @@ pub struct DescendentsIter<'a> { impl<'a> DescendentsIter<'a> { fn new(symbol_table: &'a SemanticIndex, scope_id: FileScopeId) -> Self { let scope = &symbol_table.scopes[scope_id]; - let scopes = &symbol_table.scopes[scope.descendents.clone()]; + let scopes = &symbol_table.scopes[scope.descendents()]; Self { next_id: scope_id + 1, @@ -378,7 +398,7 @@ impl<'a> Iterator for ChildrenIter<'a> { fn next(&mut self) -> Option { self.descendents - .find(|(_, scope)| scope.parent == Some(self.parent)) + .find(|(_, scope)| scope.parent() == Some(self.parent)) } } diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 9e97a1bf2ed7c..e8bd21917bbaa 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -25,7 +25,9 @@ use crate::semantic_index::symbol::{ FileScopeId, NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopeKind, ScopedSymbolId, SymbolTableBuilder, }; -use crate::semantic_index::use_def::{FlowSnapshot, ScopedConstraintId, UseDefMapBuilder}; +use crate::semantic_index::use_def::{ + EagerBindingsKey, FlowSnapshot, ScopedConstraintId, ScopedEagerBindingsId, UseDefMapBuilder, +}; use crate::semantic_index::SemanticIndex; use crate::unpack::{Unpack, UnpackValue}; use crate::visibility_constraints::{ScopedVisibilityConstraintId, VisibilityConstraintsBuilder}; @@ -91,6 +93,7 @@ pub(super) struct SemanticIndexBuilder<'db> { expressions_by_node: FxHashMap>, imported_modules: FxHashSet, attribute_assignments: FxHashMap>, + eager_bindings: FxHashMap, } impl<'db> SemanticIndexBuilder<'db> { @@ -122,6 +125,8 @@ impl<'db> SemanticIndexBuilder<'db> { imported_modules: FxHashSet::default(), attribute_assignments: FxHashMap::default(), + + eager_bindings: FxHashMap::default(), }; builder.push_scope_with_parent(NodeWithScopeRef::Module, None); @@ -134,13 +139,13 @@ impl<'db> SemanticIndexBuilder<'db> { .scope_stack .last() .map(|ScopeInfo { file_scope_id, .. }| file_scope_id) - .expect("Always to have a root scope") + .expect("SemanticIndexBuilder should have created a root scope") } fn loop_state(&self) -> LoopState { self.scope_stack .last() - .expect("Always to have a root scope") + .expect("SemanticIndexBuilder should have created a root scope") .loop_state } @@ -177,13 +182,11 @@ impl<'db> SemanticIndexBuilder<'db> { fn push_scope_with_parent(&mut self, node: NodeWithScopeRef, parent: Option) { let children_start = self.scopes.next_index() + 1; + // SAFETY: `node` is guaranteed to be a child of `self.module` #[allow(unsafe_code)] - let scope = Scope { - parent, - // SAFETY: `node` is guaranteed to be a child of `self.module` - node: unsafe { node.to_kind(self.module.clone()) }, - descendents: children_start..children_start, - }; + let node_with_kind = unsafe { node.to_kind(self.module.clone()) }; + + let scope = Scope::new(parent, node_with_kind, children_start..children_start); self.try_node_context_stack_manager.enter_nested_scope(); let file_scope_id = self.scopes.push(scope); @@ -206,13 +209,74 @@ impl<'db> SemanticIndexBuilder<'db> { } fn pop_scope(&mut self) -> FileScopeId { - let ScopeInfo { file_scope_id, .. } = - self.scope_stack.pop().expect("Root scope to be present"); - let children_end = self.scopes.next_index(); - let scope = &mut self.scopes[file_scope_id]; - scope.descendents = scope.descendents.start..children_end; self.try_node_context_stack_manager.exit_scope(); - file_scope_id + + let ScopeInfo { + file_scope_id: popped_scope_id, + .. + } = self + .scope_stack + .pop() + .expect("Root scope should be present"); + + let children_end = self.scopes.next_index(); + let popped_scope = &mut self.scopes[popped_scope_id]; + popped_scope.extend_descendents(children_end); + + if !popped_scope.is_eager() { + return popped_scope_id; + } + + // If the scope that we just popped off is an eager scope, we need to "lock" our view of + // which bindings reach each of the uses in the scope. Loop through each enclosing scope, + // looking for any that bind each symbol. + for enclosing_scope_info in self.scope_stack.iter().rev() { + let enclosing_scope_id = enclosing_scope_info.file_scope_id; + let enclosing_scope_kind = self.scopes[enclosing_scope_id].kind(); + let enclosing_symbol_table = &self.symbol_tables[enclosing_scope_id]; + + // Names bound in class scopes are never visible to nested scopes, so we never need to + // save eager scope bindings in a class scope. + if enclosing_scope_kind.is_class() { + continue; + } + + for nested_symbol in self.symbol_tables[popped_scope_id].symbols() { + // Skip this symbol if this enclosing scope doesn't contain any bindings for + // it, or if the nested scope _does_. + if nested_symbol.is_bound() { + continue; + } + let Some(enclosing_symbol_id) = + enclosing_symbol_table.symbol_id_by_name(nested_symbol.name()) + else { + continue; + }; + let enclosing_symbol = enclosing_symbol_table.symbol(enclosing_symbol_id); + if !enclosing_symbol.is_bound() { + continue; + } + + // Snapshot the bindings of this symbol that are visible at this point in this + // enclosing scope. + let key = EagerBindingsKey { + enclosing_scope: enclosing_scope_id, + enclosing_symbol: enclosing_symbol_id, + nested_scope: popped_scope_id, + }; + let eager_bindings = self.use_def_maps[enclosing_scope_id] + .snapshot_eager_bindings(enclosing_symbol_id); + self.eager_bindings.insert(key, eager_bindings); + } + + // Lazy scopes are "sticky": once we see a lazy scope we stop doing lookups + // eagerly, even if we would encounter another eager enclosing scope later on. + if !enclosing_scope_kind.is_eager() { + break; + } + } + + popped_scope_id } fn current_symbol_table(&mut self) -> &mut SymbolTableBuilder { @@ -729,6 +793,7 @@ impl<'db> SemanticIndexBuilder<'db> { self.scope_ids_by_scope.shrink_to_fit(); self.scopes_by_node.shrink_to_fit(); + self.eager_bindings.shrink_to_fit(); SemanticIndex { symbol_tables, @@ -747,6 +812,7 @@ impl<'db> SemanticIndexBuilder<'db> { .into_iter() .map(|(k, v)| (k, Arc::new(v))) .collect(), + eager_bindings: self.eager_bindings, } } } diff --git a/crates/red_knot_python_semantic/src/semantic_index/symbol.rs b/crates/red_knot_python_semantic/src/semantic_index/symbol.rs index d2f93062411eb..5e678c2526703 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/symbol.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/symbol.rs @@ -111,16 +111,7 @@ pub struct ScopeId<'db> { impl<'db> ScopeId<'db> { pub(crate) fn is_function_like(self, db: &'db dyn Db) -> bool { - // Type parameter scopes behave like function scopes in terms of name resolution; CPython - // symbol table also uses the term "function-like" for these scopes. - matches!( - self.node(db).scope_kind(), - ScopeKind::Annotation - | ScopeKind::Function - | ScopeKind::Lambda - | ScopeKind::TypeAlias - | ScopeKind::Comprehension - ) + self.node(db).scope_kind().is_function_like() } pub(crate) fn node(self, db: &dyn Db) -> &NodeWithScopeKind { @@ -178,13 +169,25 @@ impl FileScopeId { #[derive(Debug, salsa::Update)] pub struct Scope { - pub(super) parent: Option, - pub(super) node: NodeWithScopeKind, - pub(super) descendents: Range, + parent: Option, + node: NodeWithScopeKind, + descendents: Range, } impl Scope { - pub fn parent(self) -> Option { + pub(super) fn new( + parent: Option, + node: NodeWithScopeKind, + descendents: Range, + ) -> Self { + Scope { + parent, + node, + descendents, + } + } + + pub fn parent(&self) -> Option { self.parent } @@ -195,6 +198,18 @@ impl Scope { pub fn kind(&self) -> ScopeKind { self.node().scope_kind() } + + pub fn descendents(&self) -> Range { + self.descendents.clone() + } + + pub(super) fn extend_descendents(&mut self, children_end: FileScopeId) { + self.descendents = self.descendents.start..children_end; + } + + pub(crate) fn is_eager(&self) -> bool { + self.kind().is_eager() + } } #[derive(Copy, Clone, Debug, PartialEq, Eq)] @@ -209,8 +224,32 @@ pub enum ScopeKind { } impl ScopeKind { - pub const fn is_comprehension(self) -> bool { - matches!(self, ScopeKind::Comprehension) + pub(crate) fn is_eager(self) -> bool { + match self { + ScopeKind::Class | ScopeKind::Comprehension => true, + ScopeKind::Module + | ScopeKind::Annotation + | ScopeKind::Function + | ScopeKind::Lambda + | ScopeKind::TypeAlias => false, + } + } + + pub(crate) fn is_function_like(self) -> bool { + // Type parameter scopes behave like function scopes in terms of name resolution; CPython + // symbol table also uses the term "function-like" for these scopes. + matches!( + self, + ScopeKind::Annotation + | ScopeKind::Function + | ScopeKind::Lambda + | ScopeKind::TypeAlias + | ScopeKind::Comprehension + ) + } + + pub(crate) fn is_class(self) -> bool { + matches!(self, ScopeKind::Class) } } @@ -316,6 +355,18 @@ impl SymbolTableBuilder { self.table.symbols[id].insert_flags(SymbolFlags::IS_USED); } + pub(super) fn symbols(&self) -> impl Iterator { + self.table.symbols() + } + + pub(super) fn symbol_id_by_name(&self, name: &str) -> Option { + self.table.symbol_id_by_name(name) + } + + pub(super) fn symbol(&self, symbol_id: impl Into) -> &Symbol { + self.table.symbol(symbol_id) + } + pub(super) fn finish(mut self) -> SymbolTable { self.table.shrink_to_fit(); self.table diff --git a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs index e4d201170819a..ce987ef40382b 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs @@ -262,12 +262,12 @@ use self::symbol_state::{ }; use crate::semantic_index::ast_ids::ScopedUseId; use crate::semantic_index::definition::Definition; -use crate::semantic_index::symbol::ScopedSymbolId; +use crate::semantic_index::symbol::{FileScopeId, ScopedSymbolId}; use crate::semantic_index::use_def::symbol_state::DeclarationIdWithConstraint; use crate::visibility_constraints::{ ScopedVisibilityConstraintId, VisibilityConstraints, VisibilityConstraintsBuilder, }; -use ruff_index::IndexVec; +use ruff_index::{newtype_index, IndexVec}; use rustc_hash::FxHashMap; use super::constraint::Constraint; @@ -309,6 +309,10 @@ pub(crate) struct UseDefMap<'db> { /// [`SymbolState`] visible at end of scope for each symbol. public_symbols: IndexVec, + + /// Snapshot of bindings in this scope that can be used to resolve a reference in a nested + /// eager scope. + eager_bindings: EagerBindings, } impl<'db> UseDefMap<'db> { @@ -326,6 +330,15 @@ impl<'db> UseDefMap<'db> { self.bindings_iterator(self.public_symbols[symbol].bindings()) } + pub(crate) fn eager_bindings( + &self, + eager_bindings: ScopedEagerBindingsId, + ) -> Option> { + self.eager_bindings + .get(eager_bindings) + .map(|symbol_bindings| self.bindings_iterator(symbol_bindings)) + } + pub(crate) fn bindings_at_declaration( &self, declaration: Definition<'db>, @@ -383,6 +396,30 @@ impl<'db> UseDefMap<'db> { } } +/// Uniquely identifies a snapshot of bindings that can be used to resolve a reference in a nested +/// eager scope. +/// +/// An eager scope has its entire body executed immediately at the location where it is defined. +/// For any free references in the nested scope, we use the bindings that are visible at the point +/// where the nested scope is defined, instead of using the public type of the symbol. +/// +/// There is a unique ID for each distinct [`EagerBindingsKey`] in the file. +#[newtype_index] +pub(crate) struct ScopedEagerBindingsId; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) struct EagerBindingsKey { + /// The enclosing scope containing the bindings + pub(crate) enclosing_scope: FileScopeId, + /// The referenced symbol (in the enclosing scope) + pub(crate) enclosing_symbol: ScopedSymbolId, + /// The nested eager scope containing the reference + pub(crate) nested_scope: FileScopeId, +} + +/// A snapshot of bindings that can be used to resolve a reference in a nested eager scope. +type EagerBindings = IndexVec; + /// Either live bindings or live declarations for a symbol. #[derive(Debug, PartialEq, Eq, salsa::Update)] enum SymbolDefinitions { @@ -505,6 +542,10 @@ pub(super) struct UseDefMapBuilder<'db> { /// Currently live bindings and declarations for each symbol. symbol_states: IndexVec, + + /// Snapshot of bindings in this scope that can be used to resolve a reference in a nested + /// eager scope. + eager_bindings: EagerBindings, } impl Default for UseDefMapBuilder<'_> { @@ -517,6 +558,7 @@ impl Default for UseDefMapBuilder<'_> { bindings_by_use: IndexVec::new(), definitions_by_definition: FxHashMap::default(), symbol_states: IndexVec::new(), + eager_bindings: EagerBindings::default(), } } } @@ -644,6 +686,14 @@ impl<'db> UseDefMapBuilder<'db> { debug_assert_eq!(use_id, new_use); } + pub(super) fn snapshot_eager_bindings( + &mut self, + enclosing_symbol: ScopedSymbolId, + ) -> ScopedEagerBindingsId { + self.eager_bindings + .push(self.symbol_states[enclosing_symbol].bindings().clone()) + } + /// Take a snapshot of the current visible-symbols state. pub(super) fn snapshot(&self) -> FlowSnapshot { FlowSnapshot { @@ -721,6 +771,7 @@ impl<'db> UseDefMapBuilder<'db> { self.symbol_states.shrink_to_fit(); self.bindings_by_use.shrink_to_fit(); self.definitions_by_definition.shrink_to_fit(); + self.eager_bindings.shrink_to_fit(); UseDefMap { all_definitions: self.all_definitions, @@ -729,6 +780,7 @@ impl<'db> UseDefMapBuilder<'db> { bindings_by_use: self.bindings_by_use, public_symbols: self.symbol_states, definitions_by_definition: self.definitions_by_definition, + eager_bindings: self.eager_bindings, } } } diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 48261b798c533..b405b4c5cbe23 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -47,7 +47,7 @@ use crate::semantic_index::definition::{ }; use crate::semantic_index::expression::{Expression, ExpressionKind}; use crate::semantic_index::semantic_index; -use crate::semantic_index::symbol::{NodeWithScopeKind, NodeWithScopeRef, ScopeId}; +use crate::semantic_index::symbol::{FileScopeId, NodeWithScopeKind, NodeWithScopeRef, ScopeId}; use crate::semantic_index::SemanticIndex; use crate::symbol::{ builtins_module_scope, builtins_symbol, symbol, symbol_from_bindings, symbol_from_declarations, @@ -3498,7 +3498,9 @@ impl<'db> TypeInferenceBuilder<'db> { // Walk up parent scopes looking for a possible enclosing scope that may have a // definition of this name visible to us (would be `LOAD_DEREF` at runtime.) - for (enclosing_scope_file_id, _) in self.index.ancestor_scopes(file_scope_id) { + // Note that we skip the scope containing the use that we are resolving, since we + // already looked for the symbol there up above. + for (enclosing_scope_file_id, _) in self.index.ancestor_scopes(file_scope_id).skip(1) { // Class scopes are not visible to nested scopes, and we need to handle global // scope differently (because an unbound name there falls back to builtins), so // check only function-like scopes. @@ -3506,6 +3508,23 @@ impl<'db> TypeInferenceBuilder<'db> { if !enclosing_scope_id.is_function_like(db) { continue; } + + // If the reference is in a nested eager scope, we need to look for the symbol at + // the point where the previous enclosing scope was defined, instead of at the end + // of the scope. (Note that the semantic index builder takes care of only + // registering eager bindings for nested scopes that are actually eager, and for + // enclosing scopes that actually contain bindings that we should use when + // resolving the reference.) + if !self.is_deferred() { + if let Some(bindings) = self.index.eager_bindings( + enclosing_scope_file_id, + symbol_name, + file_scope_id, + ) { + return symbol_from_bindings(db, bindings); + } + } + let enclosing_symbol_table = self.index.symbol_table(enclosing_scope_file_id); let Some(enclosing_symbol) = enclosing_symbol_table.symbol_by_name(symbol_name) else { @@ -3526,10 +3545,20 @@ impl<'db> TypeInferenceBuilder<'db> { // Avoid infinite recursion if `self.scope` already is the module's global scope. .or_fall_back_to(db, || { if file_scope_id.is_global() { - Symbol::Unbound - } else { - global_symbol(db, self.file(), symbol_name) + return Symbol::Unbound; } + + if !self.is_deferred() { + if let Some(bindings) = self.index.eager_bindings( + FileScopeId::global(), + symbol_name, + file_scope_id, + ) { + return symbol_from_bindings(db, bindings); + } + } + + global_symbol(db, self.file(), symbol_name) }) // Not found in globals? Fallback to builtins // (without infinite recursion if we're already in builtins.)