[red-knot] Resolve references in eager nested scopes eagerly

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()]
 ```
 

crates/red_knot_python_semantic/resources/mdtest/scopes/eager.md

@@ -0,0 +1,203 @@
+# 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]
+
+    x = 2
+```
+
+## List comprehensions
+
+List comprehensions are evaluated eagerly.
+
+```py
+def _():
+    x = 1
+
+    # revealed: Literal[1]
+    [reveal_type(x) for a in range(0)]
+
+    x = 2
+```
+
+## Set comprehensions
+
+Set comprehensions are evaluated eagerly.
+
+```py
+def _():
+    x = 1
+
+    # revealed: Literal[1]
+    {reveal_type(x) for a in range(0)}
+
+    x = 2
+```
+
+## Dict comprehensions
+
+Dict comprehensions are evaluated eagerly.
+
+```py
+def _():
+    x = 1
+
+    # revealed: Literal[1]
+    {a: reveal_type(x) for a in range(0)}
+
+    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(0))
+
+    x = 2
+```
+
+## 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(0)]
+
+    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(0)]
+
+        class B:
+            # revealed: Literal[1]
+            [reveal_type(x) for a in range(0)]
+
+    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(0)]
+    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(0)]
+
+    x = 2
+```

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,22 @@ 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,
+        enclosing_symbol: ScopedSymbolId,
+        nested_scope: FileScopeId,
+    ) -> Option<BindingWithConstraintsIterator<'_, 'db>> {
+        let key = EagerBindingsKey {
+            enclosing_scope,
+            enclosing_symbol,
+            nested_scope,
+        };
+        let id = self.eager_bindings.get(&key)?;
+        self.use_def_maps[enclosing_scope].eager_bindings(*id)
+    }
 }
 
 pub struct AncestorsIter<'a> {
@@ -312,7 +331,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 +347,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 +397,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/ast_ids.rs

@@ -40,6 +40,7 @@ impl AstIds {
         })
     }
 
+    #[track_caller]
     fn use_id(&self, key: impl Into<ExpressionNodeKey>) -> ScopedUseId {
         self.uses_map[&key.into()]
     }