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

Author: oconnor663

crates/ty_python_semantic/src/place.rs

@@ -1309,10 +1309,9 @@ fn place_from_bindings_impl<'db>(
                                     loop_back.reachability_constraint,
                                 ));
                         }
-                        // 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")
+                        }
                     }
                 }
                 // If all the bindings in the loop are in statically false branches, it might be

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -47,8 +47,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,
@@ -840,12 +840,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 {
@@ -860,7 +861,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
@@ -871,13 +873,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.
@@ -2210,8 +2217,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
@@ -2318,8 +2327,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

@@ -937,6 +937,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

@@ -5009,11 +5009,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>,