Introduce SnapshotCompleteness to retain and utilize incomplete snapshots

Author: mtshiba

crates/ty_python_semantic/resources/mdtest/public_types.md

@@ -214,8 +214,7 @@ def outer() -> None:
     x = 1
 
     def inner() -> None:
-        # TODO: ideally, this should be `Literal[1, 2]`
-        reveal_type(x)  # revealed: None | Literal[1, 2]
+        reveal_type(x)  # revealed: Literal[1, 2]
     inner()
 
     x = 2

crates/ty_python_semantic/src/place.rs

@@ -4,9 +4,10 @@ use crate::dunder_all::dunder_all_names;
 use crate::module_resolver::{KnownModule, file_to_module};
 use crate::semantic_index::definition::{Definition, DefinitionState};
 use crate::semantic_index::place::{PlaceExprRef, ScopedPlaceId};
-use crate::semantic_index::scope::ScopeId;
+use crate::semantic_index::scope::{ScopeId, ScopeLaziness};
 use crate::semantic_index::{
-    BindingWithConstraints, BindingWithConstraintsIterator, DeclarationsIterator, place_table,
+    BindingWithConstraints, BindingWithConstraintsIterator, BoundnessAnalysis,
+    ConsideredDefinitions, DeclarationsIterator, EnclosingSnapshotResult, place_table,
 };
 use crate::semantic_index::{DeclarationWithConstraint, global_scope, use_def_map};
 use crate::types::{
@@ -293,7 +294,7 @@ pub(crate) fn explicit_global_symbol<'db>(
         global_scope(db, file),
         name,
         RequiresExplicitReExport::No,
-        ConsideredDefinitions::AllReachable,
+        ConsideredDefinitions::AllReachable(None),
     )
 }
 
@@ -700,15 +701,24 @@ fn place_by_id<'db>(
 
     let declarations = match considered_definitions {
         ConsideredDefinitions::EndOfScope => use_def.end_of_scope_declarations(place_id),
-        ConsideredDefinitions::AllReachable => use_def.all_reachable_declarations(place_id),
+        ConsideredDefinitions::AllReachable(_) => use_def.all_reachable_declarations(place_id),
     };
 
     let declared = place_from_declarations_impl(db, declarations, requires_explicit_reexport)
         .ignore_conflicting_declarations();
 
     let all_considered_bindings = || match considered_definitions {
         ConsideredDefinitions::EndOfScope => use_def.end_of_scope_bindings(place_id),
-        ConsideredDefinitions::AllReachable => use_def.all_reachable_bindings(place_id),
+        ConsideredDefinitions::AllReachable(None) => use_def.all_reachable_bindings(place_id),
+        ConsideredDefinitions::AllReachable(Some(snapshot)) => {
+            if let EnclosingSnapshotResult::FoundBindings(bindings, _) =
+                use_def.enclosing_snapshot(snapshot, ScopeLaziness::Lazy)
+            {
+                bindings
+            } else {
+                use_def.all_reachable_bindings(place_id)
+            }
+        }
     };
 
     // If a symbol is undeclared, but qualified with `typing.Final`, we use the right-hand side
