[ty] Narrow tuple expression match subjects

crates/ty_python_core/src/builder.rs

@@ -22,8 +22,8 @@ use smallvec::SmallVec;
 use ty_module_resolver::{ModuleName, resolve_module};
 
 use crate::HasTrackedScope;
-use crate::ast_ids::AstIdsBuilder;
 use crate::ast_ids::node_key::ExpressionNodeKey;
+use crate::ast_ids::{AstIdsBuilder, ScopedUseId};
 use crate::ast_node_ref::AstNodeRef;
 use crate::definition::{
     AnnotatedAssignmentDefinitionNodeRef, AssignmentDefinitionNodeRef,
@@ -42,7 +42,7 @@ use crate::place::{PlaceExpr, PlaceTableBuilder, PossiblyNarrowedPlacesBuilder,
 use crate::predicate::{
     CallableAndCallExpr, ClassPatternKind, PatternPredicate, PatternPredicateKind, Predicate,
     PredicateNode, PredicateOrLiteral, ScopedPredicateId, SequencePatternPredicateKind,
-    StarImportPlaceholderPredicate,
+    StarImportPlaceholderPredicate, SubjectElementPatternPredicate,
 };
 use crate::program::Program;
 use crate::re_exports::exported_names;
@@ -1811,7 +1811,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                         PossiblyNarrowedPlacesBuilder::new(self.db, place_table)
                             .pattern(pattern, module)
                     }
-                    PredicateNode::IsNonTerminalCall(_)
+                    PredicateNode::SubjectElementPattern(_)
+                    | PredicateNode::IsNonTerminalCall(_)
                     | PredicateNode::StarImportPlaceholder(_) => {
                         // These predicates don't narrow any places
                         PossiblyNarrowedPlaces::default()
@@ -2005,6 +2006,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         &mut self,
         subject: Expression<'db>,
         pattern: &ast::Pattern,
+        sequence_subject_targets: &[(ScopedPlaceId, ScopedUseId, ExpressionNodeKey)],
         guard: Option<&ast::Expr>,
         previous_pattern: Option<PatternPredicate<'db>>,
         is_catchall: bool,
@@ -2050,8 +2052,29 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         // predicates are still created normally for proper control flow tracking.
         let predicate_id = if is_catchall {
             ScopedPredicateId::ALWAYS_TRUE
-        } else {
+        } else if sequence_subject_targets.is_empty() {
             self.record_narrowing_constraint(predicate)
+        } else {
+            let predicate_id = self.add_predicate(predicate);
+            for &(place, use_id, target) in sequence_subject_targets {
+                let subject_element_id =
+                    self.add_predicate(PredicateOrLiteral::Predicate(Predicate {
+                        node: PredicateNode::SubjectElementPattern(
+                            SubjectElementPatternPredicate {
+                                pattern: pattern_predicate,
+                                target,
+                            },
+                        ),
+                        is_positive: true,
+                    }));
+                self.current_use_def_map_mut()
+                    .record_narrowing_constraint_for_bindings_at_use(
+                        subject_element_id,
+                        place,
+                        use_id,
+                    );
+            }
+            predicate_id
         };
         (predicate, predicate_id, pattern_predicate)
     }
@@ -3380,6 +3403,37 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     return;
                 }
 
+                // A match subject is evaluated once. Retain the bindings read by each place so
+                // that case predicates constrain those values rather than later rebindings.
+                let places = self.current_place_table();
+                let ast_ids = self.current_ast_ids();
+                let mut sequence_subject_targets =
+                    SmallVec::<[(ScopedPlaceId, ScopedUseId, ExpressionNodeKey); 2]>::new();
+                let mut subject_elements: Vec<&ast::Expr> = match subject.as_ref() {
+                    ast::Expr::List(list) => list.elts.iter().collect(),
+                    ast::Expr::Tuple(tuple) => tuple.elts.iter().collect(),
+                    _ => Vec::new(),
+                };
+                while let Some(element) = subject_elements.pop() {
+                    match element {
+                        ast::Expr::List(list) => subject_elements.extend(&list.elts),
+                        ast::Expr::Tuple(tuple) => subject_elements.extend(&tuple.elts),
+                        _ => {
+                            let Some(target) = PlaceExpr::try_from_expr(element)
+                                .and_then(|place| places.place_id((&place).into()))
+                                .zip(ast_ids.try_use_id(element))
+                            else {
+                                continue;
+                            };
+                            sequence_subject_targets.push((
+                                target.0,
+                                target.1,
+                                ExpressionNodeKey::from(element),
+                            ));
+                        }
+                    }
+                }
+
                 let mut no_case_matched = self.flow_snapshot();
 
                 let has_catchall = cases
