[red-knot] Resolve references in eager nested scopes eagerly (#16079)

We now resolve references in "eager" scopes correctly — using the
bindings and declarations that are visible at the point where the eager
scope is created, not the "public" type of the symbol (typically the
bindings visible at the end of the scope).

---------

Co-authored-by: Alex Waygood <alex.waygood@gmail.com>

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, int]
-# revealed: tuple[int, Unknown | 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, IntIterable]
-# revealed: tuple[int, Unknown | IntIterable]
 [[reveal_type((x, y)) for x in y] for y in IterableOfIterables()]
 ```
 

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

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<FileScopeId, Arc<AttributeAssignments<'db>>>,
+
+    /// Map of all of the eager bindings that appear in this file.
+    eager_bindings: FxHashMap<EagerBindingsKey, ScopedEagerBindingsId>,
 }
 
 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<FileScopeId> {
         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<BindingWithConstraintsIterator<'_, 'db>> {
+        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<Self::Item> {
         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::Item> {
         self.descendents
-            .find(|(_, scope)| scope.parent == Some(self.parent))
+            .find(|(_, scope)| scope.parent() == Some(self.parent))
     }
 }
 

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<ExpressionNodeKey, Expression<'db>>,
     imported_modules: FxHashSet<ModuleName>,
     attribute_assignments: FxHashMap<FileScopeId, AttributeAssignments<'db>>,
+    eager_bindings: FxHashMap<EagerBindingsKey, ScopedEagerBindingsId>,
 }
 
 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<FileScopeId>) {
         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 {
+        let node_with_kind = unsafe { node.to_kind(self.module.clone()) };
-            parent,
+
-            // SAFETY: `node` is guaranteed to be a child of `self.module`
+        let scope = Scope::new(parent, node_with_kind, children_start..children_start);
-            node: unsafe { node.to_kind(self.module.clone()) },
-            descendents: 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,
         }
     }
 }

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
+        self.node(db).scope_kind().is_function_like()
-        // 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
-        )
     }
 
     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<FileScopeId>,
+    parent: Option<FileScopeId>,
-    pub(super) node: NodeWithScopeKind,
+    node: NodeWithScopeKind,
-    pub(super) descendents: Range<FileScopeId>,
+    descendents: Range<FileScopeId>,
 }
 
 impl Scope {
-    pub fn parent(self) -> Option<FileScopeId> {
+    pub(super) fn new(
+        parent: Option<FileScopeId>,
+        node: NodeWithScopeKind,
+        descendents: Range<FileScopeId>,
+    ) -> Self {
+        Scope {
+            parent,
+            node,
+            descendents,
+        }
+    }
+
+    pub fn parent(&self) -> Option<FileScopeId> {
         self.parent
     }
 
@@ -195,6 +198,18 @@ impl Scope {
     pub fn kind(&self) -> ScopeKind {
         self.node().scope_kind()
     }
+
+    pub fn descendents(&self) -> Range<FileScopeId> {
+        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 {
+    pub(crate) fn is_eager(self) -> bool {
-        matches!(self, ScopeKind::Comprehension)
+        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<Item = &Symbol> {
+        self.table.symbols()
+    }
+
+    pub(super) fn symbol_id_by_name(&self, name: &str) -> Option<ScopedSymbolId> {
+        self.table.symbol_id_by_name(name)
+    }
+
+    pub(super) fn symbol(&self, symbol_id: impl Into<ScopedSymbolId>) -> &Symbol {
+        self.table.symbol(symbol_id)
+    }
+
     pub(super) fn finish(mut self) -> SymbolTable {
         self.table.shrink_to_fit();
         self.table

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<ScopedSymbolId, SymbolState>,
+
+    /// 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<BindingWithConstraintsIterator<'_, 'db>> {
+        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<ScopedEagerBindingsId, SymbolBindings>;
+
 /// 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<ScopedSymbolId, SymbolState>,
+
+    /// 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,
         }
     }
 }

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
+                        return Symbol::Unbound;
-                    } else {
-                        global_symbol(db, self.file(), symbol_name)
                     }
+
+                    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.)