@@ -763,7 +773,7 @@ fn place_by_id<'db>(
                 // Place is possibly undeclared and (possibly) bound
                 Place::Type(inferred_ty, boundness) => Place::Type(
                     UnionType::from_elements(db, [inferred_ty, declared_ty]),
-                    if boundness_analysis == BoundnessAnalysis::AssumeBound {
+                    if boundness_analysis.is_assume_bound() {
                         Boundness::Bound
                     } else {
                         boundness
@@ -782,7 +792,7 @@ fn place_by_id<'db>(
             let boundness_analysis = bindings.boundness_analysis;
             let mut inferred = place_from_bindings_impl(db, bindings, requires_explicit_reexport);
 
-            if boundness_analysis == BoundnessAnalysis::AssumeBound {
+            if boundness_analysis.is_assume_bound() {
                 if let Place::Type(ty, Boundness::PossiblyUnbound) = inferred {
                     inferred = Place::Type(ty, Boundness::Bound);
                 }
@@ -1037,7 +1047,7 @@ fn place_from_bindings_impl<'db>(
         };
 
         let boundness = match boundness_analysis {
-            BoundnessAnalysis::AssumeBound => Boundness::Bound,
+            BoundnessAnalysis::AssumeBound(_) => Boundness::Bound,
             BoundnessAnalysis::BasedOnUnboundVisibility => match unbound_visibility() {
                 Some(Truthiness::AlwaysTrue) => {
                     unreachable!(
@@ -1245,7 +1255,7 @@ fn place_from_declarations_impl<'db>(
         };
 
         let boundness = match boundness_analysis {
-            BoundnessAnalysis::AssumeBound => {
+            BoundnessAnalysis::AssumeBound(_) => {
                 if all_declarations_definitely_reachable {
                     Boundness::Bound
                 } else {
@@ -1458,48 +1468,6 @@ impl RequiresExplicitReExport {
     }
 }
 
-/// Specifies which definitions should be considered when looking up a place.
-///
-/// In the example below, the `EndOfScope` variant would consider the `x = 2` and `x = 3` definitions,
-/// while the `AllReachable` variant would also consider the `x = 1` definition.
-/// ```py
-/// def _():
-///     x = 1
-///
-///     x = 2
-///
-///     if flag():
-///         x = 3
-/// ```
-#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, salsa::Update)]
-pub(crate) enum ConsideredDefinitions {
-    /// Consider only the definitions that are "live" at the end of the scope, i.e. those
-    /// that have not been shadowed or deleted.
-    EndOfScope,
-    /// Consider all definitions that are reachable from the start of the scope.
-    AllReachable,
-}
-
-/// Specifies how the boundness of a place should be determined.
-#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, salsa::Update)]
-pub(crate) enum BoundnessAnalysis {
-    /// The place is always considered bound.
-    AssumeBound,
-    /// The boundness of the place is determined based on the visibility of the implicit
-    /// `unbound` binding. In the example below, when analyzing the visibility of the
-    /// `x = <unbound>` binding from the position of the end of the scope, it would be
-    /// `Truthiness::Ambiguous`, because it could either be visible or not, depending on the
-    /// `flag()` return value. This would result in a `Boundness::PossiblyUnbound` for `x`.
-    ///
-    /// ```py
-    /// x = <unbound>
-    ///
-    /// if flag():
-    ///     x = 1
-    /// ```
-    BasedOnUnboundVisibility,
-}
-
 /// Computes a possibly-widened type `Unknown | T_inferred` from the inferred type `T_inferred`
 /// of a symbol, unless the type is a known-instance type (e.g. `typing.Any`) or the symbol is
 /// considered non-modifiable (e.g. when the symbol is `@Final`). We need this for public uses

crates/ty_python_semantic/src/semantic_index.rs

@@ -26,6 +26,9 @@ use crate::semantic_index::scope::{
     NodeWithScopeKey, NodeWithScopeRef, Scope, ScopeId, ScopeKind, ScopeLaziness,
 };
 use crate::semantic_index::symbol::ScopedSymbolId;
+pub(crate) use crate::semantic_index::use_def::{
+    BoundnessAnalysis, ConsideredDefinitions, SnapshotCompleteness,
+};
 use crate::semantic_index::use_def::{EnclosingSnapshotKey, ScopedEnclosingSnapshotId, UseDefMap};
 use crate::semantic_model::HasTrackedScope;
 
@@ -190,7 +193,10 @@ pub(crate) fn global_scope(db: &dyn Db, file: File) -> ScopeId<'_> {
 
 pub(crate) enum EnclosingSnapshotResult<'map, 'db> {
     FoundConstraint(ScopedNarrowingConstraint),
-    FoundBindings(BindingWithConstraintsIterator<'map, 'db>),
+    FoundBindings(
+        BindingWithConstraintsIterator<'map, 'db>,
+        SnapshotCompleteness,
+    ),
     NotFound,
     NoLongerInEagerContext,
 }

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -400,8 +400,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         }
     }
 
-    /// Any lazy snapshots of the place that have been reassigned are no longer valid, so delete them.
-    /// If there is no reassignment, the state of the symbol just before the closure definition is recorded as a lazy snapshot.
+    /// Any lazy snapshots of the place that have been reassigned are no longer valid on their own, so mark as incomplete.
+    /// If there is no reassignment, the state of the symbol just before the closure definition is recorded as a complete lazy snapshot.
     /// We have not yet tracked where the closure is actually referenced, but we can assume that the closure becomes referenceable after the closure is defined.
     /// ```py
     /// def outer() -> None:
@@ -416,44 +416,44 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
     ///         reveal_type(x)  # revealed: None | Literal[1]
     ///
     ///     # Reassignment of `x` after the definition of `inner2`.
-    ///     # Sweep lazy snapshots of `x` for `inner2`.
+    ///     # lazy snapshots of `x` for `inner2` is incomplete.
     ///     x = 1
     ///
     ///     def inner() -> None:
     ///         # In this scope, `x = None` appears as being shadowed by `x = 1`.
     ///         reveal_type(x)  # revealed: Literal[1]
     ///
-    ///     # No reassignment of `x` after the definition of `inner`, so we can safely record a lazy snapshot for `inner`.
+    ///     # No reassignment of `x` after the definition of `inner`, so we can safely use a complete lazy snapshot for `inner`.
     ///     inner()
     ///     inner2()
     /// ```
-    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, snapshot_id| {
-            let popped_place_table = &self.place_tables[popped_scope_id];
-            key.nested_laziness.is_eager()
-                || (key.enclosing_scope != popped_scope_id
-                    && VisibleAncestorsIter::new(&self.scopes, key.enclosing_scope)
-                        .all(|(ancestor, _)| ancestor != popped_scope_id))
-                || !key.enclosing_place.as_symbol().is_some_and(|symbol_id| {
+    fn check_lazy_snapshots(&mut self, popped_scope_id: FileScopeId) {
+        for (key, snapshot_id) in &self.enclosing_snapshots {
+            if let Some(enclosing_symbol) = key.enclosing_place.as_symbol() {
+                if key.nested_laziness.is_lazy()
+                    && (key.enclosing_scope == popped_scope_id
+                        || VisibleAncestorsIter::new(&self.scopes, key.enclosing_scope)
+                            .any(|(ancestor, _)| ancestor == popped_scope_id))
+                {
                     let name = &self.place_tables[key.enclosing_scope]
-                        .symbol(symbol_id)
+                        .symbol(enclosing_symbol)
                         .name();
-                    popped_place_table.symbol_id(name).is_some_and(|symbol_id| {
-                        popped_place_table.symbol(symbol_id).is_reassigned()
+                    let popped_place_table = &self.place_tables[popped_scope_id];
+                    if let Some(symbol_id) = popped_place_table.symbol_id(name) {
+                        if popped_place_table.symbol(symbol_id).is_reassigned()
                             && !self.use_def_maps[popped_scope_id].bindings_are_unchanged(
                                 &self.use_def_maps[key.enclosing_scope],
                                 *snapshot_id,
                                 symbol_id,
                             )
-                    })
-                })
-        });
+                        {
+                            self.use_def_maps[key.enclosing_scope]
+                                .mark_snapshot_incomplete(*snapshot_id, enclosing_symbol);
+                        }
+                    }
+                }
+            }
+        }
     }
 
     fn sweep_nonlocal_lazy_snapshots(&mut self) {
@@ -495,7 +495,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             .pop()
             .expect("Root scope should be present");
 
-        self.sweep_lazy_snapshots(popped_scope_id);
+        self.check_lazy_snapshots(popped_scope_id);
 
         let children_end = self.scopes.next_index();
 
@@ -781,7 +781,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
     /// Records that all remaining statements in the current block are unreachable.
     fn mark_unreachable(&mut self) {
         self.current_use_def_map_mut().mark_unreachable();
-        self.sweep_lazy_snapshots(self.current_scope());
+        self.check_lazy_snapshots(self.current_scope());
     }
 
     /// Records a reachability constraint that always evaluates to "ambiguous".

crates/ty_python_semantic/src/semantic_index/scope.rs

@@ -195,6 +195,10 @@ impl ScopeLaziness {
     pub(crate) const fn is_eager(self) -> bool {
         matches!(self, ScopeLaziness::Eager)
     }
+
+    pub(crate) const fn is_lazy(self) -> bool {
+        matches!(self, ScopeLaziness::Lazy)
+    }
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]