[ty] fix lazy snapshot sweeping in nested scopes (#19908)

## Summary

This PR closes astral-sh/ty#955.

## Test Plan

New test cases in `narrowing/conditionals/nested.md`.

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md

@@ -240,6 +240,21 @@ def f(x: str | None):
 
     # When there is a reassignment, any narrowing constraints on the place are invalidated in lazy scopes.
     x = None
+
+def f(x: str | None):
+    def _():
+        if x is not None:
+            def closure():
+                reveal_type(x)  # revealed: str | None
+    x = None
+
+def f(x: str | None):
+    class C:
+        def _():
+            if x is not None:
+                def closure():
+                    reveal_type(x)  # revealed: str
+        x = None  # This assignment is not visible in the inner lazy scope, so narrowing is still valid.
 ```
 
 If a variable defined in a private scope is never reassigned, narrowing remains in effect in the
@@ -256,6 +271,12 @@ def f(const: str | None):
             reveal_type(const)  # revealed: str
 
         [reveal_type(const) for _ in range(1)]  # revealed: str
+
+def f(const: str | None):
+    def _():
+        if const is not None:
+            def closure():
+                reveal_type(const)  # revealed: str
 ```
 
 And even if there is an attribute or subscript assignment to the variable, narrowing of the variable

crates/ty_python_semantic/src/semantic_index.rs

@@ -165,7 +165,7 @@ pub(crate) fn attribute_scopes<'db, 's>(
     let index = semantic_index(db, file);
     let class_scope_id = class_body_scope.file_scope_id(db);
 
-    ChildrenIter::new(index, class_scope_id).filter_map(move |(child_scope_id, scope)| {
+    ChildrenIter::new(&index.scopes, class_scope_id).filter_map(move |(child_scope_id, scope)| {
         let (function_scope_id, function_scope) =
             if scope.node().scope_kind() == ScopeKind::TypeParams {
                 // This could be a generic method with a type-params scope.
@@ -372,18 +372,18 @@ impl<'db> SemanticIndex<'db> {
     /// Returns an iterator over the descendent scopes of `scope`.
     #[allow(unused)]
     pub(crate) fn descendent_scopes(&self, scope: FileScopeId) -> DescendantsIter<'_> {
-        DescendantsIter::new(self, scope)
+        DescendantsIter::new(&self.scopes, scope)
     }
 
     /// Returns an iterator over the direct child scopes of `scope`.
     #[allow(unused)]
     pub(crate) fn child_scopes(&self, scope: FileScopeId) -> ChildrenIter<'_> {
-        ChildrenIter::new(self, scope)
+        ChildrenIter::new(&self.scopes, scope)
     }
 
     /// Returns an iterator over all ancestors of `scope`, starting with `scope` itself.
     pub(crate) fn ancestor_scopes(&self, scope: FileScopeId) -> AncestorsIter<'_> {
-        AncestorsIter::new(self, scope)
+        AncestorsIter::new(&self.scopes, scope)
     }
 
     /// Returns an iterator over ancestors of `scope` that are visible for name resolution,
@@ -401,7 +401,7 @@ impl<'db> SemanticIndex<'db> {
     /// ```
     /// The `method` function can see the global scope but not the class scope.
     pub(crate) fn visible_ancestor_scopes(&self, scope: FileScopeId) -> VisibleAncestorsIter<'_> {
-        VisibleAncestorsIter::new(self, scope)
+        VisibleAncestorsIter::new(&self.scopes, scope)
     }
 
     /// Returns the [`definition::Definition`] salsa ingredient(s) for `definition_key`.
@@ -542,9 +542,9 @@ pub(crate) struct AncestorsIter<'a> {
 }
 
 impl<'a> AncestorsIter<'a> {
-    fn new(module_table: &'a SemanticIndex, start: FileScopeId) -> Self {
+    fn new(scopes: &'a IndexSlice<FileScopeId, Scope>, start: FileScopeId) -> Self {
         Self {
-            scopes: &module_table.scopes,
+            scopes,
             next_id: Some(start),
         }
     }
@@ -571,10 +571,10 @@ pub(crate) struct VisibleAncestorsIter<'a> {
 }
 
 impl<'a> VisibleAncestorsIter<'a> {
-    fn new(module_table: &'a SemanticIndex, start: FileScopeId) -> Self {
+    fn new(scopes: &'a IndexSlice<FileScopeId, Scope>, start: FileScopeId) -> Self {
-        let starting_scope = &module_table.scopes[start];
+        let starting_scope = &scopes[start];
         Self {
-            inner: AncestorsIter::new(module_table, start),
+            inner: AncestorsIter::new(scopes, start),
             starting_scope_kind: starting_scope.kind(),
             yielded_count: 0,
         }
@@ -617,9 +617,9 @@ pub(crate) struct DescendantsIter<'a> {
 }
 
 impl<'a> DescendantsIter<'a> {
-    fn new(index: &'a SemanticIndex, scope_id: FileScopeId) -> Self {
+    fn new(scopes: &'a IndexSlice<FileScopeId, Scope>, scope_id: FileScopeId) -> Self {
-        let scope = &index.scopes[scope_id];
+        let scope = &scopes[scope_id];
-        let scopes = &index.scopes[scope.descendants()];
+        let scopes = &scopes[scope.descendants()];
 
         Self {
             next_id: scope_id + 1,
@@ -654,8 +654,8 @@ pub(crate) struct ChildrenIter<'a> {
 }
 
 impl<'a> ChildrenIter<'a> {
-    pub(crate) fn new(module_index: &'a SemanticIndex, parent: FileScopeId) -> Self {
+    pub(crate) fn new(scopes: &'a IndexSlice<FileScopeId, Scope>, parent: FileScopeId) -> Self {
-        let descendants = DescendantsIter::new(module_index, parent);
+        let descendants = DescendantsIter::new(scopes, parent);
 
         Self {
             parent,

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -47,7 +47,7 @@ use crate::semantic_index::symbol::{ScopedSymbolId, Symbol};
 use crate::semantic_index::use_def::{
     EnclosingSnapshotKey, FlowSnapshot, ScopedEnclosingSnapshotId, UseDefMapBuilder,
 };
-use crate::semantic_index::{ExpressionsScopeMap, SemanticIndex};
+use crate::semantic_index::{ExpressionsScopeMap, SemanticIndex, VisibleAncestorsIter};
 use crate::semantic_model::HasTrackedScope;
 use crate::unpack::{EvaluationMode, Unpack, UnpackKind, UnpackPosition, UnpackValue};
 use crate::{Db, Program};
@@ -400,16 +400,28 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         }
     }
 
-    /// Any lazy snapshots of places that have been reassigned or modified are no longer valid, so delete them.
+    /// Any lazy snapshots of places that have been reassigned are no longer valid, so delete them.
     fn sweep_lazy_snapshots(&mut self, popped_scope_id: FileScopeId) {
+        // Retain only snapshots that are either eager
+        //     || (enclosing_scope != popped_scope && popped_scope is not a visible ancestor of enclosing_scope)
+        //     || enclosing_place is not a symbol or not reassigned
+        // <=> remove those that are lazy
+        //     && (enclosing_scope == popped_scope || popped_scope is a visible ancestor of enclosing_scope)
+        //     && enclosing_place is a symbol and reassigned
         self.enclosing_snapshots.retain(|key, _| {
-            let place_table = &self.place_tables[key.enclosing_scope];
+            let popped_place_table = &self.place_tables[popped_scope_id];
             key.nested_laziness.is_eager()
-                || key.enclosing_scope != popped_scope_id
+                || (key.enclosing_scope != popped_scope_id
-                || !key
+                    && VisibleAncestorsIter::new(&self.scopes, key.enclosing_scope)
-                    .enclosing_place
+                        .all(|(ancestor, _)| ancestor != popped_scope_id))
-                    .as_symbol()
+                || !key.enclosing_place.as_symbol().is_some_and(|symbol_id| {
-                    .is_some_and(|symbol_id| place_table.symbol(symbol_id).is_reassigned())
+                    let name = &self.place_tables[key.enclosing_scope]
+                        .symbol(symbol_id)
+                        .name();
+                    popped_place_table.symbol_id(name).is_some_and(|symbol_id| {
+                        popped_place_table.symbol(symbol_id).is_reassigned()
+                    })
+                })
         });
     }