diff --git "a/crates/ty_python_semantic/resources/mdtest/snapshots/unpacking.md_-_Unpacking_-_Too_few_values_to_un\342\200\246_(cef19e6b2b58e6a3).snap" "b/crates/ty_python_semantic/resources/mdtest/snapshots/unpacking.md_-_Unpacking_-_Too_few_values_to_un\342\200\246_(cef19e6b2b58e6a3).snap" index b20add36c4dc5..c255da479bee2 100644 --- "a/crates/ty_python_semantic/resources/mdtest/snapshots/unpacking.md_-_Unpacking_-_Too_few_values_to_un\342\200\246_(cef19e6b2b58e6a3).snap" +++ "b/crates/ty_python_semantic/resources/mdtest/snapshots/unpacking.md_-_Unpacking_-_Too_few_values_to_un\342\200\246_(cef19e6b2b58e6a3).snap" @@ -24,7 +24,7 @@ error[invalid-assignment]: Not enough values to unpack 1 | [a, *b, c, d] = (1, 2) # error: [invalid-assignment] | ^^^^^^^^^^^^^ ------ Got 2 | | - | Expected 3 or more + | Expected at least 3 | info: rule `invalid-assignment` is enabled by default diff --git a/crates/ty_python_semantic/resources/mdtest/unpacking.md b/crates/ty_python_semantic/resources/mdtest/unpacking.md index 1628610800397..02d4952780841 100644 --- a/crates/ty_python_semantic/resources/mdtest/unpacking.md +++ b/crates/ty_python_semantic/resources/mdtest/unpacking.md @@ -106,7 +106,7 @@ reveal_type(d) # revealed: Literal[5] ### Starred expression (1) ```py -# error: [invalid-assignment] "Not enough values to unpack: Expected 3 or more" +# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3" [a, *b, c, d] = (1, 2) reveal_type(a) # revealed: Unknown reveal_type(b) # revealed: list[Unknown] @@ -119,7 +119,7 @@ reveal_type(d) # revealed: Unknown ```py [a, *b, c] = (1, 2) reveal_type(a) # revealed: Literal[1] -reveal_type(b) # revealed: list[Unknown] +reveal_type(b) # revealed: list[Never] reveal_type(c) # revealed: Literal[2] ``` @@ -154,7 +154,7 @@ reveal_type(c) # revealed: list[Literal[3, 4]] ### Starred expression (6) ```py -# error: [invalid-assignment] "Not enough values to unpack: Expected 5 or more" +# error: [invalid-assignment] "Not enough values to unpack: Expected at least 5" (a, b, c, *d, e, f) = (1,) reveal_type(a) # revealed: Unknown reveal_type(b) # revealed: Unknown @@ -258,6 +258,155 @@ def _(value: list[int]): reveal_type(c) # revealed: int ``` +## Homogeneous tuples + +### Simple unpacking + +```py +def _(value: tuple[int, ...]): + a, b = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: int +``` + +### Nested unpacking + +```py +def _(value: tuple[tuple[int, ...], ...]): + a, (b, c) = value + reveal_type(a) # revealed: tuple[int, ...] + reveal_type(b) # revealed: int + reveal_type(c) # revealed: int +``` + +### Invalid nested unpacking + +```py +def _(value: tuple[int, ...]): + # error: [not-iterable] "Object of type `int` is not iterable" + a, (b, c) = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: Unknown + reveal_type(c) # revealed: Unknown +``` + +### Starred expression + +```py +def _(value: tuple[int, ...]): + a, *b, c = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: list[int] + reveal_type(c) # revealed: int +``` + +## Mixed tuples + +```toml +[environment] +python-version = "3.11" +``` + +### Simple unpacking (1) + +```py +def _(value: tuple[int, *tuple[str, ...]]): + a, b = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: str +``` + +### Simple unpacking (2) + +```py +def _(value: tuple[int, int, *tuple[str, ...]]): + a, b = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: int +``` + +### Simple unpacking (3) + +```py +def _(value: tuple[int, *tuple[str, ...], int]): + a, b, c = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: str + reveal_type(c) # revealed: int +``` + +### Invalid unpacked + +```py +def _(value: tuple[int, int, int, *tuple[str, ...]]): + # error: [invalid-assignment] "Too many values to unpack: Expected 2" + a, b = value + reveal_type(a) # revealed: Unknown + reveal_type(b) # revealed: Unknown +``` + +### Nested unpacking + +```py +def _(value: tuple[str, *tuple[tuple[int, ...], ...]]): + a, (b, c) = value + reveal_type(a) # revealed: str + reveal_type(b) # revealed: int + reveal_type(c) # revealed: int +``` + +### Invalid nested unpacking + +```py +def _(value: tuple[str, *tuple[int, ...]]): + # error: [not-iterable] "Object of type `int` is not iterable" + a, (b, c) = value + reveal_type(a) # revealed: str + reveal_type(b) # revealed: Unknown + reveal_type(c) # revealed: Unknown +``` + +### Starred expression (1) + +```py +def _(value: tuple[int, *tuple[str, ...]]): + a, *b, c = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: list[str] + reveal_type(c) # revealed: str +``` + +### Starred expression (2) + +```py +def _(value: tuple[int, *tuple[str, ...], int]): + a, *b, c = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: list[str] + reveal_type(c) # revealed: int +``` + +### Starred expression (3) + +```py +def _(value: tuple[int, *tuple[str, ...], int]): + a, *b, c, d = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: list[str] + reveal_type(c) # revealed: str + reveal_type(d) # revealed: int +``` + +### Starred expression (4) + +```py +def _(value: tuple[int, int, *tuple[str, ...], int]): + a, *b, c = value + reveal_type(a) # revealed: int + reveal_type(b) # revealed: list[int | str] + reveal_type(c) # revealed: int +``` + ## String ### Simple unpacking @@ -290,7 +439,7 @@ reveal_type(b) # revealed: Unknown ### Starred expression (1) ```py -# error: [invalid-assignment] "Not enough values to unpack: Expected 3 or more" +# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3" (a, *b, c, d) = "ab" reveal_type(a) # revealed: Unknown reveal_type(b) # revealed: list[Unknown] @@ -299,7 +448,7 @@ reveal_type(d) # revealed: Unknown ``` ```py -# error: [invalid-assignment] "Not enough values to unpack: Expected 3 or more" +# error: [invalid-assignment] "Not enough values to unpack: Expected at least 3" (a, b, *c, d) = "a" reveal_type(a) # revealed: Unknown reveal_type(b) # revealed: Unknown @@ -312,7 +461,7 @@ reveal_type(d) # revealed: Unknown ```py (a, *b, c) = "ab" reveal_type(a) # revealed: LiteralString -reveal_type(b) # revealed: list[Unknown] +reveal_type(b) # revealed: list[Never] reveal_type(c) # revealed: LiteralString ``` diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index d62a76f8eb6fa..2ae5e0635a91d 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -726,7 +726,7 @@ impl<'db> Type<'db> { .map(|ty| ty.materialize(db, variance.flip())), ) .build(), - Type::Tuple(tuple_type) => Type::tuple(db, tuple_type.materialize(db, variance)), + Type::Tuple(tuple_type) => Type::tuple(tuple_type.materialize(db, variance)), Type::TypeVar(type_var) => Type::TypeVar(type_var.materialize(db, variance)), Type::TypeIs(type_is) => { type_is.with_type(db, type_is.return_type(db).materialize(db, variance)) @@ -1141,7 +1141,7 @@ impl<'db> Type<'db> { match self { Type::Union(union) => Type::Union(union.normalized(db)), Type::Intersection(intersection) => Type::Intersection(intersection.normalized(db)), - Type::Tuple(tuple) => Type::tuple(db, tuple.normalized(db)), + Type::Tuple(tuple) => Type::tuple(tuple.normalized(db)), Type::Callable(callable) => Type::Callable(callable.normalized(db)), Type::ProtocolInstance(protocol) => protocol.normalized(db), Type::NominalInstance(instance) => Type::NominalInstance(instance.normalized(db)), @@ -3458,7 +3458,7 @@ impl<'db> Type<'db> { Type::BooleanLiteral(bool) => Truthiness::from(*bool), Type::StringLiteral(str) => Truthiness::from(!str.value(db).is_empty()), Type::BytesLiteral(bytes) => Truthiness::from(!bytes.value(db).is_empty()), - Type::Tuple(tuple) => match tuple.tuple(db).size_hint() { + Type::Tuple(tuple) => match tuple.tuple(db).len().size_hint() { // The tuple type is AlwaysFalse if it contains only the empty tuple (_, Some(0)) => Truthiness::AlwaysFalse, // The tuple type is AlwaysTrue if its inhabitants must always have length >=1 @@ -4312,7 +4312,7 @@ impl<'db> Type<'db> { let mut parameter = Parameter::positional_only(Some(Name::new_static("iterable"))) .with_annotated_type(instantiated); - if matches!(spec.size_hint().1, Some(0)) { + if matches!(spec.len().maximum(), Some(0)) { parameter = parameter.with_default_type(TupleType::empty(db)); } Parameters::new([parameter]) @@ -5350,7 +5350,7 @@ impl<'db> Type<'db> { } builder.build() } - Type::Tuple(tuple) => Type::Tuple(tuple.apply_type_mapping(db, type_mapping)), + Type::Tuple(tuple) => Type::tuple(tuple.apply_type_mapping(db, type_mapping)), Type::TypeIs(type_is) => type_is.with_type(db, type_is.return_type(db).apply_type_mapping(db, type_mapping)), diff --git a/crates/ty_python_semantic/src/types/builder.rs b/crates/ty_python_semantic/src/types/builder.rs index 2186310b03c28..70876773225d0 100644 --- a/crates/ty_python_semantic/src/types/builder.rs +++ b/crates/ty_python_semantic/src/types/builder.rs @@ -444,6 +444,10 @@ impl<'db> UnionBuilder<'db> { } pub(crate) fn build(self) -> Type<'db> { + self.try_build().unwrap_or(Type::Never) + } + + pub(crate) fn try_build(self) -> Option> { let mut types = vec![]; for element in self.elements { match element { @@ -460,9 +464,12 @@ impl<'db> UnionBuilder<'db> { } } match types.len() { - 0 => Type::Never, - 1 => types[0], - _ => Type::Union(UnionType::new(self.db, types.into_boxed_slice())), + 0 => None, + 1 => Some(types[0]), + _ => Some(Type::Union(UnionType::new( + self.db, + types.into_boxed_slice(), + ))), } } } diff --git a/crates/ty_python_semantic/src/types/call/arguments.rs b/crates/ty_python_semantic/src/types/call/arguments.rs index 8b749f1f4281d..f563eee7bfab3 100644 --- a/crates/ty_python_semantic/src/types/call/arguments.rs +++ b/crates/ty_python_semantic/src/types/call/arguments.rs @@ -221,10 +221,10 @@ fn expand_type<'db>(db: &'db dyn Db, ty: Type<'db>) -> Option>> { let expanded = tuple .all_elements() .map(|element| { - if let Some(expanded) = expand_type(db, element) { + if let Some(expanded) = expand_type(db, *element) { Either::Left(expanded.into_iter()) } else { - Either::Right(std::iter::once(element)) + Either::Right(std::iter::once(*element)) } }) .multi_cartesian_product() diff --git a/crates/ty_python_semantic/src/types/generics.rs b/crates/ty_python_semantic/src/types/generics.rs index 983b5caf7d1d2..1481c51e6be5b 100644 --- a/crates/ty_python_semantic/src/types/generics.rs +++ b/crates/ty_python_semantic/src/types/generics.rs @@ -286,9 +286,13 @@ impl<'db> Specialization<'db> { return tuple; } if let [element_type] = self.types(db) { - return TupleType::new(db, TupleSpec::homogeneous(*element_type)).tuple(db); + if let Some(tuple) = TupleType::new(db, TupleSpec::homogeneous(*element_type)) { + return tuple.tuple(db); + } } - TupleType::new(db, TupleSpec::homogeneous(Type::unknown())).tuple(db) + TupleType::new(db, TupleSpec::homogeneous(Type::unknown())) + .expect("tuple[Unknown, ...] should never contain Never") + .tuple(db) } /// Returns the type that a typevar is mapped to, or None if the typevar isn't part of this @@ -330,7 +334,7 @@ impl<'db> Specialization<'db> { .collect(); let tuple_inner = self .tuple_inner(db) - .map(|tuple| tuple.apply_type_mapping(db, type_mapping)); + .and_then(|tuple| tuple.apply_type_mapping(db, type_mapping)); Specialization::new(db, self.generic_context(db), types, tuple_inner) } @@ -374,7 +378,7 @@ impl<'db> Specialization<'db> { pub(crate) fn normalized(self, db: &'db dyn Db) -> Self { let types: Box<[_]> = self.types(db).iter().map(|ty| ty.normalized(db)).collect(); - let tuple_inner = self.tuple_inner(db).map(|tuple| tuple.normalized(db)); + let tuple_inner = self.tuple_inner(db).and_then(|tuple| tuple.normalized(db)); Self::new(db, self.generic_context(db), types, tuple_inner) } @@ -394,7 +398,7 @@ impl<'db> Specialization<'db> { vartype.materialize(db, variance) }) .collect(); - let tuple_inner = self.tuple_inner(db).map(|tuple| { + let tuple_inner = self.tuple_inner(db).and_then(|tuple| { // Tuples are immutable, so tuple element types are always in covariant position. tuple.materialize(db, variance) }); @@ -637,7 +641,7 @@ impl<'db> SpecializationBuilder<'db> { (TupleSpec::Fixed(formal_tuple), TupleSpec::Fixed(actual_tuple)) => { if formal_tuple.len() == actual_tuple.len() { for (formal_element, actual_element) in formal_tuple.elements().zip(actual_tuple.elements()) { - self.infer(formal_element, actual_element)?; + self.infer(*formal_element, *actual_element)?; } } } diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index a8301e04f602c..c75446f9a2bc7 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -2835,7 +2835,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // it will actually be the type of the generic parameters to `BaseExceptionGroup` or `ExceptionGroup`. let symbol_ty = if let Type::Tuple(tuple) = node_ty { let mut builder = UnionBuilder::new(self.db()); - for element in tuple.tuple(self.db()).all_elements() { + for element in tuple.tuple(self.db()).all_elements().copied() { builder = builder.add( if element.is_assignable_to(self.db(), type_base_exception) { element.to_instance(self.db()).expect( @@ -3701,7 +3701,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ast::Expr::List(ast::ExprList { elts, .. }) | ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => { let mut assigned_tys = match assigned_ty { - Some(Type::Tuple(tuple)) => Either::Left(tuple.tuple(self.db()).all_elements()), + Some(Type::Tuple(tuple)) => { + Either::Left(tuple.tuple(self.db()).all_elements().copied()) + } Some(_) | None => Either::Right(std::iter::empty()), }; @@ -6485,13 +6487,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { op, ), - (Type::Tuple(lhs), Type::Tuple(rhs), ast::Operator::Add) => Some(Type::tuple( - self.db(), - TupleType::new( + (Type::Tuple(lhs), Type::Tuple(rhs), ast::Operator::Add) => { + Some(Type::tuple(TupleType::new( self.db(), lhs.tuple(self.db()).concat(self.db(), rhs.tuple(self.db())), - ), - )), + ))) + } // We've handled all of the special cases that we support for literals, so we need to // fall back on looking for dunder methods on one of the operand types. @@ -6948,14 +6949,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // tuples. // // Ref: https://github.com/astral-sh/ruff/pull/18251#discussion_r2115909311 - let (minimum_length, _) = tuple.tuple(self.db()).size_hint(); + let (minimum_length, _) = tuple.tuple(self.db()).len().size_hint(); if minimum_length > 1 << 12 { return None; } let mut definitely_true = false; let mut definitely_false = true; - for element in tuple.tuple(self.db()).all_elements() { + for element in tuple.tuple(self.db()).all_elements().copied() { if element.is_string_literal() { if literal == element { definitely_true = true; @@ -7238,7 +7239,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut any_eq = false; let mut any_ambiguous = false; - for ty in rhs_tuple.all_elements() { + for ty in rhs_tuple.all_elements().copied() { let eq_result = self.infer_binary_type_comparison( Type::Tuple(lhs), ast::CmpOp::Eq, @@ -7450,8 +7451,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { return Ok(Type::unknown()); }; - let left_iter = left.elements(); - let right_iter = right.elements(); + let left_iter = left.elements().copied(); + let right_iter = right.elements().copied(); let mut builder = UnionBuilder::new(self.db()); @@ -7695,7 +7696,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { "tuple", value_node.into(), value_ty, - tuple.display_minimum_length(), + tuple.len().display_minimum(), int, ); Type::unknown() @@ -8856,7 +8857,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let ty = if return_todo { todo_type!("PEP 646") } else { - Type::tuple(self.db(), TupleType::new(self.db(), element_types)) + Type::tuple(TupleType::new(self.db(), element_types)) }; // Here, we store the type for the inner `int, str` tuple-expression, diff --git a/crates/ty_python_semantic/src/types/instance.rs b/crates/ty_python_semantic/src/types/instance.rs index 957bb913874b2..bc609339e882e 100644 --- a/crates/ty_python_semantic/src/types/instance.rs +++ b/crates/ty_python_semantic/src/types/instance.rs @@ -19,7 +19,7 @@ impl<'db> Type<'db> { TupleType::homogeneous(db, Type::unknown()) } (ClassType::Generic(alias), Some(KnownClass::Tuple)) => { - Self::tuple(db, TupleType::new(db, alias.specialization(db).tuple(db))) + Self::tuple(TupleType::new(db, alias.specialization(db).tuple(db))) } _ if class.class_literal(db).0.is_protocol(db) => { Self::ProtocolInstance(ProtocolInstanceType::from_class(class)) diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index dc3dbd68ac2b9..e8f2d207add91 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -184,6 +184,7 @@ impl ClassInfoConstraintFunction { tuple .tuple(db) .all_elements() + .copied() .map(|element| self.generate_constraint(db, element)), ), Type::ClassLiteral(class_literal) => { diff --git a/crates/ty_python_semantic/src/types/tuple.rs b/crates/ty_python_semantic/src/types/tuple.rs index 48d6a7d536d65..9707dbdac7d09 100644 --- a/crates/ty_python_semantic/src/types/tuple.rs +++ b/crates/ty_python_semantic/src/types/tuple.rs @@ -16,17 +16,71 @@ //! that adds that "collapse `Never`" behavior, whereas [`TupleSpec`] allows you to add any element //! types, including `Never`.) +use std::borrow::Borrow; +use std::cmp::Ordering; +use std::hash::Hash; + use itertools::{Either, EitherOrBoth, Itertools}; use crate::types::class::{ClassType, KnownClass}; -use crate::types::{Type, TypeMapping, TypeRelation, TypeVarInstance, TypeVarVariance, UnionType}; +use crate::types::{ + Type, TypeMapping, TypeRelation, TypeVarInstance, TypeVarVariance, UnionBuilder, UnionType, +}; use crate::util::subscript::{Nth, OutOfBoundsError, PyIndex, PySlice, StepSizeZeroError}; use crate::{Db, FxOrderSet}; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum TupleLength { + Fixed(usize), + Variable(usize, usize), +} + +impl TupleLength { + /// Returns the minimum and maximum length of this tuple. (The maximum length will be `None` + /// for a tuple with a variable-length portion.) + pub(crate) fn size_hint(self) -> (usize, Option) { + match self { + TupleLength::Fixed(len) => (len, Some(len)), + TupleLength::Variable(prefix, suffix) => (prefix + suffix, None), + } + } + + /// Returns the minimum length of this tuple. + pub(crate) fn minimum(self) -> usize { + match self { + TupleLength::Fixed(len) => len, + TupleLength::Variable(prefix, suffix) => prefix + suffix, + } + } + + /// Returns the maximum length of this tuple, if any. + pub(crate) fn maximum(self) -> Option { + match self { + TupleLength::Fixed(len) => Some(len), + TupleLength::Variable(_, _) => None, + } + } + + pub(crate) fn display_minimum(self) -> String { + let minimum_length = self.minimum(); + match self { + TupleLength::Fixed(_) => minimum_length.to_string(), + TupleLength::Variable(_, _) => format!("at least {minimum_length}"), + } + } + + pub(crate) fn display_maximum(self) -> String { + match self.maximum() { + Some(maximum) => maximum.to_string(), + None => "unlimited".to_string(), + } + } +} + /// # Ordering /// Ordering is based on the tuple's salsa-assigned id and not on its elements. /// The id may change between runs, or when the tuple was garbage collected and recreated. -#[salsa::interned(debug)] +#[salsa::interned(debug, constructor=new_internal)] #[derive(PartialOrd, Ord)] pub struct TupleType<'db> { #[returns(ref)] @@ -37,35 +91,53 @@ pub struct TupleType<'db> { impl get_size2::GetSize for TupleType<'_> {} impl<'db> Type<'db> { - pub(crate) fn tuple(db: &'db dyn Db, tuple: TupleType<'db>) -> Self { - // If a fixed-length (i.e., mandatory) element of the tuple is `Never`, then it's not - // possible to instantiate the tuple as a whole. - if tuple.tuple(db).fixed_elements().any(|ty| ty.is_never()) { + pub(crate) fn tuple(tuple: Option>) -> Self { + let Some(tuple) = tuple else { return Type::Never; - } + }; Self::Tuple(tuple) } } impl<'db> TupleType<'db> { + pub(crate) fn new(db: &'db dyn Db, tuple_key: T) -> Option + where + T: Borrow> + Hash + salsa::plumbing::interned::Lookup>, + TupleSpec<'db>: salsa::plumbing::interned::HashEqLike, + { + // If a fixed-length (i.e., mandatory) element of the tuple is `Never`, then it's not + // possible to instantiate the tuple as a whole. + let tuple = tuple_key.borrow(); + if tuple.fixed_elements().any(Type::is_never) { + return None; + } + + // If the variable-length portion is Never, it can only be instantiated with zero elements. + // That means this isn't a variable-length tuple after all! + if let TupleSpec::Variable(tuple) = tuple { + if tuple.variable.is_never() { + let tuple = TupleSpec::Fixed(FixedLengthTuple::from_elements( + tuple.prefix.iter().chain(&tuple.suffix).copied(), + )); + return Some(TupleType::new_internal::<_, TupleSpec<'db>>(db, tuple)); + } + } + + Some(TupleType::new_internal(db, tuple_key)) + } + pub(crate) fn empty(db: &'db dyn Db) -> Type<'db> { - Type::tuple( + Type::tuple(TupleType::new( db, - TupleType::new(db, TupleSpec::from(FixedLengthTupleSpec::empty())), - ) + TupleSpec::from(FixedLengthTuple::empty()), + )) } pub(crate) fn from_elements( db: &'db dyn Db, types: impl IntoIterator>, ) -> Type<'db> { - Type::tuple( - db, - TupleType::new( - db, - TupleSpec::from(FixedLengthTupleSpec::from_elements(types)), - ), - ) + Type::tuple(TupleType::new(db, TupleSpec::from_elements(types))) } #[cfg(test)] @@ -75,14 +147,14 @@ impl<'db> TupleType<'db> { variable: Type<'db>, suffix: impl IntoIterator>, ) -> Type<'db> { - Type::tuple( + Type::tuple(TupleType::new( db, - TupleType::new(db, VariableLengthTupleSpec::mixed(prefix, variable, suffix)), - ) + VariableLengthTuple::mixed(prefix, variable, suffix), + )) } pub(crate) fn homogeneous(db: &'db dyn Db, element: Type<'db>) -> Type<'db> { - Type::tuple(db, TupleType::new(db, TupleSpec::homogeneous(element))) + Type::tuple(TupleType::new(db, TupleSpec::homogeneous(element))) } pub(crate) fn to_class_type(self, db: &'db dyn Db) -> Option> { @@ -102,11 +174,11 @@ impl<'db> TupleType<'db> { /// /// See [`Type::normalized`] for more details. #[must_use] - pub(crate) fn normalized(self, db: &'db dyn Db) -> Self { + pub(crate) fn normalized(self, db: &'db dyn Db) -> Option { TupleType::new(db, self.tuple(db).normalized(db)) } - pub(crate) fn materialize(self, db: &'db dyn Db, variance: TypeVarVariance) -> Self { + pub(crate) fn materialize(self, db: &'db dyn Db, variance: TypeVarVariance) -> Option { TupleType::new(db, self.tuple(db).materialize(db, variance)) } @@ -114,7 +186,7 @@ impl<'db> TupleType<'db> { self, db: &'db dyn Db, type_mapping: &TypeMapping<'a, 'db>, - ) -> Self { + ) -> Option { TupleType::new(db, self.tuple(db).apply_type_mapping(db, type_mapping)) } @@ -149,35 +221,47 @@ impl<'db> TupleType<'db> { } } -/// A fixed-length tuple spec. +/// A tuple spec describes the contents of a tuple type, which might be fixed- or variable-length. /// /// Tuple specs are used for more than just `tuple` instances, so they allow `Never` to appear as a /// fixed-length element type. [`TupleType`] adds that additional invariant (since a tuple that /// must contain an element that can't be instantiated, can't be instantiated itself). -#[derive(Clone, Debug, Default, Eq, Hash, PartialEq, salsa::Update)] -pub struct FixedLengthTupleSpec<'db>(Vec>); +pub(crate) type TupleSpec<'db> = Tuple>; + +/// A fixed-length tuple. +/// +/// Our tuple representation can hold instances of any Rust type. For tuples containing Python +/// types, use [`TupleSpec`], which defines some additional type-specific methods. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub struct FixedLengthTuple(Vec); -impl<'db> FixedLengthTupleSpec<'db> { - pub(crate) fn empty() -> Self { - Self::default() +impl FixedLengthTuple { + fn empty() -> Self { + Self(Vec::new()) } pub(crate) fn with_capacity(capacity: usize) -> Self { Self(Vec::with_capacity(capacity)) } - pub(crate) fn from_elements(elements: impl IntoIterator>) -> Self { + fn from_elements(elements: impl IntoIterator) -> Self { Self(elements.into_iter().collect()) } - pub(crate) fn elements_slice(&self) -> &[Type<'db>] { + pub(crate) fn elements_slice(&self) -> &[T] { &self.0 } - pub(crate) fn elements( - &self, - ) -> impl DoubleEndedIterator> + ExactSizeIterator + '_ { - self.0.iter().copied() + pub(crate) fn elements(&self) -> impl DoubleEndedIterator + ExactSizeIterator + '_ { + self.0.iter() + } + + pub(crate) fn all_elements(&self) -> impl Iterator { + self.0.iter() + } + + pub(crate) fn into_all_elements_with_kind(self) -> impl Iterator> { + self.0.into_iter().map(TupleElement::Fixed) } /// Returns the length of this tuple. @@ -189,26 +273,57 @@ impl<'db> FixedLengthTupleSpec<'db> { self.0.is_empty() } - fn concat(&self, other: &TupleSpec<'db>) -> TupleSpec<'db> { + pub(crate) fn push(&mut self, element: T) { + self.0.push(element); + } +} + +impl<'db> FixedLengthTuple> { + fn concat(&self, other: &Tuple>) -> Tuple> { match other { - TupleSpec::Fixed(other) => TupleSpec::Fixed(FixedLengthTupleSpec::from_elements( - self.elements().chain(other.elements()), + TupleSpec::Fixed(other) => TupleSpec::Fixed(FixedLengthTuple::from_elements( + self.elements().chain(other.elements()).copied(), )), - TupleSpec::Variable(other) => VariableLengthTupleSpec::mixed( - self.elements().chain(other.prefix_elements()), + TupleSpec::Variable(other) => VariableLengthTuple::mixed( + self.elements().chain(other.prefix_elements()).copied(), other.variable, - other.suffix_elements(), + other.suffix_elements().copied(), ), } } - pub(crate) fn push(&mut self, element: Type<'db>) { - self.0.push(element); - } + fn resize( + &self, + db: &'db dyn Db, + new_length: TupleLength, + ) -> Result>, ResizeTupleError> { + match new_length { + TupleLength::Fixed(new_length) => match self.len().cmp(&new_length) { + Ordering::Less => Err(ResizeTupleError::TooFewValues), + Ordering::Greater => Err(ResizeTupleError::TooManyValues), + Ordering::Equal => Ok(Tuple::Fixed(self.clone())), + }, + + TupleLength::Variable(prefix, suffix) => { + // The number of rhs values that will be consumed by the starred target. + let Some(variable) = self.len().checked_sub(prefix + suffix) else { + return Err(ResizeTupleError::TooFewValues); + }; - pub(crate) fn extend_from_slice(&mut self, elements: &[Type<'db>]) { - self.0.extend_from_slice(elements); + // Extract rhs values into the prefix, then into the starred target, then into the + // suffix. + let mut elements = self.elements().copied(); + let prefix = elements.by_ref().take(prefix).collect(); + let variable = UnionType::from_elements(db, elements.by_ref().take(variable)); + let suffix = elements.by_ref().take(suffix).collect(); + Ok(Tuple::Variable(VariableLengthTuple { + prefix, + variable, + suffix, + })) + } + } } #[must_use] @@ -241,18 +356,18 @@ impl<'db> FixedLengthTupleSpec<'db> { fn has_relation_to( &self, db: &'db dyn Db, - other: &TupleSpec<'db>, + other: &Tuple>, relation: TypeRelation, ) -> bool { match other { - TupleSpec::Fixed(other) => { + Tuple::Fixed(other) => { self.0.len() == other.0.len() && (self.0.iter()) .zip(&other.0) .all(|(self_ty, other_ty)| self_ty.has_relation_to(db, *other_ty, relation)) } - TupleSpec::Variable(other) => { + Tuple::Variable(other) => { // This tuple must have enough elements to match up with the other tuple's prefix // and suffix, and each of those elements must pairwise satisfy the relation. let mut self_iter = self.0.iter(); @@ -292,7 +407,20 @@ impl<'db> FixedLengthTupleSpec<'db> { } } -impl<'db> PyIndex<'db> for &FixedLengthTupleSpec<'db> { +#[allow(unsafe_code)] +unsafe impl salsa::Update for FixedLengthTuple +where + T: salsa::Update, +{ + unsafe fn maybe_update(old_pointer: *mut Self, new_value: Self) -> bool { + unsafe { + let old_value = &mut *old_pointer; + Vec::maybe_update(&mut old_value.0, new_value.0) + } + } +} + +impl<'db> PyIndex<'db> for &FixedLengthTuple> { type Item = Type<'db>; fn py_index(self, db: &'db dyn Db, index: i32) -> Result { @@ -300,7 +428,7 @@ impl<'db> PyIndex<'db> for &FixedLengthTupleSpec<'db> { } } -impl<'db> PySlice<'db> for FixedLengthTupleSpec<'db> { +impl<'db> PySlice<'db> for FixedLengthTuple> { type Item = Type<'db>; fn py_slice( @@ -314,54 +442,73 @@ impl<'db> PySlice<'db> for FixedLengthTupleSpec<'db> { } } -/// A variable-length tuple spec. +/// A variable-length tuple. /// -/// The tuple spec can contain a fixed-length heterogeneous prefix and/or suffix. All of the -/// elements of the variable-length portion must be of the same type. +/// The tuple can contain a fixed-length heterogeneous prefix and/or suffix. All of the elements of +/// the variable-length portion must be the same. /// -/// Tuple specs are used for more than just `tuple` instances, so they allow `Never` to appear as a -/// fixed-length element type. [`TupleType`] adds that additional invariant (since a tuple that -/// must contain an element that can't be instantiated, can't be instantiated itself). -#[derive(Clone, Debug, Eq, Hash, PartialEq, salsa::Update)] -pub struct VariableLengthTupleSpec<'db> { - pub(crate) prefix: Vec>, - pub(crate) variable: Type<'db>, - pub(crate) suffix: Vec>, +/// Our tuple representation can hold instances of any Rust type. For tuples containing Python +/// types, use [`TupleSpec`], which defines some additional type-specific methods. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub struct VariableLengthTuple { + pub(crate) prefix: Vec, + pub(crate) variable: T, + pub(crate) suffix: Vec, } -impl<'db> VariableLengthTupleSpec<'db> { +impl VariableLengthTuple { /// Creates a new tuple spec containing zero or more elements of a given type, with no prefix /// or suffix. - fn homogeneous(ty: Type<'db>) -> TupleSpec<'db> { + fn homogeneous(ty: T) -> Tuple { Self::mixed([], ty, []) } fn mixed( - prefix: impl IntoIterator>, - variable: Type<'db>, - suffix: impl IntoIterator>, - ) -> TupleSpec<'db> { - // If the variable-length portion is Never, it can only be instantiated with zero elements. - // That means this isn't a variable-length tuple after all! - if variable.is_never() { - return TupleSpec::Fixed(FixedLengthTupleSpec::from_elements( - prefix.into_iter().chain(suffix), - )); - } - - TupleSpec::Variable(Self { + prefix: impl IntoIterator, + variable: T, + suffix: impl IntoIterator, + ) -> Tuple { + Tuple::Variable(Self { prefix: prefix.into_iter().collect(), variable, suffix: suffix.into_iter().collect(), }) } - fn prefix_elements( - &self, - ) -> impl DoubleEndedIterator> + ExactSizeIterator + '_ { - self.prefix.iter().copied() + fn prefix_elements(&self) -> impl DoubleEndedIterator + ExactSizeIterator + '_ { + self.prefix.iter() + } + + fn suffix_elements(&self) -> impl DoubleEndedIterator + ExactSizeIterator + '_ { + self.suffix.iter() } + fn fixed_elements(&self) -> impl Iterator + '_ { + self.prefix_elements().chain(self.suffix_elements()) + } + + fn all_elements(&self) -> impl Iterator + '_ { + (self.prefix_elements()) + .chain(std::iter::once(&self.variable)) + .chain(self.suffix_elements()) + } + + fn into_all_elements_with_kind(self) -> impl Iterator> { + (self.prefix.into_iter().map(TupleElement::Prefix)) + .chain(std::iter::once(TupleElement::Variable(self.variable))) + .chain(self.suffix.into_iter().map(TupleElement::Suffix)) + } + + fn len(&self) -> TupleLength { + TupleLength::Variable(self.prefix.len(), self.suffix.len()) + } + + fn push(&mut self, element: T) { + self.suffix.push(element); + } +} + +impl<'db> VariableLengthTuple> { /// Returns the prefix of the prenormalization of this tuple. /// /// This is used in our subtyping and equivalence checks below to handle different tuple types @@ -387,16 +534,12 @@ impl<'db> VariableLengthTupleSpec<'db> { variable: Option>, ) -> impl Iterator> + 'a { let variable = variable.unwrap_or(self.variable); - self.prefix_elements().chain( - self.suffix_elements() - .take_while(move |element| element.is_equivalent_to(db, variable)), - ) - } - - fn suffix_elements( - &self, - ) -> impl DoubleEndedIterator> + ExactSizeIterator + '_ { - self.suffix.iter().copied() + self.prefix_elements() + .chain( + self.suffix_elements() + .take_while(move |element| element.is_equivalent_to(db, variable)), + ) + .copied() } /// Returns the suffix of the prenormalization of this tuple. @@ -426,49 +569,74 @@ impl<'db> VariableLengthTupleSpec<'db> { let variable = variable.unwrap_or(self.variable); self.suffix_elements() .skip_while(move |element| element.is_equivalent_to(db, variable)) + .copied() } - fn fixed_elements(&self) -> impl Iterator> + '_ { - self.prefix_elements().chain(self.suffix_elements()) - } - - fn all_elements(&self) -> impl Iterator> + '_ { - self.prefix_elements() - .chain(std::iter::once(self.variable)) - .chain(self.suffix_elements()) - } - - /// Returns the minimum length of this tuple. - pub(crate) fn minimum_length(&self) -> usize { - self.prefix.len() + self.suffix.len() - } - - fn concat(&self, db: &'db dyn Db, other: &TupleSpec<'db>) -> TupleSpec<'db> { + fn concat(&self, db: &'db dyn Db, other: &Tuple>) -> Tuple> { match other { - TupleSpec::Fixed(other) => VariableLengthTupleSpec::mixed( - self.prefix_elements(), + TupleSpec::Fixed(other) => VariableLengthTuple::mixed( + self.prefix_elements().copied(), self.variable, - self.suffix_elements().chain(other.elements()), + self.suffix_elements().chain(other.elements()).copied(), ), - TupleSpec::Variable(other) => { + Tuple::Variable(other) => { let variable = UnionType::from_elements( db, - self.suffix_elements() + (self.suffix_elements().copied()) .chain([self.variable, other.variable]) - .chain(other.prefix_elements()), + .chain(other.prefix_elements().copied()), ); - VariableLengthTupleSpec::mixed( - self.prefix_elements(), + VariableLengthTuple::mixed( + self.prefix_elements().copied(), variable, - other.suffix_elements(), + other.suffix_elements().copied(), ) } } } - fn push(&mut self, element: Type<'db>) { - self.suffix.push(element); + fn resize( + &self, + db: &'db dyn Db, + new_length: TupleLength, + ) -> Result>, ResizeTupleError> { + match new_length { + TupleLength::Fixed(new_length) => { + // The number of elements that will get their value from our variable-length + // portion. + let Some(variable_count) = new_length.checked_sub(self.len().minimum()) else { + return Err(ResizeTupleError::TooManyValues); + }; + Ok(Tuple::Fixed(FixedLengthTuple::from_elements( + (self.prefix_elements().copied()) + .chain(std::iter::repeat_n(self.variable, variable_count)) + .chain(self.suffix_elements().copied()), + ))) + } + + TupleLength::Variable(prefix_length, suffix_length) => { + // "Overflow" are elements of our prefix/suffix that will be folded into the + // result's variable-length portion. "Underflow" are elements of the result + // prefix/suffix that will come from our variable-length portion. + let self_prefix_length = self.prefix.len(); + let prefix_underflow = prefix_length.saturating_sub(self_prefix_length); + let self_suffix_length = self.suffix.len(); + let suffix_overflow = self_suffix_length.saturating_sub(suffix_length); + let suffix_underflow = suffix_length.saturating_sub(self_suffix_length); + let prefix = (self.prefix_elements().copied().take(prefix_length)) + .chain(std::iter::repeat_n(self.variable, prefix_underflow)); + let variable = UnionType::from_elements( + db, + (self.prefix_elements().copied().skip(prefix_length)) + .chain(std::iter::once(self.variable)) + .chain(self.suffix_elements().copied().take(suffix_overflow)), + ); + let suffix = std::iter::repeat_n(self.variable, suffix_underflow) + .chain(self.suffix_elements().copied().skip(suffix_overflow)); + Ok(VariableLengthTuple::mixed(prefix, variable, suffix)) + } + } } #[must_use] @@ -523,11 +691,11 @@ impl<'db> VariableLengthTupleSpec<'db> { fn has_relation_to( &self, db: &'db dyn Db, - other: &TupleSpec<'db>, + other: &Tuple>, relation: TypeRelation, ) -> bool { match other { - TupleSpec::Fixed(other) => { + Tuple::Fixed(other) => { // The `...` length specifier of a variable-length tuple type is interpreted // differently depending on the type of the variable-length elements. // @@ -546,7 +714,7 @@ impl<'db> VariableLengthTupleSpec<'db> { // In addition, the other tuple must have enough elements to match up with this // tuple's prefix and suffix, and each of those elements must pairwise satisfy the // relation. - let mut other_iter = other.elements(); + let mut other_iter = other.elements().copied(); for self_ty in self.prenormalized_prefix_elements(db, None) { let Some(other_ty) = other_iter.next() else { return false; @@ -568,7 +736,7 @@ impl<'db> VariableLengthTupleSpec<'db> { true } - TupleSpec::Variable(other) => { + Tuple::Variable(other) => { // When prenormalizing below, we assume that a dynamic variable-length portion of // one tuple materializes to the variable-length portion of the other tuple. let self_prenormalize_variable = match self.variable { @@ -653,7 +821,22 @@ impl<'db> VariableLengthTupleSpec<'db> { } } -impl<'db> PyIndex<'db> for &VariableLengthTupleSpec<'db> { +#[allow(unsafe_code)] +unsafe impl salsa::Update for VariableLengthTuple +where + T: salsa::Update, +{ + unsafe fn maybe_update(old_pointer: *mut Self, new_value: Self) -> bool { + let old_value = unsafe { &mut *old_pointer }; + let mut changed = false; + changed |= unsafe { Vec::maybe_update(&mut old_value.prefix, new_value.prefix) }; + changed |= unsafe { T::maybe_update(&mut old_value.variable, new_value.variable) }; + changed |= unsafe { Vec::maybe_update(&mut old_value.suffix, new_value.suffix) }; + changed + } +} + +impl<'db> PyIndex<'db> for &VariableLengthTuple> { type Item = Type<'db>; fn py_index(self, db: &'db dyn Db, index: i32) -> Result { @@ -671,7 +854,7 @@ impl<'db> PyIndex<'db> for &VariableLengthTupleSpec<'db> { Ok(UnionType::from_elements( db, std::iter::once(self.variable) - .chain(self.suffix_elements().take(index_past_prefix)), + .chain(self.suffix_elements().copied().take(index_past_prefix)), )) } @@ -687,7 +870,7 @@ impl<'db> PyIndex<'db> for &VariableLengthTupleSpec<'db> { let index_past_suffix = index_from_end - self.suffix.len() + 1; Ok(UnionType::from_elements( db, - (self.prefix_elements().rev()) + (self.prefix_elements().rev().copied()) .take(index_past_suffix) .rev() .chain(std::iter::once(self.variable)), @@ -697,106 +880,125 @@ impl<'db> PyIndex<'db> for &VariableLengthTupleSpec<'db> { } } -/// A tuple spec that might be fixed- or variable-length. +/// A tuple that might be fixed- or variable-length. /// -/// Tuple specs are used for more than just `tuple` instances, so they allow `Never` to appear as a -/// fixed-length element type. [`TupleType`] adds that additional invariant (since a tuple that -/// must contain an element that can't be instantiated, can't be instantiated itself). -#[derive(Clone, Debug, Eq, Hash, PartialEq, salsa::Update)] -pub enum TupleSpec<'db> { - Fixed(FixedLengthTupleSpec<'db>), - Variable(VariableLengthTupleSpec<'db>), +/// Our tuple representation can hold instances of any Rust type. For tuples containing Python +/// types, use [`TupleSpec`], which defines some additional type-specific methods. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub enum Tuple { + Fixed(FixedLengthTuple), + Variable(VariableLengthTuple), } -impl<'db> TupleSpec<'db> { - pub(crate) fn with_capacity(capacity: usize) -> Self { - TupleSpec::Fixed(FixedLengthTupleSpec::with_capacity(capacity)) +impl Tuple { + pub(crate) fn homogeneous(element: T) -> Self { + VariableLengthTuple::homogeneous(element) } - pub(crate) fn homogeneous(element: Type<'db>) -> Self { - VariableLengthTupleSpec::homogeneous(element) + pub(crate) fn from_elements(elements: impl IntoIterator) -> Self { + FixedLengthTuple::from_elements(elements).into() + } + + pub(crate) fn with_capacity(capacity: usize) -> Self { + Tuple::Fixed(FixedLengthTuple::with_capacity(capacity)) } /// Returns an iterator of all of the fixed-length element types of this tuple. - pub(crate) fn fixed_elements(&self) -> impl Iterator> + '_ { + pub(crate) fn fixed_elements(&self) -> impl Iterator + '_ { match self { - TupleSpec::Fixed(tuple) => Either::Left(tuple.elements()), - TupleSpec::Variable(tuple) => Either::Right(tuple.fixed_elements()), + Tuple::Fixed(tuple) => Either::Left(tuple.elements()), + Tuple::Variable(tuple) => Either::Right(tuple.fixed_elements()), } } /// Returns an iterator of all of the element types of this tuple. Does not deduplicate the /// elements, and does not distinguish between fixed- and variable-length elements. - pub(crate) fn all_elements(&self) -> impl Iterator> + '_ { + pub(crate) fn all_elements(&self) -> impl Iterator + '_ { match self { - TupleSpec::Fixed(tuple) => Either::Left(tuple.elements()), - TupleSpec::Variable(tuple) => Either::Right(tuple.all_elements()), + Tuple::Fixed(tuple) => Either::Left(tuple.all_elements()), + Tuple::Variable(tuple) => Either::Right(tuple.all_elements()), } } - pub(crate) fn display_minimum_length(&self) -> String { + pub(crate) fn into_all_elements_with_kind(self) -> impl Iterator> { match self { - TupleSpec::Fixed(tuple) => tuple.len().to_string(), - TupleSpec::Variable(tuple) => format!("at least {}", tuple.minimum_length()), + Tuple::Fixed(tuple) => Either::Left(tuple.into_all_elements_with_kind()), + Tuple::Variable(tuple) => Either::Right(tuple.into_all_elements_with_kind()), } } pub(crate) const fn is_variadic(&self) -> bool { - matches!(self, TupleSpec::Variable(_)) + matches!(self, Tuple::Variable(_)) } - /// Returns the minimum and maximum length of this tuple. (The maximum length will be `None` - /// for a tuple with a variable-length portion.) - pub(crate) fn size_hint(&self) -> (usize, Option) { + /// Returns the length of this tuple. + pub(crate) fn len(&self) -> TupleLength { match self { - TupleSpec::Fixed(tuple) => { - let len = tuple.len(); - (len, Some(len)) - } - TupleSpec::Variable(tuple) => (tuple.minimum_length(), None), + Tuple::Fixed(tuple) => TupleLength::Fixed(tuple.len()), + Tuple::Variable(tuple) => tuple.len(), } } pub(crate) fn is_empty(&self) -> bool { match self { - TupleSpec::Fixed(tuple) => tuple.is_empty(), - TupleSpec::Variable(_) => false, + Tuple::Fixed(tuple) => tuple.is_empty(), + Tuple::Variable(_) => false, } } + pub(crate) fn push(&mut self, element: T) { + match self { + Tuple::Fixed(tuple) => tuple.push(element), + Tuple::Variable(tuple) => tuple.push(element), + } + } +} + +impl<'db> Tuple> { /// Concatenates another tuple to the end of this tuple, returning a new tuple. pub(crate) fn concat(&self, db: &'db dyn Db, other: &Self) -> Self { match self { - TupleSpec::Fixed(tuple) => tuple.concat(other), - TupleSpec::Variable(tuple) => tuple.concat(db, other), + Tuple::Fixed(tuple) => tuple.concat(other), + Tuple::Variable(tuple) => tuple.concat(db, other), } } - pub(crate) fn push(&mut self, element: Type<'db>) { + /// Resizes this tuple to a different length, if possible. If this tuple cannot satisfy the + /// desired minimum or maximum length, we return an error. If we return an `Ok` result, the + /// [`len`][Self::len] of the resulting tuple is guaranteed to be equal to `new_length`. + pub(crate) fn resize( + &self, + db: &'db dyn Db, + new_length: TupleLength, + ) -> Result { match self { - TupleSpec::Fixed(tuple) => tuple.push(element), - TupleSpec::Variable(tuple) => tuple.push(element), + Tuple::Fixed(tuple) => tuple.resize(db, new_length), + Tuple::Variable(tuple) => tuple.resize(db, new_length), } } - fn normalized(&self, db: &'db dyn Db) -> Self { + pub(crate) fn normalized(&self, db: &'db dyn Db) -> Self { match self { - TupleSpec::Fixed(tuple) => TupleSpec::Fixed(tuple.normalized(db)), - TupleSpec::Variable(tuple) => tuple.normalized(db), + Tuple::Fixed(tuple) => Tuple::Fixed(tuple.normalized(db)), + Tuple::Variable(tuple) => tuple.normalized(db), } } - fn materialize(&self, db: &'db dyn Db, variance: TypeVarVariance) -> Self { + pub(crate) fn materialize(&self, db: &'db dyn Db, variance: TypeVarVariance) -> Self { match self { - TupleSpec::Fixed(tuple) => TupleSpec::Fixed(tuple.materialize(db, variance)), - TupleSpec::Variable(tuple) => tuple.materialize(db, variance), + Tuple::Fixed(tuple) => Tuple::Fixed(tuple.materialize(db, variance)), + Tuple::Variable(tuple) => tuple.materialize(db, variance), } } - fn apply_type_mapping<'a>(&self, db: &'db dyn Db, type_mapping: &TypeMapping<'a, 'db>) -> Self { + pub(crate) fn apply_type_mapping<'a>( + &self, + db: &'db dyn Db, + type_mapping: &TypeMapping<'a, 'db>, + ) -> Self { match self { - TupleSpec::Fixed(tuple) => TupleSpec::Fixed(tuple.apply_type_mapping(db, type_mapping)), - TupleSpec::Variable(tuple) => tuple.apply_type_mapping(db, type_mapping), + Tuple::Fixed(tuple) => Tuple::Fixed(tuple.apply_type_mapping(db, type_mapping)), + Tuple::Variable(tuple) => tuple.apply_type_mapping(db, type_mapping), } } @@ -806,35 +1008,34 @@ impl<'db> TupleSpec<'db> { typevars: &mut FxOrderSet>, ) { match self { - TupleSpec::Fixed(tuple) => tuple.find_legacy_typevars(db, typevars), - TupleSpec::Variable(tuple) => tuple.find_legacy_typevars(db, typevars), + Tuple::Fixed(tuple) => tuple.find_legacy_typevars(db, typevars), + Tuple::Variable(tuple) => tuple.find_legacy_typevars(db, typevars), } } fn has_relation_to(&self, db: &'db dyn Db, other: &Self, relation: TypeRelation) -> bool { match self { - TupleSpec::Fixed(self_tuple) => self_tuple.has_relation_to(db, other, relation), - TupleSpec::Variable(self_tuple) => self_tuple.has_relation_to(db, other, relation), + Tuple::Fixed(self_tuple) => self_tuple.has_relation_to(db, other, relation), + Tuple::Variable(self_tuple) => self_tuple.has_relation_to(db, other, relation), } } fn is_equivalent_to(&self, db: &'db dyn Db, other: &Self) -> bool { match (self, other) { - (TupleSpec::Fixed(self_tuple), TupleSpec::Fixed(other_tuple)) => { + (Tuple::Fixed(self_tuple), Tuple::Fixed(other_tuple)) => { self_tuple.is_equivalent_to(db, other_tuple) } - (TupleSpec::Variable(self_tuple), TupleSpec::Variable(other_tuple)) => { + (Tuple::Variable(self_tuple), Tuple::Variable(other_tuple)) => { self_tuple.is_equivalent_to(db, other_tuple) } - (TupleSpec::Fixed(_), TupleSpec::Variable(_)) - | (TupleSpec::Variable(_), TupleSpec::Fixed(_)) => false, + (Tuple::Fixed(_), Tuple::Variable(_)) | (Tuple::Variable(_), Tuple::Fixed(_)) => false, } } fn is_disjoint_from(&self, db: &'db dyn Db, other: &Self) -> bool { // Two tuples with an incompatible number of required elements must always be disjoint. - let (self_min, self_max) = self.size_hint(); - let (other_min, other_max) = other.size_hint(); + let (self_min, self_max) = self.len().size_hint(); + let (other_min, other_max) = other.len().size_hint(); if self_max.is_some_and(|max| max < other_min) { return true; } @@ -846,22 +1047,22 @@ impl<'db> TupleSpec<'db> { #[allow(clippy::items_after_statements)] fn any_disjoint<'db>( db: &'db dyn Db, - a: impl IntoIterator>, - b: impl IntoIterator>, + a: impl IntoIterator>, + b: impl IntoIterator>, ) -> bool { a.into_iter().zip(b).any(|(self_element, other_element)| { - self_element.is_disjoint_from(db, other_element) + self_element.is_disjoint_from(db, *other_element) }) } match (self, other) { - (TupleSpec::Fixed(self_tuple), TupleSpec::Fixed(other_tuple)) => { + (Tuple::Fixed(self_tuple), Tuple::Fixed(other_tuple)) => { if any_disjoint(db, self_tuple.elements(), other_tuple.elements()) { return true; } } - (TupleSpec::Variable(self_tuple), TupleSpec::Variable(other_tuple)) => { + (Tuple::Variable(self_tuple), Tuple::Variable(other_tuple)) => { if any_disjoint( db, self_tuple.prefix_elements(), @@ -878,8 +1079,8 @@ impl<'db> TupleSpec<'db> { } } - (TupleSpec::Fixed(fixed), TupleSpec::Variable(variable)) - | (TupleSpec::Variable(variable), TupleSpec::Fixed(fixed)) => { + (Tuple::Fixed(fixed), Tuple::Variable(variable)) + | (Tuple::Variable(variable), Tuple::Fixed(fixed)) => { if any_disjoint(db, fixed.elements(), variable.prefix_elements()) { return true; } @@ -897,31 +1098,161 @@ impl<'db> TupleSpec<'db> { fn is_single_valued(&self, db: &'db dyn Db) -> bool { match self { - TupleSpec::Fixed(tuple) => tuple.is_single_valued(db), - TupleSpec::Variable(_) => false, + Tuple::Fixed(tuple) => tuple.is_single_valued(db), + Tuple::Variable(_) => false, } } } -impl<'db> From> for TupleSpec<'db> { - fn from(tuple: FixedLengthTupleSpec<'db>) -> Self { - TupleSpec::Fixed(tuple) +impl From> for Tuple { + fn from(tuple: FixedLengthTuple) -> Self { + Tuple::Fixed(tuple) } } -impl<'db> From> for TupleSpec<'db> { - fn from(tuple: VariableLengthTupleSpec<'db>) -> Self { - TupleSpec::Variable(tuple) +impl From> for Tuple { + fn from(tuple: VariableLengthTuple) -> Self { + Tuple::Variable(tuple) } } -impl<'db> PyIndex<'db> for &TupleSpec<'db> { +#[allow(unsafe_code)] +unsafe impl salsa::Update for Tuple +where + T: salsa::Update, +{ + unsafe fn maybe_update(old_pointer: *mut Self, new_value: Self) -> bool { + let old_value = unsafe { &mut *old_pointer }; + match (old_value, new_value) { + (Tuple::Fixed(old), Tuple::Fixed(new)) => unsafe { + FixedLengthTuple::maybe_update(old, new) + }, + (Tuple::Variable(old), Tuple::Variable(new)) => unsafe { + VariableLengthTuple::maybe_update(old, new) + }, + (old_value, new_value) => { + *old_value = new_value; + true + } + } + } +} + +impl<'db> PyIndex<'db> for &Tuple> { type Item = Type<'db>; fn py_index(self, db: &'db dyn Db, index: i32) -> Result { match self { - TupleSpec::Fixed(tuple) => tuple.py_index(db, index), - TupleSpec::Variable(tuple) => tuple.py_index(db, index), + Tuple::Fixed(tuple) => tuple.py_index(db, index), + Tuple::Variable(tuple) => tuple.py_index(db, index), + } + } +} + +pub(crate) enum TupleElement { + Fixed(T), + Prefix(T), + Variable(T), + Suffix(T), +} + +/// Unpacks tuple values in an unpacking assignment. +/// +/// You provide a [`TupleLength`] specifying how many assignment targets there are, and which one +/// (if any) is a starred target. You then call [`unpack_tuple`][TupleUnpacker::unpack_tuple] to +/// unpack the values from a rhs tuple into those targets. If the rhs is a union, call +/// `unpack_tuple` separately for each element of the union. We will automatically wrap the types +/// assigned to the starred target in `list`. +pub(crate) struct TupleUnpacker<'db> { + db: &'db dyn Db, + targets: Tuple>, +} + +impl<'db> TupleUnpacker<'db> { + pub(crate) fn new(db: &'db dyn Db, len: TupleLength) -> Self { + let new_builders = |len: usize| std::iter::repeat_with(|| UnionBuilder::new(db)).take(len); + let targets = match len { + TupleLength::Fixed(len) => { + Tuple::Fixed(FixedLengthTuple::from_elements(new_builders(len))) + } + TupleLength::Variable(prefix, suffix) => VariableLengthTuple::mixed( + new_builders(prefix), + UnionBuilder::new(db), + new_builders(suffix), + ), + }; + Self { db, targets } + } + + /// Unpacks a single rhs tuple into the target tuple that we are building. If you want to + /// unpack a single type into each target, call this method with a homogeneous tuple. + /// + /// The lengths of the targets and the rhs have to be compatible, but not necessarily + /// identical. The lengths only have to be identical if both sides are fixed-length; if either + /// side is variable-length, we will pull multiple values out of the rhs variable-length + /// portion, and assign multiple values to the starred target, as needed. + pub(crate) fn unpack_tuple( + &mut self, + values: &Tuple>, + ) -> Result<(), ResizeTupleError> { + let values = values.resize(self.db, self.targets.len())?; + match (&mut self.targets, &values) { + (Tuple::Fixed(targets), Tuple::Fixed(values)) => { + targets.unpack_tuple(values); + } + (Tuple::Variable(targets), Tuple::Variable(values)) => { + targets.unpack_tuple(self.db, values); + } + _ => panic!("should have ensured that tuples are the same length"), } + Ok(()) + } + + /// Returns the unpacked types for each target. If you called + /// [`unpack_tuple`][TupleUnpacker::unpack_tuple] multiple times, each target type will be the + /// union of the type unpacked into that target from each of the rhs tuples. If there is a + /// starred target, we will each unpacked type in `list`. + pub(crate) fn into_types(self) -> impl Iterator> { + self.targets + .into_all_elements_with_kind() + .map(|builder| match builder { + TupleElement::Variable(builder) => builder.try_build().unwrap_or_else(|| { + KnownClass::List.to_specialized_instance(self.db, [Type::unknown()]) + }), + TupleElement::Fixed(builder) + | TupleElement::Prefix(builder) + | TupleElement::Suffix(builder) => { + builder.try_build().unwrap_or_else(Type::unknown) + } + }) } } + +impl<'db> FixedLengthTuple> { + fn unpack_tuple(&mut self, values: &FixedLengthTuple>) { + // We have already verified above that the two tuples have the same length. + for (target, value) in self.0.iter_mut().zip(values.elements().copied()) { + target.add_in_place(value); + } + } +} + +impl<'db> VariableLengthTuple> { + fn unpack_tuple(&mut self, db: &'db dyn Db, values: &VariableLengthTuple>) { + // We have already verified above that the two tuples have the same length. + for (target, value) in (self.prefix.iter_mut()).zip(values.prefix_elements().copied()) { + target.add_in_place(value); + } + self.variable + .add_in_place(KnownClass::List.to_specialized_instance(db, [values.variable])); + for (target, value) in (self.suffix.iter_mut()).zip(values.suffix_elements().copied()) { + target.add_in_place(value); + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) enum ResizeTupleError { + TooFewValues, + TooManyValues, +} diff --git a/crates/ty_python_semantic/src/types/unpacker.rs b/crates/ty_python_semantic/src/types/unpacker.rs index 3674e5564740a..b9df27a78f99d 100644 --- a/crates/ty_python_semantic/src/types/unpacker.rs +++ b/crates/ty_python_semantic/src/types/unpacker.rs @@ -1,5 +1,4 @@ use std::borrow::Cow; -use std::cmp::Ordering; use ruff_db::parsed::ParsedModuleRef; use rustc_hash::FxHashMap; @@ -9,13 +8,12 @@ use ruff_python_ast::{self as ast, AnyNodeRef}; use crate::Db; use crate::semantic_index::ast_ids::{HasScopedExpressionId, ScopedExpressionId}; use crate::semantic_index::place::ScopeId; -use crate::types::tuple::{FixedLengthTupleSpec, TupleSpec, TupleType}; -use crate::types::{Type, TypeCheckDiagnostics, infer_expression_types, todo_type}; +use crate::types::tuple::{ResizeTupleError, Tuple, TupleLength, TupleUnpacker}; +use crate::types::{Type, TypeCheckDiagnostics, infer_expression_types}; use crate::unpack::{UnpackKind, UnpackValue}; use super::context::InferContext; use super::diagnostic::INVALID_ASSIGNMENT; -use super::{KnownClass, UnionType}; /// Unpacks the value expression type to their respective targets. pub(crate) struct Unpacker<'db, 'ast> { @@ -115,18 +113,13 @@ impl<'db, 'ast> Unpacker<'db, 'ast> { } ast::Expr::List(ast::ExprList { elts, .. }) | ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => { - // Initialize the vector of target types, one for each target. - // - // This is mainly useful for the union type where the target type at index `n` is - // going to be a union of types from every union type element at index `n`. - // - // For example, if the type is `tuple[int, int] | tuple[int, str]` and the target - // has two elements `(a, b)`, then - // * The type of `a` will be a union of `int` and `int` which are at index 0 in the - // first and second tuple respectively which resolves to an `int`. - // * Similarly, the type of `b` will be a union of `int` and `str` which are at - // index 1 in the first and second tuple respectively which will be `int | str`. - let mut target_types = vec![vec![]; elts.len()]; + let target_len = match elts.iter().position(ast::Expr::is_starred_expr) { + Some(starred_index) => { + TupleLength::Variable(starred_index, elts.len() - (starred_index + 1)) + } + None => TupleLength::Fixed(elts.len()), + }; + let mut unpacker = TupleUnpacker::new(self.db(), target_len); let unpack_types = match value_ty { Type::Union(union_ty) => union_ty.elements(self.db()), @@ -134,205 +127,75 @@ impl<'db, 'ast> Unpacker<'db, 'ast> { }; for ty in unpack_types.iter().copied() { - // Deconstruct certain types to delegate the inference back to the tuple type - // for correct handling of starred expressions. - let ty = match ty { + let tuple = match ty { + Type::Tuple(tuple_ty) => Cow::Borrowed(tuple_ty.tuple(self.db())), Type::StringLiteral(string_literal_ty) => { // We could go further and deconstruct to an array of `StringLiteral` // with each individual character, instead of just an array of // `LiteralString`, but there would be a cost and it's not clear that // it's worth it. - TupleType::from_elements( - self.db(), - std::iter::repeat_n( - Type::LiteralString, - string_literal_ty.python_len(self.db()), - ), - ) + Cow::Owned(Tuple::from_elements(std::iter::repeat_n( + Type::LiteralString, + string_literal_ty.python_len(self.db()), + ))) + } + Type::LiteralString => Cow::Owned(Tuple::homogeneous(Type::LiteralString)), + _ => { + // TODO: Update our iterator protocol machinery to return a tuple + // describing the returned values in more detail, when we can. + Cow::Owned(Tuple::homogeneous( + ty.try_iterate(self.db()).unwrap_or_else(|err| { + err.report_diagnostic(&self.context, ty, value_expr); + err.fallback_element_type(self.db()) + }), + )) } - _ => ty, }; - if let Type::Tuple(tuple_ty) = ty { - let tuple = self.tuple_ty_elements(target, elts, tuple_ty, value_expr); - - let length_mismatch = match elts.len().cmp(&tuple.len()) { - Ordering::Less => { - if let Some(builder) = - self.context.report_lint(&INVALID_ASSIGNMENT, target) - { + if let Err(err) = unpacker.unpack_tuple(tuple.as_ref()) { + unpacker + .unpack_tuple(&Tuple::homogeneous(Type::unknown())) + .expect("adding a homogeneous tuple should always succeed"); + if let Some(builder) = self.context.report_lint(&INVALID_ASSIGNMENT, target) + { + match err { + ResizeTupleError::TooManyValues => { let mut diag = builder.into_diagnostic("Too many values to unpack"); diag.set_primary_message(format_args!( "Expected {}", - elts.len(), + target_len.display_minimum(), + )); + diag.annotate(self.context.secondary(value_expr).message( + format_args!("Got {}", tuple.len().display_minimum()), )); - diag.annotate( - self.context - .secondary(value_expr) - .message(format_args!("Got {}", tuple.len())), - ); } - true - } - Ordering::Greater => { - if let Some(builder) = - self.context.report_lint(&INVALID_ASSIGNMENT, target) - { + ResizeTupleError::TooFewValues => { let mut diag = builder.into_diagnostic("Not enough values to unpack"); diag.set_primary_message(format_args!( "Expected {}", - elts.len(), + target_len.display_minimum(), + )); + diag.annotate(self.context.secondary(value_expr).message( + format_args!("Got {}", tuple.len().display_maximum()), )); - diag.annotate( - self.context - .secondary(value_expr) - .message(format_args!("Got {}", tuple.len())), - ); - } - true - } - Ordering::Equal => false, - }; - - for (index, ty) in tuple.elements().enumerate() { - if let Some(element_types) = target_types.get_mut(index) { - if length_mismatch { - element_types.push(Type::unknown()); - } else { - element_types.push(ty); } } } - } else { - let ty = if ty.is_literal_string() { - Type::LiteralString - } else { - ty.try_iterate(self.db()).unwrap_or_else(|err| { - err.report_diagnostic(&self.context, ty, value_expr); - err.fallback_element_type(self.db()) - }) - }; - // Both `elts` and `target_types` are guaranteed to have the same length. - for (element, target_type) in elts.iter().zip(&mut target_types) { - if element.is_starred_expr() { - target_type.push( - KnownClass::List.to_specialized_instance(self.db(), [ty]), - ); - } else { - target_type.push(ty); - } - } } } - for (index, element) in elts.iter().enumerate() { - // SAFETY: `target_types` is initialized with the same length as `elts`. - let element_ty = match target_types[index].as_slice() { - [] => Type::unknown(), - types => UnionType::from_elements(self.db(), types), - }; - self.unpack_inner(element, value_expr, element_ty); + // We constructed unpacker above using the length of elts, so the zip should + // consume the same number of elements from each. + for (target, value_ty) in elts.iter().zip(unpacker.into_types()) { + self.unpack_inner(target, value_expr, value_ty); } } _ => {} } } - /// Returns the [`Type`] elements inside the given [`TupleType`] taking into account that there - /// can be a starred expression in the `elements`. - /// - /// `value_expr` is an AST reference to the value being unpacked. It is - /// only used for diagnostics. - fn tuple_ty_elements( - &self, - expr: &ast::Expr, - targets: &[ast::Expr], - tuple_ty: TupleType<'db>, - value_expr: AnyNodeRef<'_>, - ) -> Cow<'_, FixedLengthTupleSpec<'db>> { - let TupleSpec::Fixed(tuple) = tuple_ty.tuple(self.db()) else { - let todo = todo_type!("Unpack variable-length tuple"); - return Cow::Owned(FixedLengthTupleSpec::from_elements(targets.iter().map( - |target| { - if target.is_starred_expr() { - KnownClass::List.to_specialized_instance(self.db(), [todo]) - } else { - todo - } - }, - ))); - }; - - // If there is a starred expression, it will consume all of the types at that location. - let Some(starred_index) = targets.iter().position(ast::Expr::is_starred_expr) else { - // Otherwise, the types will be unpacked 1-1 to the targets. - return Cow::Borrowed(tuple); - }; - - if tuple.len() >= targets.len() - 1 { - // This branch is only taken when there are enough elements in the tuple type to - // combine for the starred expression. So, the arithmetic and indexing operations are - // safe to perform. - let mut element_types = FixedLengthTupleSpec::with_capacity(targets.len()); - let tuple_elements = tuple.elements_slice(); - - // Insert all the elements before the starred expression. - // SAFETY: Safe because of the length check above. - element_types.extend_from_slice(&tuple_elements[..starred_index]); - - // The number of target expressions that are remaining after the starred expression. - // For example, in `(a, *b, c, d) = ...`, the index of starred element `b` is 1 and the - // remaining elements after that are 2. - let remaining = targets.len() - (starred_index + 1); - - // This index represents the position of the last element that belongs to the starred - // expression, in an exclusive manner. For example, in `(a, *b, c) = (1, 2, 3, 4)`, the - // starred expression `b` will consume the elements `Literal[2]` and `Literal[3]` and - // the index value would be 3. - let starred_end_index = tuple.len() - remaining; - - // SAFETY: Safe because of the length check above. - let starred_element_types = &tuple_elements[starred_index..starred_end_index]; - - element_types.push(KnownClass::List.to_specialized_instance( - self.db(), - [if starred_element_types.is_empty() { - Type::unknown() - } else { - UnionType::from_elements(self.db(), starred_element_types) - }], - )); - - // Insert the types remaining that aren't consumed by the starred expression. - // SAFETY: Safe because of the length check above. - element_types.extend_from_slice(&tuple_elements[starred_end_index..]); - - Cow::Owned(element_types) - } else { - if let Some(builder) = self.context.report_lint(&INVALID_ASSIGNMENT, expr) { - let mut diag = builder.into_diagnostic("Not enough values to unpack"); - diag.set_primary_message(format_args!("Expected {} or more", targets.len() - 1)); - diag.annotate( - self.context - .secondary(value_expr) - .message(format_args!("Got {}", tuple.len())), - ); - } - - Cow::Owned(FixedLengthTupleSpec::from_elements(targets.iter().map( - |target| { - if target.is_starred_expr() { - KnownClass::List.to_specialized_instance(self.db(), [Type::unknown()]) - } else { - Type::unknown() - } - }, - ))) - } - } - pub(crate) fn finish(mut self) -> UnpackResult<'db> { self.targets.shrink_to_fit(); UnpackResult {