[ty] filter out pre-loop bindings from loop headers

Author: oconnor663

crates/ty_python_semantic/src/place.rs

@@ -1208,24 +1208,25 @@ fn loop_header_reachability_impl<'db>(
 
         match use_def.definition(live_binding.binding) {
             DefinitionState::Defined(def) => {
+                debug_assert_ne!(
+                    def, definition,
+                    "loop headers only include bindings from within the loop"
+                );
                 has_defined_bindings = true;
-                if def != definition {
-                    reachable_bindings.insert(ReachableLoopBinding {
-                        definition: def,
-                        narrowing_constraint: live_binding.narrowing_constraint,
-                    });
-                }
+                reachable_bindings.insert(ReachableLoopBinding {
+                    definition: def,
+                    narrowing_constraint: live_binding.narrowing_constraint,
+                });
             }
             // `del` in the loop body is always visible to code after the loop via the
             // normal control flow merge. Updating `deleted_reachability` here is
             // necessary for prior uses in the loop to see it.
             DefinitionState::Deleted => {
                 deleted_reachability = deleted_reachability.or(reachability);
             }
-            // If UNBOUND is visible at loop-back, then it was visible before the loop.
-            // Loop header definitions don't shadow preexisting bindings, so we don't
-            // need to do anything with this.
-            DefinitionState::Undefined => {}
+            DefinitionState::Undefined => {
+                unreachable!("loop headers only include bindings from within the loop")
+            }
         }
     }
 
@@ -1242,7 +1243,7 @@ pub(crate) struct LoopHeaderReachability<'db> {
     /// Whether any reachable loop-back binding is a defined binding.
     pub(crate) has_defined_bindings: bool,
     pub(crate) deleted_reachability: Truthiness,
-    /// Reachable, defined loop-back bindings (excluding the loop header definition itself).
+    /// Reachable loop-back bindings that are not `del`s.
     pub(crate) reachable_bindings: FxIndexSet<ReachableLoopBinding<'db>>,
 }
 

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -48,8 +48,8 @@ use crate::semantic_index::scope::{
 use crate::semantic_index::scope::{Scope, ScopeId, ScopeKind, ScopeLaziness};
 use crate::semantic_index::symbol::{ScopedSymbolId, Symbol};
 use crate::semantic_index::use_def::{
-    EnclosingSnapshotKey, FlowSnapshot, PreviousDefinitions, ScopedEnclosingSnapshotId,
-    UseDefMapBuilder,
+    EnclosingSnapshotKey, FlowSnapshot, PreviousDefinitions, ScopedDefinitionId,
+    ScopedEnclosingSnapshotId, UseDefMapBuilder,
 };
 use crate::semantic_index::{
     ExpressionsScopeMap, LoopHeader, LoopToken, SemanticIndex, VisibleAncestorsIter,
@@ -845,12 +845,13 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
     }
 
     /// Create loop header definitions for all places that are bound within a loop. Return the
-    /// `LoopToken` referenced by those definitions, and the set of bound place IDs.
+    /// `LoopToken` referenced by those definitions, the set of bound place IDs, and the lower
+    /// bound `ScopedDefinitionId` for definitions created within the loop.
     fn synthesize_loop_header_definitions(
         &mut self,
         loop_stmt: LoopStmtRef<'ast>,
         bound_places: Vec<PlaceExpr>,
-    ) -> (LoopToken<'db>, FxHashSet<ScopedPlaceId>) {
+    ) -> (LoopToken<'db>, FxHashSet<ScopedPlaceId>, ScopedDefinitionId) {
         let loop_token = LoopToken::new(self.db);
         let mut bound_place_ids: FxHashSet<ScopedPlaceId> = FxHashSet::default();
         for place_expr in bound_places {
@@ -865,7 +866,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                 self.push_additional_definition(place_id, loop_header_ref);
             }
         }
-        (loop_token, bound_place_ids)
+        let loop_min_definition_id = self.current_use_def_map_mut().next_definition_id();
+        (loop_token, bound_place_ids, loop_min_definition_id)
     }
 
     /// Build a `LoopHeader` that tracks all the variables bound in a loop, which will be visible
@@ -876,13 +878,18 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         &mut self,
         loop_header_places: &FxHashSet<ScopedPlaceId>,
         loop_token: LoopToken<'db>,
+        loop_min_definition_id: ScopedDefinitionId,
     ) {
         let mut loop_header = LoopHeader::new();
         let use_def = self.current_use_def_map_mut();
-        // Collect bindings.
+        // Collect all the bindings within the loop that reached a loop back edge. Use the minimum
+        // definition ID to filter out all the pre-loop bindings. The loop header doesn't shadow
+        // them, so there's no need to duplicate them.
         for place_id in loop_header_places {
             for live_binding in use_def.loop_back_bindings(*place_id) {
-                loop_header.add_binding(*place_id, live_binding);
+                if live_binding.binding >= loop_min_definition_id {
+                    loop_header.add_binding(*place_id, live_binding);
+                }
             }
         }
         // Mark the reachability and narrowing constraints as used.
@@ -2225,8 +2232,10 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
 
                 // Collect all the loop-back bindings (including the `continue` states we just
                 // merged) and populate the `LoopHeader`.
-                if let Some((loop_token, bound_place_ids)) = maybe_loop_header_info {
-                    self.populate_loop_header(&bound_place_ids, loop_token);
+                if let Some((loop_token, bound_place_ids, loop_min_definition_id)) =
+                    maybe_loop_header_info
+                {
+                    self.populate_loop_header(&bound_place_ids, loop_token, loop_min_definition_id);
                 }
 
                 // We execute the `else` branch once the condition evaluates to false. This could
@@ -2333,8 +2342,10 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
 
                 // Collect all the loop-back bindings (including the `continue` states we just
                 // merged) and populate the `LoopHeader`.
-                if let Some((loop_token, bound_place_ids)) = maybe_loop_header_info {
-                    self.populate_loop_header(&bound_place_ids, loop_token);
+                if let Some((loop_token, bound_place_ids, loop_min_definition_id)) =
+                    maybe_loop_header_info
+                {
+                    self.populate_loop_header(&bound_place_ids, loop_token, loop_min_definition_id);
                 }
 
                 // We may execute the `else` clause without ever executing the body, so merge in

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -990,6 +990,10 @@ impl<'db> UseDefMapBuilder<'db> {
         }
     }
 
+    pub(super) fn next_definition_id(&self) -> ScopedDefinitionId {
+        self.all_definitions.next_index()
+    }
+
     pub(super) fn record_binding(
         &mut self,
         place: ScopedPlaceId,

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -5218,11 +5218,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
     /// Infer the type for a loop header definition.
     ///
-    /// The loop header sees all bindings that loop-back, either by reaching the end of the loop
-    /// body or a `continue` statement. This can include bindings from before the loop too, though
-    /// that's technically redundant, since the loop header definition itself doesn't shadow those
-    /// bindings. See `struct LoopHeader` in the semantic index for more on how all this fits
-    /// together.
+    /// The loop header sees all the bindings that originate in the loop and are visible at a
+    /// loop-back edge (either the end of the loop body or a `continue` statement). See `struct
+    /// LoopHeader` in the semantic index for more on how all this fits together.
     fn infer_loop_header_definition(
         &mut self,
         loop_header_kind: &LoopHeaderDefinitionKind<'db>,