@@ -3402,6 +3456,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                         .add_pattern_narrowing_constraint(
                             subject_expr,
                             &case.pattern,
+                            &sequence_subject_targets,
                             case.guard.as_deref(),
                             previous_pattern,
                             is_catchall,

crates/ty_python_core/src/predicate.rs

@@ -11,6 +11,7 @@ use ruff_db::files::File;
 use ruff_index::{FrozenIndexVec, Idx, IndexVec};
 use ruff_python_ast::{Singleton, name::Name};
 
+use crate::ast_ids::ExpressionNodeKey;
 use crate::db::Db;
 use crate::expression::Expression;
 use crate::global_scope;
@@ -130,9 +131,20 @@ pub enum PredicateNode<'db> {
     /// positives.
     IsNonTerminalCall(CallableAndCallExpr<'db>),
     Pattern(PatternPredicate<'db>),
+    SubjectElementPattern(SubjectElementPatternPredicate<'db>),
     StarImportPlaceholder(StarImportPlaceholderPredicate<'db>),
 }
 
+/// A pattern predicate applied to one expression in a sequence-display subject.
+///
+/// The full pattern determines the predicate's truth value, while `target` selects the subject
+/// occurrence whose aligned pattern constraint should be applied to a binding.
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
+pub struct SubjectElementPatternPredicate<'db> {
+    pub pattern: PatternPredicate<'db>,
+    pub target: ExpressionNodeKey,
+}
+
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
 pub enum ClassPatternKind {
     Irrefutable,

crates/ty_python_core/src/use_def.rs

@@ -1380,6 +1380,32 @@ impl<'db> UseDefMapBuilder<'db> {
         self.record_narrowing_constraint_node_for_places(atom, places);
     }
 
+    /// Records a narrowing constraint on the current live bindings that were read by the
+    /// corresponding earlier uses.
+    pub(super) fn record_narrowing_constraint_for_bindings_at_use(
+        &mut self,
+        predicate: ScopedPredicateId,
+        place: ScopedPlaceId,
+        use_id: ScopedUseId,
+    ) {
+        if predicate == ScopedPredicateId::ALWAYS_TRUE
+            || predicate == ScopedPredicateId::ALWAYS_FALSE
+        {
+            return;
+        }
+
+        let constraint = self.narrowing_constraints.add_atom(predicate);
+        let state = match place {
+            ScopedPlaceId::Symbol(symbol) => &mut self.symbol_states[symbol],
+            ScopedPlaceId::Member(member) => &mut self.member_states[member],
+        };
+        state.record_narrowing_constraint_for_bindings_at_use(
+            &mut self.narrowing_constraints,
+            constraint,
+            &self.bindings_by_use[use_id],
+        );
+    }
+
     /// Records a negated narrowing constraint for only the specified places.
     ///
     /// The positive and negative constraints use the same predicate ID. This lets `P or not P`

crates/ty_python_core/src/use_def/place_state.rs

@@ -517,6 +517,24 @@ impl PlaceState {
             .record_narrowing_constraint(narrowing_constraints, constraint);
     }
 
+    /// Add the given constraint to live bindings that were also present at an earlier use.
+    pub(super) fn record_narrowing_constraint_for_bindings_at_use(
+        &mut self,
+        narrowing_constraints: &mut NarrowingConstraintsBuilder,
+        constraint: ScopedNarrowingConstraint,
+        bindings_at_use: &Bindings,
+    ) {
+        for binding in &mut self.bindings.live_bindings {
+            if bindings_at_use
+                .iter()
+                .any(|binding_at_use| binding_at_use.binding() == binding.binding())
+            {
+                binding.narrowing_constraint = narrowing_constraints
+                    .add_and_constraint(binding.narrowing_constraint, constraint);
+            }
+        }
+    }
+
     /// Add given reachability constraint to all live bindings.
     pub(super) fn record_reachability_constraint(
         &mut self,

crates/ty_python_semantic/resources/mdtest/narrow/match.md

@@ -343,6 +343,194 @@ def test_match_value_sequence(value: object) -> None:
             reveal_type(value[0])  # revealed: object
 ```
 
+## Sequence display subjects
+
+A tuple or list display has no place of its own to narrow. A successful sequence pattern instead
+narrows the corresponding narrowable elements. If a multi-element pattern fails, we do not know
+which element failed to match.
+
+```py
+class TupleSubjectA: ...
+class TupleSubjectA1(TupleSubjectA): ...
+class TupleSubjectB: ...
+class TupleSubjectB1(TupleSubjectB): ...
+
+def match_tuple_expression_subject(a: TupleSubjectA, b: TupleSubjectB) -> None:
+    match a, b:
+        case [TupleSubjectA1(), TupleSubjectB1()]:
+            reveal_type(a)  # revealed: TupleSubjectA1
+            reveal_type(b)  # revealed: TupleSubjectB1
+        case _:
+            reveal_type(a)  # revealed: TupleSubjectA
+            reveal_type(b)  # revealed: TupleSubjectB
+
+    reveal_type(a)  # revealed: TupleSubjectA
+    reveal_type(b)  # revealed: TupleSubjectB
+
+def match_list_expression_subject(a: TupleSubjectA, b: TupleSubjectB) -> None:
+    match [a, b]:
+        case [TupleSubjectA1(), TupleSubjectB1()]:
+            reveal_type(a)  # revealed: TupleSubjectA1
+            reveal_type(b)  # revealed: TupleSubjectB1
+```
+
+## Nested sequence display subjects
+
+Element narrowing recurses through nested tuple and list displays. Attributes and subscripts are
+narrowed when they occupy a fixed position. Dictionary displays and starred subject elements do not
+yet have a fixed element-to-pattern correspondence.
+
+```py
+class TupleSubjectA: ...
+class TupleSubjectA1(TupleSubjectA): ...
+class TupleSubjectB: ...
+class TupleSubjectB1(TupleSubjectB): ...
+
+class SequenceSubjectContainer:
+    a: TupleSubjectA
+
+def match_nested_sequence_expression_subject(
+    container: SequenceSubjectContainer,
+    values: list[TupleSubjectB],
+) -> None:
+    match [[container.a], values[0], object()]:
+        case [[TupleSubjectA1()], TupleSubjectB1(), _]:
+            reveal_type(container.a)  # revealed: TupleSubjectA1
+            reveal_type(values[0])  # revealed: TupleSubjectB1
+
+def match_mapping_expression_subject(value: object) -> None:
+    match [{"value": value}]:
+        case [{"value": int()}]:
+            reveal_type(value)  # revealed: object
+
+def match_starred_list_expression_subject(
+    a: TupleSubjectA,
+    values: list[object],
+) -> None:
+    match [a, *values]:
+        case [TupleSubjectA1()]:
+            reveal_type(a)  # revealed: TupleSubjectA
+```
+
+## Sequence pattern forms for display subjects
+
+Element narrowing respects later cases, OR patterns, impossible alternatives, repeated subject
+expressions, and starred sequence patterns.
+
+```py
+class TupleSubjectA: ...
+class TupleSubjectA1(TupleSubjectA): ...
+class TupleSubjectA2(TupleSubjectA): ...
+class TupleSubjectB: ...
+class TupleSubjectB1(TupleSubjectB): ...
+class TupleSubjectB2(TupleSubjectB): ...
+
+def match_tuple_expression_later_case(a: TupleSubjectA, b: TupleSubjectB) -> None:
+    match a, b:
+        case [TupleSubjectA2(), TupleSubjectB2()]:
+            pass
+        case [TupleSubjectA1(), TupleSubjectB1()]:
+            reveal_type(a)  # revealed: TupleSubjectA1
+            reveal_type(b)  # revealed: TupleSubjectB1
+
+def match_tuple_expression_or_pattern(a: TupleSubjectA, b: TupleSubjectB) -> None:
+    match a, b:
+        case [TupleSubjectA1(), TupleSubjectB1()] | [*_]:
+            # The second alternative does not constrain either tuple element.
+            reveal_type(a)  # revealed: TupleSubjectA
+            reveal_type(b)  # revealed: TupleSubjectB
+
+def match_tuple_expression_constrained_or_pattern(
+    a: TupleSubjectA,
+    b: TupleSubjectB,
+) -> None:
+    match a, b:
+        case [TupleSubjectA1(), TupleSubjectB1()] | [TupleSubjectA2(), TupleSubjectB2()]:
+            reveal_type(a)  # revealed: TupleSubjectA1 | TupleSubjectA2
+            reveal_type(b)  # revealed: TupleSubjectB1 | TupleSubjectB2
+
+def match_tuple_expression_or_impossible_alternative(
+    a: TupleSubjectA,
+    b: TupleSubjectB,
+) -> None:
+    match a, b:
+        case [TupleSubjectA1()] | [TupleSubjectA2(), TupleSubjectB1()]:
+            reveal_type(a)  # revealed: TupleSubjectA2
+            reveal_type(b)  # revealed: TupleSubjectB1
+
+def match_repeated_tuple_expression_subject(a: TupleSubjectA) -> None:
+    match a, a:
+        case [TupleSubjectA1(), TupleSubjectA()]:
+            reveal_type(a)  # revealed: TupleSubjectA1
+
+def match_tuple_expression_starred_pattern(
+    a: TupleSubjectA,
+    middle: object,
+    b: TupleSubjectB,
+) -> None:
+    match a, middle, b:
+        case [TupleSubjectA1(), *_, TupleSubjectB1()]:
+            reveal_type(a)  # revealed: TupleSubjectA1
+            reveal_type(middle)  # revealed: object
+            reveal_type(b)  # revealed: TupleSubjectB1
+```
+
+## Subject-time bindings in display subjects
+
+Each element constraint applies to the binding read while that subject element was evaluated. It
+does not constrain a binding introduced by a later subject element, pattern capture, or guard.
+
+```py
+from typing import final
+
+class TupleSubjectA: ...
+class TupleSubjectA1(TupleSubjectA): ...
+class TupleSubjectA2(TupleSubjectA): ...
+class TupleSubjectB: ...
+class TupleSubjectB1(TupleSubjectB): ...
+class ReboundTupleSubject: ...
+
+@final
+class ReboundTupleSubject1(ReboundTupleSubject): ...
+
+@final
+class ReboundTupleSubject2(ReboundTupleSubject): ...
+
+def match_tuple_expression_rebound_subject(a: ReboundTupleSubject) -> None:
+    match a, (a := ReboundTupleSubject2()), a:
+        case [ReboundTupleSubject1(), ReboundTupleSubject2(), ReboundTupleSubject2()]:
+            reveal_type(a)  # revealed: ReboundTupleSubject2
+            1 + "x"  # error: [unsupported-operator]
+
+def match_tuple_expression_multiple_bindings(flag: bool, b: TupleSubjectB) -> None:
+    if flag:
+        a: TupleSubjectA = TupleSubjectA1()
+    else:
+        a = TupleSubjectA2()
+
+    match a, b:
+        case [TupleSubjectA1(), TupleSubjectB1()]:
+            reveal_type(a)  # revealed: TupleSubjectA1
+            reveal_type(b)  # revealed: TupleSubjectB1
+
+def match_tuple_expression_subject_capture(a: TupleSubjectA, b: TupleSubjectB) -> None:
+    match a, b:
+        case [TupleSubjectA1(), a]:
+            reveal_type(a)  # revealed: @Todo(`match` pattern definition types)
+
+def match_tuple_expression_guard_rebinding(
+    a: TupleSubjectA,
+    b: TupleSubjectB,
+    flag: bool,
+) -> None:
+    match a, b:
+        case [TupleSubjectA1(), TupleSubjectB1()] if (a := TupleSubjectA2()) and flag:
+            pass
+        case [TupleSubjectA1(), TupleSubjectB1()]:
+            reveal_type(a)  # revealed: TupleSubjectA1 | TupleSubjectA2
+            reveal_type(b)  # revealed: TupleSubjectB1
+```
+
 ## Value patterns
 
 Value patterns are evaluated by equality, which is overridable. Therefore successfully matching on