diff --git a/crates/ty_python_semantic/resources/mdtest/generics/legacy/callables.md b/crates/ty_python_semantic/resources/mdtest/generics/legacy/callables.md index 61f87dcc8408c4..31ea3b6614d039 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/legacy/callables.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/legacy/callables.md @@ -386,6 +386,95 @@ def singleton(flag: bool = False) -> Callable[[Callable[[int], S]], Callable[[in return wrapper ``` +## Return type inference from partially annotated overloads + +The catch-all overload returns `object`, which is preserved when inferring a return type from the +whole callback even though the literal-specific overloads have unannotated return types. + +```py +from typing import Callable, Literal, TypeVar, overload +from typing_extensions import assert_type + +R = TypeVar("R") +T = TypeVar("T") + +def infer_return(callback: Callable[[T], R]) -> R: + raise NotImplementedError + +@overload +def callback(value: Literal["a"]): ... +@overload +def callback(value: Literal["b"]): ... +@overload +def callback(value: Literal["c"]): ... +@overload +def callback(value: Literal["d", "e"]): ... +@overload +def callback(value: Literal["f", "g"]): ... +@overload +def callback(value: Literal["h", "i"]): ... +@overload +def callback(value: Literal["j", "k"]): ... +@overload +def callback(value: object) -> object: ... +def callback(value): + raise NotImplementedError + +assert_type(infer_return(callback), object) +``` + +## Generic inference after projection budget exhaustion + +The literal-specific overloads below produce more alternative bindings than generic inference can +project within its limits. The precise type of `default=0` does not replace the missing callback +evidence: we recover with `Unknown` in either argument order. + +```py +from typing import Callable, Literal, TypeVar, overload +from typing_extensions import assert_type +from ty_extensions._internal import Unknown + +R = TypeVar("R") +T = TypeVar("T") + +def infer_return(callback: Callable[[T], R], default: R) -> R: + raise NotImplementedError + +@overload +def callback(value: Literal[0, 1]): ... +@overload +def callback(value: Literal[2, 3]): ... +@overload +def callback(value: Literal[4, 5]): ... +@overload +def callback(value: Literal[6, 7]): ... +@overload +def callback(value: Literal[8, 9]): ... +@overload +def callback(value: Literal[10, 11]): ... +@overload +def callback(value: Literal[12, 13]): ... +@overload +def callback(value: Literal[14, 15]): ... +@overload +def callback(value: Literal[16, 17]): ... +@overload +def callback(value: Literal[18, 19]): ... +@overload +def callback(value: Literal[20, 21]): ... +@overload +def callback(value: Literal[22, 23]): ... +@overload +def callback(value: Literal[24, 25]): ... +@overload +def callback(value: object) -> object: ... +def callback(value): + raise NotImplementedError + +assert_type(infer_return(callback, 0), Unknown) +assert_type(infer_return(default=0, callback=callback), Unknown) +``` + ## Multiple occurrences of a higher-order generic callable If a generic callable is used more than once in a higher-order call, each occurrence should get its diff --git a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md index 3ef5f1a90e0609..cd43054f86429b 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md @@ -1577,6 +1577,42 @@ def _(values: list[Recursive], sink: Callable[[object], None]) -> None: reveal_type(first_recursive(values, sink)) ``` +## Inferring from multiple intersection arguments + +Each argument below satisfies `Source[T]` in two ways. Combining independent alternatives must +remain bounded, and the merged inference result retains evidence from all four arguments. Reordering +the arguments does not change that result. + +```py +from typing import assert_type +from ty_extensions import Intersection + +class Source[T]: + def get(self) -> T: + raise NotImplementedError + +class A: ... +class B: ... +class C: ... +class D: ... +class E: ... +class F: ... +class G: ... +class H: ... + +def first[T](a: Source[T], b: Source[T], c: Source[T], d: Source[T]) -> T: + return a.get() + +def _( + a: Intersection[Source[A], Source[B]], + b: Intersection[Source[C], Source[D]], + c: Intersection[Source[E], Source[F]], + d: Intersection[Source[G], Source[H]], +) -> None: + assert_type(first(a, b, c, d), A | B | C | D | E | F | G | H) + assert_type(first(d, c, b, a), A | B | C | D | E | F | G | H) +``` + ## Typevars in a union ```py diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index 95f906dc6a5c1a..09469890d25972 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/crates/ty_python_semantic/src/types/call/bind.rs @@ -3060,8 +3060,8 @@ impl<'db> Bindings<'db> { let constraints = ConstraintSetBuilder::new(); let set = constraints.load(db, env, tracked.constraints(db)); - let result = match set.solutions(db, env, &constraints, inferable) { - Solutions::Constrained(paths) => Type::heterogeneous_tuple( + let result = match set.solutions(db, env, inferable) { + Ok(Solutions::Constrained(paths)) => Type::heterogeneous_tuple( db, env, paths.into_vec().into_iter().map(|path| { @@ -3074,8 +3074,9 @@ impl<'db> Bindings<'db> { )) }), ), - Solutions::Unsatisfiable => Type::none(db, env), - Solutions::Unconstrained => Type::empty_tuple(db, env), + Ok(Solutions::Unsatisfiable) => Type::none(db, env), + Ok(Solutions::Unconstrained) => Type::empty_tuple(db, env), + Err(_) => Type::unknown(), }; overload.set_return_type(result); } @@ -3097,8 +3098,8 @@ impl<'db> Bindings<'db> { let constraints = ConstraintSetBuilder::new(); let set = constraints.load(db, env, tracked.constraints(db)); - let result = match set.solutions(db, env, &constraints, inferable) { - Solutions::Constrained(paths) => Type::heterogeneous_tuple( + let result = match set.solutions(db, env, inferable) { + Ok(Solutions::Constrained(paths)) => Type::heterogeneous_tuple( db, env, paths.into_vec().into_iter().map(|path| { @@ -3110,8 +3111,9 @@ impl<'db> Bindings<'db> { )) }), ), - Solutions::Unsatisfiable => Type::none(db, env), - Solutions::Unconstrained => Type::empty_tuple(db, env), + Ok(Solutions::Unsatisfiable) => Type::none(db, env), + Ok(Solutions::Unconstrained) => Type::empty_tuple(db, env), + Err(_) => Type::unknown(), }; overload.set_return_type(result); } @@ -5859,7 +5861,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> { self.inferable_typevars, ); - // Use `solutions_with` to determine per-typevar variance from the raw + // Use `solve_with` to determine per-typevar variance from the raw // lower/upper bounds on each BDD path. let mut variance_map: FxHashMap, TypeVarVariance> = FxHashMap::default(); diff --git a/crates/ty_python_semantic/src/types/constraints.rs b/crates/ty_python_semantic/src/types/constraints.rs index a606eae4bd7918..3ffb1fa4aada7e 100644 --- a/crates/ty_python_semantic/src/types/constraints.rs +++ b/crates/ty_python_semantic/src/types/constraints.rs @@ -106,6 +106,7 @@ use ty_python_core::rank::RankBitBox; use ty_static::EnvVars; use crate::types::class::GenericAlias; +use crate::types::constraints::projection::{ProjectionError, SolutionBudget}; use crate::types::constraints::support::{Support, SupportId}; use crate::types::typevar::{BoundTypeVarIdentity, TypeVarInstance, TypeVarSet}; use crate::types::variance::VarianceInferable; @@ -119,6 +120,7 @@ use crate::types::{ }; use crate::{Db, FxIndexMap, FxIndexSet, FxOrderSet, ProgramEnvironment}; +pub(crate) mod projection; mod solutions; mod support; @@ -840,46 +842,6 @@ impl<'db, 'c> ConstraintSet<'db, 'c> { .negate(db, builder) } - /// Computes default solutions for each BDD path. - pub(crate) fn solutions( - self, - db: &'db dyn Db, - env: &ProgramEnvironment<'db>, - builder: &'c ConstraintSetBuilder<'db>, - inferable: TypeVarSet<'db>, - ) -> Solutions<'db> { - self.solutions_with(db, env, builder, inferable, |_variance, path_bound| { - PathBounds::default_solve(db, env, builder, path_bound) - }) - } - - /// Computes solutions using a caller-provided selector for each typevar on each BDD path. - /// - /// The selector receives the typevar's variance and explicit lower and upper bounds. Its - /// outcome distinguishes missing evidence, invalid paths, and exhausted solution budgets. - /// The caller is responsible for combining the resulting paths (typically via union). - pub(crate) fn solutions_with( - self, - db: &'db dyn Db, - env: &ProgramEnvironment<'db>, - builder: &'c ConstraintSetBuilder<'db>, - inferable: TypeVarSet<'db>, - choose: impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, - ) -> Solutions<'db> { - self.verify_builder(builder); - let mut storage = builder.storage.borrow_mut(); - let path_bounds = PathBounds::compute( - db, - env, - &mut storage, - self.node, - inferable, - self.source_order, - ); - drop(storage); - path_bounds.solve_with(choose) - } - pub(crate) fn display( self, db: &'db dyn Db, @@ -3230,19 +3192,20 @@ impl NodeId { result } - fn remove_noninferable<'db>( + fn remove_noninferable<'db, L: SolutionLimits>( self, db: &'db dyn Db, env: &ProgramEnvironment<'db>, storage: &mut ConstraintSetStorage<'db>, inferable: TypeVarSet<'db>, source_order: Option, - ) -> (Self, Option) { + limits: &mut L, + ) -> ControlFlow)> { match self.node() { - Node::AlwaysTrue => (ALWAYS_TRUE, None), - Node::AlwaysFalse => (ALWAYS_FALSE, None), + Node::AlwaysTrue => ControlFlow::Continue((ALWAYS_TRUE, None)), + Node::AlwaysFalse => ControlFlow::Continue((ALWAYS_FALSE, None)), Node::Interior(interior) => { - interior.remove_noninferable(db, env, storage, inferable, source_order) + interior.remove_noninferable(db, env, storage, inferable, source_order, limits) } } } @@ -3743,6 +3706,50 @@ pub(crate) enum PathBounds<'db> { Constrained(Box<[Box<[PathBound<'db>]>]>), } +/// Limits shared by the preprocessing and collection walks used to extract solutions. +trait SolutionLimits { + type Break; + + fn visit_node(&mut self) -> ControlFlow { + ControlFlow::Continue(()) + } + + fn satisfied_path(&mut self) -> ControlFlow { + ControlFlow::Continue(()) + } +} + +struct UnboundedSolutionLimits; + +impl SolutionLimits for UnboundedSolutionLimits { + type Break = Infallible; +} + +struct BoundedSolutionLimits { + remaining_paths: usize, + remaining_visits: usize, +} + +impl SolutionLimits for BoundedSolutionLimits { + type Break = ProjectionError; + + fn visit_node(&mut self) -> ControlFlow { + let Some(remaining) = self.remaining_visits.checked_sub(1) else { + return ControlFlow::Break(ProjectionError::TraversalBudgetExceeded); + }; + self.remaining_visits = remaining; + ControlFlow::Continue(()) + } + + fn satisfied_path(&mut self) -> ControlFlow { + let Some(remaining) = self.remaining_paths.checked_sub(1) else { + return ControlFlow::Break(ProjectionError::PathBudgetExceeded); + }; + self.remaining_paths = remaining; + ControlFlow::Continue(()) + } +} + impl<'db> PathBounds<'db> { /// Computes sorted BDD paths and accumulates per-typevar lower/upper bounds for each path. /// @@ -3756,6 +3763,59 @@ impl<'db> PathBounds<'db> { inferable: TypeVarSet<'db>, source_order: Option, ) -> Self { + let ControlFlow::Continue(result) = Self::compute_with_limits( + db, + env, + storage, + node, + inferable, + source_order, + &mut UnboundedSolutionLimits, + ); + result + } + + /// Computes complete path bounds within limits shared by preprocessing and collection. + /// + /// Visits include the concrete-conjunction fast path and both BDD walks. The path limit + /// counts materialized constrained paths; an unconstrained or unsatisfiable result needs no + /// path allowance. No partially collected family is returned when either limit is exhausted. + fn compute_bounded( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + storage: &mut ConstraintSetStorage<'db>, + node: NodeId, + inferable: TypeVarSet<'db>, + source_order: Option, + budget: SolutionBudget, + ) -> Result { + let mut limits = BoundedSolutionLimits { + remaining_paths: budget.paths, + remaining_visits: budget.visits, + }; + match Self::compute_with_limits( + db, + env, + storage, + node, + inferable, + source_order, + &mut limits, + ) { + ControlFlow::Continue(result) => Ok(result), + ControlFlow::Break(error) => Err(error), + } + } + + fn compute_with_limits( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + storage: &mut ConstraintSetStorage<'db>, + node: NodeId, + inferable: TypeVarSet<'db>, + source_order: Option, + limits: &mut L, + ) -> ControlFlow { let mut source_orders = storage.calculate_source_orders(source_order); if let Some(path_bounds) = Self::compute_simple_bound_conjunction( db, @@ -3764,24 +3824,34 @@ impl<'db> PathBounds<'db> { &source_orders, node, inferable, - ) { - return path_bounds; + limits, + )? { + return ControlFlow::Continue(path_bounds); } let (node, derived_source_order) = - node.remove_noninferable(db, env, storage, inferable, source_order); + node.remove_noninferable(db, env, storage, inferable, source_order, limits)?; source_orders.extend(storage.calculate_source_orders(derived_source_order)); let interior = match node.node() { - Node::AlwaysTrue => return PathBounds::Unconstrained, - Node::AlwaysFalse => return PathBounds::Unsatisfiable, + Node::AlwaysTrue => { + limits.visit_node()?; + return ControlFlow::Continue(PathBounds::Unconstrained); + } + Node::AlwaysFalse => { + limits.visit_node()?; + return ControlFlow::Continue(PathBounds::Unsatisfiable); + } Node::Interior(interior) => interior, }; let mut walker = SolutionWalker::new(source_orders); + // Sequent discovery must also happen in source order. Sorting the collected paths is + // too late: sequent pairs are not commutative, and TDD traversal order can otherwise + // discard gradual evidence before solution extraction. let path_source_order = storage.ordered_source_order(source_order, derived_source_order); let mut path = interior.path_assignments(db, env, storage, path_source_order); - walker.visit_node(db, env, storage, &mut path, node); - walker.finish(db, env, storage) + walker.visit_node(db, env, storage, &mut path, node, limits)?; + ControlFlow::Continue(walker.finish(db, env, storage)) } /// Accumulates a conjunction of concrete bound constraints without constructing a @@ -3790,41 +3860,47 @@ impl<'db> PathBounds<'db> { /// There are no relationships to derive between these constraints, as the upper and lower /// bounds do not contain typevars. The normal solution-selection logic still validates each /// accumulated bound against the typevar's declared bound or constraints. - fn compute_simple_bound_conjunction( + fn compute_simple_bound_conjunction( db: &'db dyn Db, env: &ProgramEnvironment<'db>, storage: &mut ConstraintSetStorage<'db>, source_orders: &FxIndexSet, node: NodeId, inferable: TypeVarSet<'db>, - ) -> Option { - match node.node() { - Node::AlwaysTrue => return Some(PathBounds::Unconstrained), - Node::AlwaysFalse => return Some(PathBounds::Unsatisfiable), - Node::Interior(_) => {} - } - + limits: &mut L, + ) -> ControlFlow> { let mut constraints = Vec::default(); let mut current = node; loop { + limits.visit_node()?; match current.node() { - Node::AlwaysTrue => break, - Node::AlwaysFalse => return None, + Node::AlwaysTrue => { + if constraints.is_empty() { + return ControlFlow::Continue(Some(PathBounds::Unconstrained)); + } + limits.satisfied_path()?; + break; + } + Node::AlwaysFalse => { + return ControlFlow::Continue( + constraints.is_empty().then_some(PathBounds::Unsatisfiable), + ); + } Node::Interior(_) => { let interior = storage.interior_node_data(current); if interior.if_uncertain != ALWAYS_FALSE || interior.if_false != ALWAYS_FALSE { - return None; + return ControlFlow::Continue(None); } let constraint = storage.constraint_data(interior.constraint); if !constraint.typevar.is_inferable(db, inferable) { - return None; + return ControlFlow::Continue(None); } if iter::chain(constraint.bounds.lower, constraint.bounds.upper).any(|bound| { bound.has_typevar(db, env) || bound.has_unspecialized_type_var(db, env) }) { - return None; + return ControlFlow::Continue(None); } current = interior.if_true; @@ -3856,7 +3932,7 @@ impl<'db> PathBounds<'db> { .drain(..) .map(|(bound_typevar, bounds)| bounds.finish(db, env, bound_typevar)) .collect(); - Some(PathBounds::Constrained(Box::new([path]))) + ControlFlow::Continue(Some(PathBounds::Constrained(Box::new([path])))) } pub(crate) fn solve( @@ -3876,50 +3952,72 @@ impl<'db> PathBounds<'db> { /// the path's available bindings, but marks the resulting path family as incomplete. pub(crate) fn solve_with( &self, - mut choose: impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, + choose: impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, ) -> Solutions<'db> { + let Ok(solutions) = self.try_solve_with(choose, |_| Ok::<(), Infallible>(())); + solutions + } + + /// Checks each retained solution before collecting it or solving the next path. + fn try_solve_with( + &self, + mut choose: impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, + mut check_solution: impl FnMut(&Solution<'db>) -> Result<(), E>, + ) -> Result, E> { let paths = match self { - PathBounds::Unsatisfiable => return Solutions::Unsatisfiable, - PathBounds::Unconstrained => return Solutions::Unconstrained, + PathBounds::Unsatisfiable => return Ok(Solutions::Unsatisfiable), + PathBounds::Unconstrained => return Ok(Solutions::Unconstrained), PathBounds::Constrained(paths) => paths, }; let mut solutions = Vec::with_capacity(paths.len()); let mut exceeded_budget = false; - 'paths: for path in paths { - let mut solution = Vec::with_capacity(path.len()); - let mut path_exceeded_budget = false; - for path_bound in path { - let variance = path_bound.variance(); - - let ty = match choose(variance, path_bound) { - PathBoundSolution::Solved(ty) => Some(ty), - PathBoundSolution::Unsolved => None, - PathBoundSolution::Unsatisfiable => continue 'paths, - PathBoundSolution::BudgetExceeded { fallback } => { - path_exceeded_budget = true; - fallback - } - }; - if let Some(ty) = ty { - solution.push(TypeVarSolution { - bound_typevar: path_bound.bound_typevar, - solution: ty, - }); - } - } + for path in paths { + let Some((solution, path_exceeded_budget)) = Self::solve_path_with(path, &mut choose) + else { + continue; + }; + check_solution(&solution)?; exceeded_budget |= path_exceeded_budget; solutions.push(solution); } if solutions.is_empty() { - return Solutions::Unsatisfiable; + return Ok(Solutions::Unsatisfiable); } - Solutions::Constrained(if exceeded_budget { + Ok(Solutions::Constrained(if exceeded_budget { SolutionPaths::BudgetExceeded(solutions) } else { SolutionPaths::Complete(solutions) - }) + })) + } + + /// Solves one complete path, retaining whether any of its bindings used a fallback. + /// A later unsatisfiable bound rejects the path even if an earlier bound exhausted its budget. + fn solve_path_with( + path: &[PathBound<'db>], + choose: &mut impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, + ) -> Option<(Solution<'db>, bool)> { + let mut solution = Vec::with_capacity(path.len()); + let mut exceeded_budget = false; + for path_bound in path { + let ty = match choose(path_bound.variance(), path_bound) { + PathBoundSolution::Solved(ty) => Some(ty), + PathBoundSolution::Unsolved => None, + PathBoundSolution::Unsatisfiable => return None, + PathBoundSolution::BudgetExceeded { fallback } => { + exceeded_budget = true; + fallback + } + }; + if let Some(ty) = ty { + solution.push(TypeVarSolution { + bound_typevar: path_bound.bound_typevar, + solution: ty, + }); + } + } + Some((solution, exceeded_budget)) } /// The default solution selection logic for a single typevar on a single BDD path. @@ -4337,11 +4435,12 @@ impl InteriorNode { bound_typevars: TypeVarSet<'db>, source_order: Option, ) -> (NodeId, Option) { - self.abstract_inner( + let ControlFlow::Continue(result) = self.abstract_inner( db, env, storage, source_order, + &mut UnboundedSolutionLimits, // Remove any node that constrains one of `bound_typevars`, or that has a lower/upper // bound that mentions one of them. Removed constraints are still added to `path`, so // the sequent map can propagate any derived constraints that do not mention the @@ -4353,17 +4452,19 @@ impl InteriorNode { typevar.is_inferable(db, bound_typevars) }) }, - ) + ); + result } - fn remove_noninferable<'db>( + fn remove_noninferable<'db, L: SolutionLimits>( self, db: &'db dyn Db, env: &ProgramEnvironment<'db>, storage: &mut ConstraintSetStorage<'db>, inferable: TypeVarSet<'db>, source_order: Option, - ) -> (NodeId, Option) { + limits: &mut L, + ) -> ControlFlow)> { let is_bare_inferable_typevar = |ty: Type<'_>| { ty.as_typevar() .is_some_and(|bound_typevar| bound_typevar.is_inferable(db, inferable)) @@ -4373,6 +4474,7 @@ impl InteriorNode { env, storage, source_order, + limits, // We only want to keep constraints on inferable typevars. If the constraint's typevar // is itself inferable, we keep it. We also need to keep some constraints in // non-inferable typevars, if their lower or upper bound is a bare inferable typevar. @@ -4397,16 +4499,18 @@ impl InteriorNode { ) } - fn abstract_inner<'db, F>( + fn abstract_inner<'db, F, L>( self, db: &'db dyn Db, env: &ProgramEnvironment<'db>, storage: &mut ConstraintSetStorage<'db>, source_order: Option, + limits: &mut L, should_remove: F, - ) -> (NodeId, Option) + ) -> ControlFlow)> where F: FnMut(&ConstraintSetStorage<'_>, ConstraintId) -> bool, + L: SolutionLimits, { #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum Disposition { @@ -4414,17 +4518,23 @@ impl InteriorNode { Remove, } - struct AbstractVisitor { + struct AbstractVisitor<'a, F, L> { should_remove: F, + limits: &'a mut L, } - impl PathVisitor for AbstractVisitor + impl PathVisitor for AbstractVisitor<'_, F, L> where F: FnMut(&ConstraintSetStorage<'_>, ConstraintId) -> bool, + L: SolutionLimits, { type Result = (NodeId, Option); type Interior = (Disposition, ConstraintId); - type Break = Infallible; + type Break = L::Break; + + fn visit_node(&mut self) -> ControlFlow { + self.limits.visit_node() + } fn visit_satisfied<'db>( &mut self, @@ -4564,9 +4674,11 @@ impl InteriorNode { } let mut path = self.path_assignments(db, env, storage, source_order); - let mut visitor = AbstractVisitor { should_remove }; - let ControlFlow::Continue(result) = path.visit(db, env, storage, self.node(), &mut visitor); - result + let mut visitor = AbstractVisitor { + should_remove, + limits, + }; + path.visit(db, env, storage, self.node(), &mut visitor) } fn path_assignments<'db>( @@ -6005,6 +6117,12 @@ trait PathVisitor { type Interior; type Break; + /// Called before visiting any interior or terminal node. Returning `Break` prevents the + /// traversal from entering the node or deriving facts from its outgoing edges. + fn visit_node(&mut self) -> ControlFlow { + ControlFlow::Continue(()) + } + /// Called when we reach the end of a satisfied path. `path` will contain all of the /// assignments on this path. The `Result` value that you return will be propagated back up as /// we "unwind" this path. @@ -6412,6 +6530,7 @@ impl PathAssignments { where V: PathVisitor, { + visitor.visit_node()?; match node.node() { Node::AlwaysTrue if negated => visitor.visit_unsatisfied(db, storage, self), Node::AlwaysTrue => visitor.visit_satisfied(db, storage, self), @@ -7125,6 +7244,48 @@ mod tests { class.to_instance(db, &db.program_environment()) } + fn bounded_path_bounds<'db>( + db: &'db TestDb, + set: ConstraintSet<'db, '_>, + inferable: TypeVarSet<'db>, + max_paths: usize, + max_visits: usize, + ) -> Result, ProjectionError> { + PathBounds::compute_bounded( + db, + &db.program_environment(), + &mut set.builder.storage.borrow_mut(), + set.node, + inferable, + set.source_order, + SolutionBudget { + paths: max_paths, + visits: max_visits, + ..SolutionBudget::default() + }, + ) + } + + #[derive(Default)] + struct CountSolutionLimits { + visits: usize, + paths: usize, + } + + impl SolutionLimits for CountSolutionLimits { + type Break = Infallible; + + fn visit_node(&mut self) -> ControlFlow { + self.visits += 1; + ControlFlow::Continue(()) + } + + fn satisfied_path(&mut self) -> ControlFlow { + self.paths += 1; + ControlFlow::Continue(()) + } + } + #[test] fn type_mapping_updates_constraint_bounds() { // (list[U] ≤ T ≤ list[U])[U ↦ int] = (list[int] ≤ T ≤ list[int]) @@ -7504,6 +7665,115 @@ mod tests { } } + #[test] + fn bounded_path_fast_paths_respect_limits() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let builder = ConstraintSetBuilder::new(); + let t = create_typevar(db, "T"); + let inferable = TypeVarSet::from_typevars(db, [t]); + + for (set, expected) in [ + (ConstraintSet::always(&builder), PathBounds::Unconstrained), + (ConstraintSet::never(&builder), PathBounds::Unsatisfiable), + ] { + assert_eq!( + bounded_path_bounds(db, set, inferable, 0, 0), + Err(ProjectionError::TraversalBudgetExceeded) + ); + assert_eq!(bounded_path_bounds(db, set, inferable, 0, 1), Ok(expected)); + } + + let set = create_constraint(db, &builder, t, KnownClass::Int); + let expected = PathBounds::compute( + db, + &env, + &mut builder.storage.borrow_mut(), + set.node, + inferable, + set.source_order, + ); + assert_eq!( + bounded_path_bounds(db, set, inferable, 0, 2), + Err(ProjectionError::PathBudgetExceeded) + ); + assert_eq!( + bounded_path_bounds(db, set, inferable, 1, 1), + Err(ProjectionError::TraversalBudgetExceeded) + ); + assert_eq!(bounded_path_bounds(db, set, inferable, 1, 2), Ok(expected)); + } + + #[test] + fn bounded_path_collection_shares_preprocessing_visits() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let builder = ConstraintSetBuilder::new(); + let t = create_typevar(db, "T"); + let hidden = create_typevar(db, "Hidden"); + let visible = create_constraint(db, &builder, t, KnownClass::Int); + let hidden_alternatives = + create_constraint(db, &builder, hidden, KnownClass::Str).or(db, &builder, || { + create_constraint(db, &builder, hidden, KnownClass::Bytes) + }); + let set = visible.and(db, &builder, || hidden_alternatives); + let inferable = TypeVarSet::from_typevars(db, [t]); + let mut storage = builder.storage.borrow_mut(); + let source_orders = storage.calculate_source_orders(set.source_order); + let mut preprocessing = CountSolutionLimits::default(); + let ControlFlow::Continue(fast_path) = PathBounds::compute_simple_bound_conjunction( + db, + &env, + &mut storage, + &source_orders, + set.node, + inferable, + &mut preprocessing, + ); + assert_eq!(fast_path, None); + let ControlFlow::Continue(_) = set.node.remove_noninferable( + db, + &env, + &mut storage, + inferable, + set.source_order, + &mut preprocessing, + ); + + let mut complete = CountSolutionLimits::default(); + let ControlFlow::Continue(expected) = PathBounds::compute_with_limits( + db, + &env, + &mut storage, + set.node, + inferable, + set.source_order, + &mut complete, + ); + assert_eq!(complete.paths, 1); + assert!(complete.visits > preprocessing.visits); + drop(storage); + + assert_eq!( + bounded_path_bounds(db, set, inferable, 1, preprocessing.visits), + Err(ProjectionError::TraversalBudgetExceeded) + ); + assert_eq!( + bounded_path_bounds(db, set, inferable, 1, complete.visits - 1), + Err(ProjectionError::TraversalBudgetExceeded) + ); + assert_eq!( + bounded_path_bounds(db, set, inferable, 1, complete.visits), + Ok(expected) + ); + assert_eq!( + bounded_path_bounds(db, set, inferable, 0, complete.visits), + Err(ProjectionError::PathBudgetExceeded) + ); + } + #[test] fn simple_lower_bound_conjunction_skips_sequent_analysis() { let db = setup_db(); @@ -7527,13 +7797,15 @@ mod tests { ) }; - let solutions = set.solutions(db, &env, &builder, inferable); + let solutions = set.solutions(db, &env, inferable); assert_eq!( solutions, - Solutions::Constrained(SolutionPaths::Complete(vec![vec![TypeVarSolution { - bound_typevar: t, - solution: UnionType::from_elements(db, &env, [int, str]), - }]])) + Ok(Solutions::Constrained(SolutionPaths::Complete(vec![vec![ + TypeVarSolution { + bound_typevar: t, + solution: UnionType::from_elements(db, &env, [int, str]), + } + ]]))) ); let storage = builder.storage.borrow(); @@ -7564,7 +7836,7 @@ mod tests { ) }; - let Solutions::Constrained(solutions) = set.solutions(db, &env, &builder, inferable) else { + let Ok(Solutions::Constrained(solutions)) = set.solutions(db, &env, inferable) else { panic!("expected constrained solutions"); }; let solutions = solutions.into_vec(); @@ -7608,8 +7880,8 @@ mod tests { }; assert_eq!( - set.solutions(db, &env, &builder, inferable), - Solutions::Unsatisfiable + set.solutions(db, &env, inferable), + Ok(Solutions::Unsatisfiable) ); let storage = builder.storage.borrow(); @@ -8147,9 +8419,9 @@ class E: ... drop(storage); let set = ConstraintSet::from_node(&builder, node, source_order); - let solutions = set.solutions(db, &env, &builder, inferable); + let solutions = set.solutions(db, &env, inferable); let mut merged = FxHashMap::default(); - if let Solutions::Constrained(paths) = &solutions { + if let Ok(Solutions::Constrained(paths)) = &solutions { for path in paths.as_slice() { for binding in path { merged @@ -8179,9 +8451,9 @@ class E: ... }) .join(", "); let paths = match &solutions { - Solutions::Unsatisfiable => String::from("unsatisfiable"), - Solutions::Unconstrained => String::from("unconstrained"), - Solutions::Constrained(paths) => paths + Ok(Solutions::Unsatisfiable) => String::from("unsatisfiable"), + Ok(Solutions::Unconstrained) => String::from("unconstrained"), + Ok(Solutions::Constrained(paths)) => paths .as_slice() .iter() .map(|path| { @@ -8196,6 +8468,7 @@ class E: ... .join(", ") }) .join("; "), + Err(error) => format!("error: {error:?}"), }; signatures.insert(format!( "never={} always={} merged=[{merged}] paths=[{paths}]", @@ -8928,6 +9201,56 @@ class E: ... } } + #[test] + fn solution_walker_break_restores_path_assignments() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let t = create_typevar(db, "T"); + let builder = ConstraintSetBuilder::new(); + let t_int = create_constraint(db, &builder, t, KnownClass::Int); + let t_str = create_constraint(db, &builder, t, KnownClass::Str); + let set = t_int.or(db, &builder, || t_str); + let source_orders = builder + .storage + .borrow() + .calculate_source_orders(set.source_order); + let expected = PathBounds::compute( + db, + &env, + &mut builder.storage.borrow_mut(), + set.node, + TypeVarSet::from_typevars(db, [t]), + set.source_order, + ); + + // Both limits interrupt an edge with path-local assignments: the visit limit stops + // below the root, and the path limit stops after collecting the first alternative. + for (remaining_paths, remaining_visits, error) in [ + (usize::MAX, 1, ProjectionError::TraversalBudgetExceeded), + (1, usize::MAX, ProjectionError::PathBudgetExceeded), + ] { + let mut path = path_assignments_for(db, &env, &builder, set.node, set.source_order); + let mut storage = builder.storage.borrow_mut(); + let mut limits = BoundedSolutionLimits { + remaining_paths, + remaining_visits, + }; + let mut walker = SolutionWalker::new(source_orders.clone()); + assert_eq!( + walker.visit_node(db, &env, &mut storage, &mut path, set.node, &mut limits), + ControlFlow::Break(error) + ); + drop(walker); + + let mut limits = UnboundedSolutionLimits; + let mut walker = SolutionWalker::new(source_orders.clone()); + let ControlFlow::Continue(()) = + walker.visit_node(db, &env, &mut storage, &mut path, set.node, &mut limits); + assert_eq!(walker.finish(db, &env, &mut storage), expected); + } + } + /// Double negation of a TDD with uncertain branches is semantically equivalent to the /// original (though the structure may differ since negation produces flat TDDs). #[test] diff --git a/crates/ty_python_semantic/src/types/constraints/projection.rs b/crates/ty_python_semantic/src/types/constraints/projection.rs new file mode 100644 index 00000000000000..86ad764a680724 --- /dev/null +++ b/crates/ty_python_semantic/src/types/constraints/projection.rs @@ -0,0 +1,252 @@ +//! Bounded projections of correlated constraint solutions. + +use rustc_hash::FxHashSet; + +use super::{ConstraintSet, PathBound, PathBoundSolution, PathBounds, Solutions, TypeVarSolution}; +use crate::types::typevar::TypeVarSet; +use crate::types::{Type, TypeVarVariance}; +use crate::{Db, ProgramEnvironment}; + +/// Limits for one projection, including preprocessing, path collection, and its result. +#[derive(Clone, Copy, Debug)] +pub(crate) struct SolutionBudget { + /// Satisfied paths collected before per-variable solution selection can reject them. + pub(crate) paths: usize, + /// Interior and terminal visits, shared by preprocessing and path collection. + pub(crate) visits: usize, + /// Set-theoretic terms contributed to the result, including terms exposed by aliases. + pub(crate) type_terms: usize, +} + +impl Default for SolutionBudget { + fn default() -> Self { + // Allow long, simple conjunctions and sizable existing unions without allowing their + // alternatives to expand into an equally large family of specializations. + Self { + paths: 4_096, + visits: 32_768, + type_terms: 8_192, + } + } +} + +/// Why an exact projection could not be completed. +/// +/// None of these outcomes proves that the constraint set is unsatisfiable. In particular, a +/// caller must not use the prefix visited before a limit was reached as the complete answer. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ProjectionError { + PathBudgetExceeded, + TraversalBudgetExceeded, + TypeBudgetExceeded, + IncompleteSolution, +} + +/// An exact projection of all retained solution paths. +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum SolutionProjection { + Unsatisfiable, + Unconstrained, + Constrained(T), +} + +/// A shared limit on the type terms consumed while constructing a projection. +/// +/// Union projections charge each contribution before adding it. An intersection projection must +/// additionally use `IntersectionType::bounded_from_elements`, since distributing intersections +/// over unions can multiply, rather than add, the number of terms. +pub(crate) struct ProjectionTypeBudget { + remaining: usize, +} + +impl ProjectionTypeBudget { + fn new(remaining: usize) -> Self { + Self { remaining } + } + + /// Charges the set-theoretic terms that a type constructor may flatten or inspect. Aliases + /// are included so a large union cannot evade the limit by being hidden behind a name. + pub(crate) fn charge_type<'db>( + &mut self, + db: &'db dyn Db, + ty: Type<'db>, + ) -> Result<(), ProjectionError> { + self.charge_type_inner(db, ty, &mut FxHashSet::default()) + } + + fn charge_type_inner<'db>( + &mut self, + db: &'db dyn Db, + ty: Type<'db>, + seen_aliases: &mut FxHashSet>, + ) -> Result<(), ProjectionError> { + self.remaining = self + .remaining + .checked_sub(1) + .ok_or(ProjectionError::TypeBudgetExceeded)?; + match ty { + Type::Union(union) => { + for element in union.elements(db) { + self.charge_type_inner(db, *element, seen_aliases)?; + } + } + Type::Intersection(intersection) => { + for element in intersection + .iter_positive(db) + .chain(intersection.iter_negative(db)) + { + self.charge_type_inner(db, element, seen_aliases)?; + } + } + Type::TypeAlias(alias) if seen_aliases.insert(ty) => { + self.charge_type_inner(db, alias.value_type(db), seen_aliases)?; + } + _ => {} + } + Ok(()) + } +} + +impl<'db> ConstraintSet<'db, '_> { + /// Computes default solutions for each BDD path within the default projection budget. + pub(crate) fn solutions( + self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + inferable: TypeVarSet<'db>, + ) -> Result, ProjectionError> { + let builder = self.builder; + self.solutions_with( + db, + env, + inferable, + SolutionBudget::default(), + |_variance, path_bound| PathBounds::default_solve(db, env, builder, path_bound), + ) + } + + fn bounded_path_bounds( + self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + inferable: TypeVarSet<'db>, + budget: SolutionBudget, + ) -> Result, ProjectionError> { + PathBounds::compute_bounded( + db, + env, + &mut self.builder.storage.borrow_mut(), + self.node, + inferable, + self.source_order, + budget, + ) + } + + /// Computes solutions using a caller-provided selector within the given projection budget. + /// + /// The selector receives the typevar's variance and explicit lower and upper bounds. Its + /// outcome distinguishes missing evidence, invalid paths, and exhausted solution budgets. + /// The caller is responsible for combining the resulting paths (typically via union). + /// + /// Per-variable budget exhaustion preserves available fallback bindings and marks the path + /// family as [`SolutionPaths::BudgetExceeded`](super::SolutionPaths::BudgetExceeded). + /// Exhausting a limit in the supplied [`SolutionBudget`] instead returns an error without a + /// partial path family. + pub(crate) fn solutions_with( + self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + inferable: TypeVarSet<'db>, + budget: SolutionBudget, + choose: impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, + ) -> Result, ProjectionError> { + let path_bounds = self.bounded_path_bounds(db, env, inferable, budget)?; + let mut type_budget = ProjectionTypeBudget::new(budget.type_terms); + path_bounds.try_solve_with(choose, |solution| { + for binding in solution { + type_budget.charge_type(db, binding.solution)?; + } + Ok(()) + }) + } + + /// Folds complete, correlated solutions without first allocating every solved path. + /// + /// Raw paths are collected within the traversal limits and sorted in the same source order + /// as [`Self::solutions_with`]. The storage borrow is released before invoking either + /// callback, so they can safely use the constraint builder. Each call to `fold` receives the + /// complete bindings for one retained path, including an empty slice for a valid path on + /// which no variable was solved. + /// + /// The accumulator is returned only if the entire projection succeeds. `fold` must charge + /// newly accumulated types to its supplied budget and use bounded constructors for operations + /// that can expand them. It should combine alternatives commutatively when their order is not + /// meaningful to its consumer. Existing limitations in solution extraction still apply; this + /// API does not make an order-sensitive selector or fold order-independent. + #[expect(clippy::too_many_arguments)] + pub(crate) fn try_fold_solutions( + self, + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + inferable: TypeVarSet<'db>, + budget: SolutionBudget, + choose: impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, + initial: T, + fold: impl FnMut( + T, + &[TypeVarSolution<'db>], + &mut ProjectionTypeBudget, + ) -> Result, + ) -> Result, ProjectionError> { + let path_bounds = self.bounded_path_bounds(db, env, inferable, budget)?; + + path_bounds.try_fold_with( + choose, + initial, + &mut ProjectionTypeBudget::new(budget.type_terms), + fold, + ) + } +} + +impl<'db> PathBounds<'db> { + fn try_fold_with( + &self, + mut choose: impl FnMut(TypeVarVariance, &PathBound<'db>) -> PathBoundSolution<'db>, + mut accumulated: T, + budget: &mut ProjectionTypeBudget, + mut fold: impl FnMut( + T, + &[TypeVarSolution<'db>], + &mut ProjectionTypeBudget, + ) -> Result, + ) -> Result, ProjectionError> { + let paths = match self { + Self::Unsatisfiable => return Ok(SolutionProjection::Unsatisfiable), + Self::Unconstrained => return Ok(SolutionProjection::Unconstrained), + Self::Constrained(paths) => paths, + }; + + let mut retained = false; + for path in paths { + let Some((solution, incomplete)) = Self::solve_path_with(path, &mut choose) else { + continue; + }; + if incomplete { + return Err(ProjectionError::IncompleteSolution); + } + accumulated = fold(accumulated, &solution, budget)?; + retained = true; + } + + Ok(if retained { + SolutionProjection::Constrained(accumulated) + } else { + SolutionProjection::Unsatisfiable + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/ty_python_semantic/src/types/constraints/projection/tests.rs b/crates/ty_python_semantic/src/types/constraints/projection/tests.rs new file mode 100644 index 00000000000000..76fbfea965eb93 --- /dev/null +++ b/crates/ty_python_semantic/src/types/constraints/projection/tests.rs @@ -0,0 +1,633 @@ +use itertools::Itertools; +use ruff_db::files::system_path_to_file; +use ruff_db::system::DbWithWritableSystem; +use ruff_python_ast::name::Name; +use rustc_hash::FxHashSet; +use ty_python_core::ProgramFile; + +use super::{ProjectionError, ProjectionTypeBudget, SolutionBudget, SolutionProjection}; +use crate::db::tests::{TestDb, setup_db}; +use crate::place::global_symbol; +use crate::types::constraints::{ + ConstraintSet, ConstraintSetBuilder, IteratorConstraintsExtension, PathBound, + PathBoundSolution, PathBounds, Solution, SolutionPaths, Solutions, TypeVarSolution, +}; +use crate::types::typevar::TypeVarSet; +use crate::types::{ + BoundTypeVarInstance, IntersectionType, KnownClass, Type, TypeVarVariance, UnionType, +}; + +type Paths<'db> = FxHashSet>; + +fn create_typevar<'db>(db: &'db TestDb, name: &'static str) -> BoundTypeVarInstance<'db> { + BoundTypeVarInstance::synthetic( + db, + &db.program_environment(), + Name::new_static(name), + TypeVarVariance::Invariant, + ) +} + +fn known_instance(db: &TestDb, class: KnownClass) -> Type<'_> { + class.to_instance(db, &db.program_environment()) +} + +fn exact<'db, 'c>( + db: &'db TestDb, + builder: &'c ConstraintSetBuilder<'db>, + typevar: BoundTypeVarInstance<'db>, + ty: Type<'db>, +) -> ConstraintSet<'db, 'c> { + ConstraintSet::constrain_typevar(db, &db.program_environment(), builder, typevar, ty, ty) +} + +fn binary_choice<'db, 'c>( + db: &'db TestDb, + builder: &'c ConstraintSetBuilder<'db>, + typevar: BoundTypeVarInstance<'db>, + alternatives: [Type<'db>; 2], +) -> ConstraintSet<'db, 'c> { + alternatives + .into_iter() + .when_any(db, builder, |ty| exact(db, builder, typevar, ty)) +} + +fn binding<'db>( + bound_typevar: BoundTypeVarInstance<'db>, + solution: Type<'db>, +) -> TypeVarSolution<'db> { + TypeVarSolution { + bound_typevar, + solution, + } +} + +fn collect_paths<'db, 'c>( + db: &'db TestDb, + builder: &'c ConstraintSetBuilder<'db>, + set: ConstraintSet<'db, 'c>, + typevars: &[BoundTypeVarInstance<'db>], + budget: SolutionBudget, +) -> Result>, ProjectionError> { + let env = db.program_environment(); + set.try_fold_solutions( + db, + &env, + TypeVarSet::from_typevars(db, typevars.iter().copied()), + budget, + |_, bound| PathBounds::default_solve(db, &env, builder, bound), + Paths::default(), + |mut paths, path, budget| { + for binding in path { + budget.charge_type(db, binding.solution)?; + } + let mut path = path.to_vec(); + path.sort_by_key(|binding| { + typevars + .iter() + .position(|typevar| *typevar == binding.bound_typevar) + }); + paths.insert(path); + Ok(paths) + }, + ) +} + +#[test] +fn path_limit_is_checked_before_solving() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let t = create_typevar(db, "T"); + let u = create_typevar(db, "U"); + let int = known_instance(db, KnownClass::Int); + let str = known_instance(db, KnownClass::Str); + let builder = ConstraintSetBuilder::new(); + let set = binary_choice(db, &builder, t, [int, str]) + .and(db, &builder, || binary_choice(db, &builder, u, [int, str])); + let inferable = TypeVarSet::from_typevars(db, [t, u]); + + for max_paths in [0, 3, 4] { + let mut selected = 0; + let mut folded = 0; + let result = set.try_fold_solutions( + db, + &env, + inferable, + SolutionBudget { + paths: max_paths, + ..SolutionBudget::default() + }, + |_, bound| { + selected += 1; + PathBounds::default_solve(db, &env, &builder, bound) + }, + 0, + |count, _, _| { + folded += 1; + Ok(count + 1) + }, + ); + + if max_paths < 4 { + assert_eq!(result, Err(ProjectionError::PathBudgetExceeded)); + assert_eq!(selected, 0); + assert_eq!(folded, 0); + } else { + assert_eq!(result, Ok(SolutionProjection::Constrained(4))); + assert_eq!(selected, 8); + } + } +} + +#[test] +fn terminal_projections_need_no_paths_or_types() { + let db = setup_db(); + let db = &db; + let t = create_typevar(db, "T"); + let builder = ConstraintSetBuilder::new(); + + // Terminal answers do not allocate any path or construct any type. + let terminal_budget = SolutionBudget { + paths: 0, + visits: 1, + type_terms: 0, + }; + for (set, expected) in [ + ( + ConstraintSet::always(&builder), + SolutionProjection::Unconstrained, + ), + ( + ConstraintSet::never(&builder), + SolutionProjection::Unsatisfiable, + ), + ] { + assert_eq!( + collect_paths(db, &builder, set, &[t], terminal_budget), + Ok(expected) + ); + } +} + +#[test] +fn source_and_interning_order_do_not_change_correlated_projection() { + let db = setup_db(); + let db = &db; + let t = create_typevar(db, "T"); + let u = create_typevar(db, "U"); + let int = known_instance(db, KnownClass::Int); + let str = known_instance(db, KnownClass::Str); + let bytes = known_instance(db, KnownClass::Bytes); + let bool = known_instance(db, KnownClass::Bool); + let atoms = [(t, int), (t, str), (u, bytes), (u, bool)]; + // These alternatives do not admit the crossed pairings of T and U. + let expected = FxHashSet::from_iter([ + vec![binding(t, int), binding(u, bytes)], + vec![binding(t, str), binding(u, bool)], + ]); + + for interning_order in (0..atoms.len()).permutations(atoms.len()) { + for reverse_source in [false, true] { + let builder = ConstraintSetBuilder::new(); + for index in &interning_order { + let (typevar, ty) = atoms[*index]; + exact(db, &builder, typevar, ty); + } + let [t_int, t_str, u_bytes, u_bool] = + atoms.map(|(typevar, ty)| exact(db, &builder, typevar, ty)); + let set = if reverse_source { + u_bool + .and(db, &builder, || t_str) + .or(db, &builder, || u_bytes.and(db, &builder, || t_int)) + } else { + t_int + .and(db, &builder, || u_bytes) + .or(db, &builder, || t_str.and(db, &builder, || u_bool)) + }; + + assert_eq!( + collect_paths(db, &builder, set, &[t, u], SolutionBudget::default()), + Ok(SolutionProjection::Constrained(expected.clone())), + "interning order {interning_order:?}, reverse source {reverse_source}" + ); + assert_eq!( + collect_paths( + db, + &builder, + set, + &[t, u], + SolutionBudget { + paths: 1, + ..SolutionBudget::default() + }, + ), + Err(ProjectionError::PathBudgetExceeded) + ); + } + } +} + +#[test] +fn four_independent_binary_arguments_have_sixteen_solutions() { + let db = setup_db(); + let db = &db; + let typevars = ["T", "U", "V", "W"].map(|name| create_typevar(db, name)); + let alternatives = + [[1, 2], [3, 4], [5, 6], [7, 8]].map(|choices| choices.map(Type::int_literal)); + let builder = ConstraintSetBuilder::new(); + + // Four arguments that independently admit two specializations produce sixteen whole-call + // solutions. The limit applies before constructing any of their projected return types. + let set = + typevars + .into_iter() + .zip(alternatives) + .when_all(db, &builder, |(typevar, alternatives)| { + binary_choice(db, &builder, typevar, alternatives) + }); + let expected = alternatives + .into_iter() + .multi_cartesian_product() + .map(|choices| { + typevars + .into_iter() + .zip(choices) + .map(|(typevar, ty)| binding(typevar, ty)) + .collect() + }) + .collect(); + + assert_eq!( + collect_paths( + db, + &builder, + set, + &typevars, + SolutionBudget { + paths: 16, + ..SolutionBudget::default() + }, + ), + Ok(SolutionProjection::Constrained(expected)) + ); + assert_eq!( + collect_paths( + db, + &builder, + set, + &typevars, + SolutionBudget { + paths: 15, + ..SolutionBudget::default() + }, + ), + Err(ProjectionError::PathBudgetExceeded) + ); +} + +#[test] +fn incomplete_solution_discards_the_projection() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let t = create_typevar(db, "T"); + let int = known_instance(db, KnownClass::Int); + let str = known_instance(db, KnownClass::Str); + let builder = ConstraintSetBuilder::new(); + let inferable = TypeVarSet::from_typevars(db, [t]); + let budget = SolutionBudget { + type_terms: 2, + ..SolutionBudget::default() + }; + + for alternatives in [[int, str], [str, int]] { + let set = binary_choice(db, &builder, t, alternatives); + let choose = |_, bound: &PathBound<'_>| { + if bound.lower == Some(str) { + PathBoundSolution::BudgetExceeded { + fallback: Some(str), + } + } else { + PathBoundSolution::Solved(int) + } + }; + + assert_eq!( + set.solutions_with(db, &env, inferable, budget, choose), + Ok(Solutions::Constrained(SolutionPaths::BudgetExceeded( + alternatives.map(|ty| vec![binding(t, ty)]).into() + ))) + ); + assert_eq!( + set.try_fold_solutions(db, &env, inferable, budget, choose, 0, |count, _, _| Ok( + count + 1 + ),), + Err(ProjectionError::IncompleteSolution) + ); + } +} + +#[test] +fn rejected_exhausted_path_does_not_poison_valid_sibling() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let t = create_typevar(db, "T"); + let u = create_typevar(db, "U"); + let int = known_instance(db, KnownClass::Int); + let str = known_instance(db, KnownClass::Str); + let bytes = known_instance(db, KnownClass::Bytes); + let inferable = TypeVarSet::from_typevars(db, [t, u]); + let budget = SolutionBudget { + type_terms: 1, + ..SolutionBudget::default() + }; + + for reverse_bounds in [false, true] { + for reverse_paths in [false, true] { + let builder = ConstraintSetBuilder::new(); + let t_str = exact(db, &builder, t, str); + let u_bytes = exact(db, &builder, u, bytes); + let rejected = if reverse_bounds { + u_bytes.and(db, &builder, || t_str) + } else { + t_str.and(db, &builder, || u_bytes) + }; + let valid = exact(db, &builder, t, int); + let set = if reverse_paths { + valid.or(db, &builder, || rejected) + } else { + rejected.or(db, &builder, || valid) + }; + + // Only the valid sibling consumes the budget, even when the rejected path had + // already selected a type or retained a fallback before finding its contradiction. + for rejected_binding in [ + PathBoundSolution::Solved(str), + PathBoundSolution::BudgetExceeded { + fallback: Some(str), + }, + ] { + let choose = |_, bound: &PathBound<'_>| { + if bound.bound_typevar == u { + PathBoundSolution::Unsatisfiable + } else if bound.lower == Some(str) { + rejected_binding + } else { + PathBoundSolution::Solved(int) + } + }; + assert_eq!( + set.solutions_with(db, &env, inferable, budget, choose), + Ok(Solutions::Constrained(SolutionPaths::Complete(vec![vec![ + binding(t, int), + ]]))) + ); + assert_eq!( + set.try_fold_solutions( + db, + &env, + inferable, + budget, + choose, + Vec::new(), + |mut paths, path, budget| { + for binding in path { + budget.charge_type(db, binding.solution)?; + } + paths.push(path.to_vec()); + Ok(paths) + }, + ), + Ok(SolutionProjection::Constrained(vec![vec![binding(t, int)]])) + ); + } + } + } +} + +#[test] +fn valid_unsolved_path_is_not_unconstrained() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let t = create_typevar(db, "T"); + let builder = ConstraintSetBuilder::new(); + let set = exact(db, &builder, t, known_instance(db, KnownClass::Int)); + let inferable = TypeVarSet::from_typevars(db, [t]); + let budget = SolutionBudget { + type_terms: 0, + ..SolutionBudget::default() + }; + + for (selected, collected, projected) in [ + ( + PathBoundSolution::Unsolved, + Solutions::Constrained(SolutionPaths::Complete(vec![vec![]])), + Ok(SolutionProjection::Constrained(1)), + ), + ( + PathBoundSolution::BudgetExceeded { fallback: None }, + Solutions::Constrained(SolutionPaths::BudgetExceeded(vec![vec![]])), + Err(ProjectionError::IncompleteSolution), + ), + ( + PathBoundSolution::Unsatisfiable, + Solutions::Unsatisfiable, + Ok(SolutionProjection::Unsatisfiable), + ), + ] { + assert_eq!( + set.solutions_with(db, &env, inferable, budget, |_, _| selected), + Ok(collected) + ); + assert_eq!( + set.try_fold_solutions( + db, + &env, + inferable, + budget, + |_, _| selected, + 0, + |count, path, _| { + assert!(path.is_empty()); + Ok(count + 1) + }, + ), + projected + ); + } +} + +#[test] +fn type_budget_is_charged_before_constructing_a_union() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let t = create_typevar(db, "T"); + let int = known_instance(db, KnownClass::Int); + let str = known_instance(db, KnownClass::Str); + let bytes = known_instance(db, KnownClass::Bytes); + let builder = ConstraintSetBuilder::new(); + let set = binary_choice(db, &builder, t, [int, str]) + .or(db, &builder, || exact(db, &builder, t, bytes)); + let inferable = TypeVarSet::from_typevars(db, [t]); + + for max_type_terms in [0, 1, 2, 3] { + let budget = SolutionBudget { + type_terms: max_type_terms, + ..SolutionBudget::default() + }; + let mut selected = 0; + let collected = set.solutions_with(db, &env, inferable, budget, |_, bound| { + selected += 1; + PathBounds::default_solve(db, &env, &builder, bound) + }); + // One additional path is selected to discover that it exceeds the budget; later + // paths are not solved. + assert_eq!(selected, (max_type_terms + 1).min(3)); + + let mut constructed = 0; + let result = set.try_fold_solutions( + db, + &env, + inferable, + budget, + |_, bound| PathBounds::default_solve(db, &env, &builder, bound), + Type::Never, + |accumulated, path, budget| { + assert_eq!(path.len(), 1); + let ty = path[0].solution; + budget.charge_type(db, ty)?; + constructed += 1; + Ok(UnionType::from_two_elements(db, &env, accumulated, ty)) + }, + ); + + assert_eq!(constructed, max_type_terms); + if max_type_terms < 3 { + assert_eq!(collected, Err(ProjectionError::TypeBudgetExceeded)); + assert_eq!(result, Err(ProjectionError::TypeBudgetExceeded)); + } else { + assert_eq!( + collected, + Ok(Solutions::Constrained(SolutionPaths::Complete( + [int, str, bytes].map(|ty| vec![binding(t, ty)]).into() + ))) + ); + assert_eq!( + result, + Ok(SolutionProjection::Constrained(UnionType::from_elements( + db, + &env, + [int, str, bytes], + ))) + ); + } + } +} + +#[test] +fn type_budget_charges_nested_set_theoretic_terms() -> anyhow::Result<()> { + let mut db = setup_db(); + db.write_dedented( + "/src/a.py", + r#" +type Alias = int | str +type Recursive = int | Recursive +"#, + )?; + let db = &db; + let env = db.program_environment(); + let file = system_path_to_file(db, "/src/a.py")?; + let file = ProgramFile::new(db, file, env.program(db)); + let alias = |name| { + global_symbol(db, file, name) + .place + .expect_type() + .as_type_alias() + .map(Type::TypeAlias) + .ok_or_else(|| anyhow::anyhow!("expected alias {name}")) + }; + let int = known_instance(db, KnownClass::Int); + let str = known_instance(db, KnownClass::Str); + let union = UnionType::from_two_elements(db, &env, int, str); + let intersection = + IntersectionType::from_elements(db, &env, [int, Type::int_literal(1).negate(db, &env)]); + + // Existing set operations count their members; aliases cannot hide those members. A + // recursive alias is charged again at the cycle, but its body is expanded only once. + for (ty, terms) in [ + (union, 3), + (intersection, 3), + (alias("Alias")?, 4), + (alias("Recursive")?, 4), + ] { + assert_eq!( + ProjectionTypeBudget::new(terms - 1).charge_type(db, ty), + Err(ProjectionError::TypeBudgetExceeded) + ); + assert_eq!(ProjectionTypeBudget::new(terms).charge_type(db, ty), Ok(())); + } + Ok(()) +} + +#[test] +fn intersection_construction_failure_discards_the_projection() -> anyhow::Result<()> { + let mut db = setup_db(); + db.write_dedented( + "/src/a.py", + r#" +class A: ... +class B: ... +class C: ... +class D: ... +class E: ... +"#, + )?; + let db = &db; + let env = db.program_environment(); + let file = system_path_to_file(db, "/src/a.py")?; + let file = ProgramFile::new(db, file, env.program(db)); + let instance = |name| { + global_symbol(db, file, name) + .place + .expect_type() + .to_instance_approximation(db, &env) + .ok_or_else(|| anyhow::anyhow!("expected class {name}")) + }; + let left = UnionType::from_elements(db, &env, [instance("A")?, instance("B")?]); + let right = + UnionType::from_elements(db, &env, [instance("C")?, instance("D")?, instance("E")?]); + let t = create_typevar(db, "T"); + let builder = ConstraintSetBuilder::new(); + + // These classes can overlap, so distributing the intersection requires six DNF terms. + // Charging the input alone does not prevent that expansion; the fold also needs a bounded + // intersection constructor. + for alternatives in [[left, right], [right, left]] { + let paths = PathBounds::Constrained( + alternatives + .map(|ty| Box::new([PathBound::exact(t, ty)]) as Box<[_]>) + .into(), + ); + + assert_eq!( + paths.try_fold_with( + |_, bound| PathBounds::default_solve(db, &env, &builder, bound), + Type::object(), + &mut ProjectionTypeBudget::new(7), + |accumulated, path, budget| { + assert_eq!(path.len(), 1); + let ty = path[0].solution; + budget.charge_type(db, ty)?; + IntersectionType::bounded_from_elements(db, &env, [accumulated, ty]) + .ok_or(ProjectionError::TypeBudgetExceeded) + }, + ), + Err(ProjectionError::TypeBudgetExceeded) + ); + } + Ok(()) +} diff --git a/crates/ty_python_semantic/src/types/constraints/solutions.rs b/crates/ty_python_semantic/src/types/constraints/solutions.rs index 26860aaa66e65c..6394a725584f04 100644 --- a/crates/ty_python_semantic/src/types/constraints/solutions.rs +++ b/crates/ty_python_semantic/src/types/constraints/solutions.rs @@ -1,8 +1,9 @@ use std::marker::PhantomData; +use std::ops::ControlFlow; use crate::types::constraints::{ ALWAYS_FALSE, ALWAYS_TRUE, ConstraintBoundsBuilder, ConstraintId, ConstraintSetStorage, NodeId, - PathAssignments, PathBounds, + PathAssignments, PathBounds, SolutionLimits, }; use crate::types::{BoundTypeVarInstance, Type}; use crate::{Db, FxIndexMap, FxIndexSet, ProgramEnvironment}; @@ -22,59 +23,49 @@ impl<'db> SolutionWalker<'db> { } } - pub(super) fn visit_node( + pub(super) fn visit_node( &mut self, db: &'db dyn Db, env: &ProgramEnvironment<'db>, storage: &mut ConstraintSetStorage<'db>, path: &mut PathAssignments, node: NodeId, - ) { + limits: &mut L, + ) -> ControlFlow { + limits.visit_node()?; if node == ALWAYS_FALSE { - return; + return ControlFlow::Continue(()); } // If the current node is ALWAYS_TRUE, we can immediately report the current solution. if node == ALWAYS_TRUE { + limits.satisfied_path()?; self.found_satisfied_path(path); - return; + return ControlFlow::Continue(()); } // At this point we actually have to walk the outgoing edges of this node. let interior = storage.interior_node_data(node); - path.walk_edge( - db, - env, - storage, - interior.constraint.when_true(), - |storage, path, _new_range, found_conflict| { - if !found_conflict { - self.visit_node(db, env, storage, path, interior.if_true); - } - }, - ); - path.walk_edge( - db, - env, - storage, - interior.constraint.when_unconstrained(), - |storage, path, _new_range, found_conflict| { - if !found_conflict { - self.visit_node(db, env, storage, path, interior.if_uncertain); - } - }, - ); - path.walk_edge( - db, - env, - storage, - interior.constraint.when_false(), - |storage, path, _new_range, found_conflict| { - if !found_conflict { - self.visit_node(db, env, storage, path, interior.if_false); - } - }, - ); + let constraint = interior.constraint; + for (assignment, child) in [ + (constraint.when_true(), interior.if_true), + (constraint.when_unconstrained(), interior.if_uncertain), + (constraint.when_false(), interior.if_false), + ] { + path.walk_edge( + db, + env, + storage, + assignment, + |storage, path, _new_range, found_conflict| { + if !found_conflict { + self.visit_node(db, env, storage, path, child, limits)?; + } + ControlFlow::Continue(()) + }, + )?; + } + ControlFlow::Continue(()) } fn found_satisfied_path(&mut self, path: &PathAssignments) { diff --git a/crates/ty_python_semantic/src/types/generics.rs b/crates/ty_python_semantic/src/types/generics.rs index 79e94f530269af..aa9ff9f23aa58b 100644 --- a/crates/ty_python_semantic/src/types/generics.rs +++ b/crates/ty_python_semantic/src/types/generics.rs @@ -11,6 +11,7 @@ use smallvec::SmallVec; use crate::types::callable::walk_callable_type; use crate::types::class::ClassType; use crate::types::class_base::ClassBase; +use crate::types::constraints::projection::{SolutionBudget, SolutionProjection}; use crate::types::constraints::{ ConstraintBounds, ConstraintSet, ConstraintSetBuilder, IteratorConstraintsExtension, PathBound, PathBoundSolution, PathBounds, SolutionPaths, Solutions, @@ -2324,10 +2325,17 @@ pub(crate) struct SpecializationBuilder<'db, 'c> { generic_context: GenericContext<'db>, inferable: TypeVarSet<'db>, pending: ConstraintSet<'db, 'c>, - types: FxHashMap, UnionAccumulator<'db>>, + types: LegacyTypeMappings<'db>, paramspec_seen: FxHashSet>, } +/// The legacy mapping is usable only if no accepted relation was omitted in its entirety. +/// Missing evidence from one argument can make the other arguments' inferred types too narrow. +enum LegacyTypeMappings<'db> { + Available(FxHashMap, UnionAccumulator<'db>>), + BudgetExceeded, +} + /// The result of type variable inference before choosing how to handle unsolved type variables. /// /// A `Some` entry means inference solved the corresponding type variable to that type. A `None` @@ -2396,11 +2404,14 @@ impl<'db> TypeVarInference<'db> { /// /// Failed paths can occur alongside valid paths, but their declaration failures matter only when /// every path is rejected. Preserve that evidence exclusively for unsatisfiable constraint sets. -/// Budget-exhausted paths retain their fallback bindings and completeness information. +/// Per-variable budget exhaustion retains fallback bindings and completeness information. If +/// collecting the whole relation exceeds a limit, no partial family is available for projection. enum ConstraintSetAnalysis<'db> { Unsatisfiable(SmallVec<[ConstraintFailure<'db>; 1]>), Unconstrained, Constrained(SolutionPaths<'db>), + /// A collection or result limit was exceeded. No partial family is safe to project. + BudgetExceeded, } impl<'db> ConstraintSetAnalysis<'db> { @@ -2521,7 +2532,7 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { generic_context, inferable: generic_context.inferable_typevars(db), pending: ConstraintSet::from_bool(constraints, true), - types: FxHashMap::default(), + types: LegacyTypeMappings::Available(FxHashMap::default()), paramspec_seen: FxHashSet::default(), } } @@ -2648,40 +2659,56 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { // was not enough: `solutions_with` still performed the expensive path traversal, and the // skipped projection changed precision in LiteralString tests. See the // `ty_micro[pydantic_core_schema_dict]` benchmark for a minimized reproducer. - let solutions = match self.pending.solutions_with( + let mut types = match self.pending.try_fold_solutions( db, self.env, - self.constraints, self.inferable, + SolutionBudget::default(), |_variance, path_bound| { let typevar = path_bound.bound_typevar; if let Some(ty) = choose(typevar, Some(path_bound)) { return PathBoundSolution::Solved(ty); } - PathBounds::default_solve(db, self.env, self.constraints, path_bound) + // The legacy projection accepts per-variable fallback bindings. + match PathBounds::default_solve(db, self.env, self.constraints, path_bound) { + PathBoundSolution::BudgetExceeded { fallback } => { + fallback.map_or(PathBoundSolution::Unsolved, PathBoundSolution::Solved) + } + outcome => outcome, + } + }, + FxHashMap::default(), + |mut types: FxHashMap<_, _>, solution, budget| { + for binding in solution { + budget.charge_type(db, binding.solution)?; + let identity = binding.bound_typevar.identity(db); + types + .entry(identity) + .and_modify(|existing| { + *existing = UnionType::from_two_elements( + db, + self.env, + *existing, + binding.solution, + ); + }) + .or_insert(binding.solution); + } + Ok(types) }, ) { - Solutions::Unsatisfiable => return Err(()), - Solutions::Unconstrained => { + Ok(SolutionProjection::Unsatisfiable) => return Err(()), + Ok(SolutionProjection::Unconstrained) => { return Ok(self.solve_hash_map_with(generic_context, choose)); } - Solutions::Constrained(solutions) => solutions.into_vec(), - }; - - let mut types = FxHashMap::default(); - for solution in solutions { - for binding in solution { - let identity = binding.bound_typevar.identity(db); - types - .entry(identity) - .and_modify(|existing| { - *existing = - UnionType::from_two_elements(db, self.env, *existing, binding.solution); - }) - .or_insert(binding.solution); + Err(_) => { + // A partial mapping may be narrower than the unvisited alternatives. Recover + // without choosing a witness or repeating the exhausted traversal/construction. + return Ok(self.unknown_type_mappings(generic_context)); } - } + Ok(SolutionProjection::Constrained(types)) => types, + }; // Sequent-map transitivity can add relationships between inferable typevars to path // bounds. Those relationships are important while solving, but should not become recursive @@ -2854,12 +2881,14 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { choose: &mut impl FnMut(BoundTypeVarInstance<'db>, Option<&PathBound<'db>>) -> Option>, ) -> FxHashMap, Type<'db>> { let db = self.db; + let LegacyTypeMappings::Available(types) = &mut self.types else { + return self.unknown_type_mappings(generic_context); + }; generic_context .variables_inner(db) .iter() .filter_map(|(identity, variable)| { - let mapped_ty = self - .types + let mapped_ty = types .get_mut(identity) .map(|accumulator| accumulator.get_or_build(db, self.env)); let chosen = match mapped_ty { @@ -2874,6 +2903,20 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { .collect() } + fn unknown_type_mappings( + &self, + generic_context: GenericContext<'db>, + ) -> FxHashMap, Type<'db>> { + let db = self.db; + let unknown = generic_context.unknown_specialization(db, None); + generic_context + .variables_inner(db) + .keys() + .copied() + .zip(unknown.types(db).iter().copied()) + .collect() + } + fn insert_hash_map_type_mapping( &mut self, bound_typevar: BoundTypeVarInstance<'db>, @@ -2881,7 +2924,10 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { ) { let db = self.db; let identity = bound_typevar.identity(db); - match self.types.entry(identity) { + let LegacyTypeMappings::Available(types) = &mut self.types else { + return; + }; + match types.entry(identity) { Entry::Occupied(mut entry) => { match bound_typevar.kind(db) { TypeVarKind::LegacyParamSpec | TypeVarKind::Pep695ParamSpec => { @@ -2957,13 +3003,14 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { ty: Type<'db>, ) -> bool { let db = self.db; - self.types - .get_mut(&bound_typevar) - .is_some_and(|inferred_ty| { - inferred_ty - .get_or_build(db, self.env) - .is_assignable_to(db, self.env, ty) - }) + let LegacyTypeMappings::Available(types) = &mut self.types else { + return false; + }; + types.get_mut(&bound_typevar).is_some_and(|inferred_ty| { + inferred_ty + .get_or_build(db, self.env) + .is_assignable_to(db, self.env, ty) + }) } /// Add a type mapping for a bound typevar using the given variance to determine how the @@ -3000,8 +3047,8 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { let solutions = set.solutions_with( db, self.env, - self.constraints, self.inferable, + SolutionBudget::default(), |_variance, path_bound| { let solution = PathBounds::preliminary_solve(db, self.env, self.constraints, path_bound); @@ -3015,9 +3062,10 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { ); match solutions { - Solutions::Unsatisfiable => ConstraintSetAnalysis::Unsatisfiable(failures), - Solutions::Unconstrained => ConstraintSetAnalysis::Unconstrained, - Solutions::Constrained(solutions) => ConstraintSetAnalysis::Constrained(solutions), + Ok(Solutions::Unsatisfiable) => ConstraintSetAnalysis::Unsatisfiable(failures), + Ok(Solutions::Unconstrained) => ConstraintSetAnalysis::Unconstrained, + Ok(Solutions::Constrained(solutions)) => ConstraintSetAnalysis::Constrained(solutions), + Err(_) => ConstraintSetAnalysis::BudgetExceeded, } } @@ -3025,10 +3073,18 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { /// /// This projection loses correlations between alternatives, so callers must only request it /// after they have accepted the corresponding relation. + /// Omitting an accepted relation makes the legacy mapping unavailable for precise recovery. /// /// TODO: Remove this compatibility path once [`build_with`][Self::build_with] and all other /// inference consumers can build specializations solely from the call-wide constraint set. fn project_for_legacy_fallback(&mut self, analysis: &ConstraintSetAnalysis<'db>) { + if matches!(analysis, ConstraintSetAnalysis::BudgetExceeded) { + self.types = LegacyTypeMappings::BudgetExceeded; + return; + } + if matches!(self.types, LegacyTypeMappings::BudgetExceeded) { + return; + } let ConstraintSetAnalysis::Constrained(solutions) = analysis else { return; }; @@ -3291,8 +3347,9 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { let mapping_when = mapping.when_constraint_set_assignable_to_owned(db, env, formal); let mapping_when = self.constraints.load(db, env, &mapping_when); // Logically equivalent constraints can still infer different solutions, such as `Any` - // instead of `object`; preserve the original constraints when gradual evidence differs. - let mapping_solutions = mapping_when.solutions(db, env, self.constraints, self.inferable); + // instead of `object`; preserve the original constraints when gradual evidence differs + // or either solution collection exceeds its budget. + let mapping_solutions = mapping_when.solutions(db, env, self.inferable).ok()?; if !typed_dicts.into_iter().all(|element| { let element_when = self.constraints.load( db, @@ -3302,8 +3359,9 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { element_when .iff(db, self.constraints, mapping_when) .is_always_satisfied(db, env) - && element_when.solutions(db, env, self.constraints, self.inferable) - == mapping_solutions + && element_when + .solutions(db, env, self.inferable) + .is_ok_and(|solutions| solutions == mapping_solutions) }) { return None; } @@ -3393,8 +3451,8 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> { .map_or(Ok(()), Err); }; - // One accepted alternative proves that their disjunction is satisfiable; solving - // the combined TDD again would eagerly repeat the expensive path enumeration. + // Retain every alternative that was not proved unsatisfiable. Solving the + // combined TDD here would repeat their potentially expensive path traversals. self.record_constraint_set(combined); for (_, analysis) in accepted { self.project_for_legacy_fallback(&analysis); @@ -4205,17 +4263,54 @@ mod tests { let set = ConstraintSet::constrain_typevar(db, &env, &constraints, typevar, int, int); let analysis = builder.analyze_constraint_set(set); - assert!(builder.types.is_empty()); + assert!(matches!(&builder.types, LegacyTypeMappings::Available(types) if types.is_empty())); assert!(builder.pending.is_always_satisfied(db, &env)); builder.record_constraint_set(set); - assert!(builder.types.is_empty()); + assert!(matches!(&builder.types, LegacyTypeMappings::Available(types) if types.is_empty())); assert!(!builder.pending.is_always_satisfied(db, &env)); builder.project_for_legacy_fallback(&analysis); assert!(builder.inferred_type_is_assignable_to(typevar.identity(db), int)); } + #[test] + fn exhausted_projection_keeps_legacy_mapping_unavailable() { + let db = setup_db(); + let db = &db; + let env = db.program_environment(); + let typevar = BoundTypeVarInstance::synthetic( + db, + &env, + Name::new_static("T"), + TypeVarVariance::Invariant, + ); + let context = GenericContext::from_typevar_instances(db, &env, [typevar]); + let constraints = ConstraintSetBuilder::new(); + let mut builder = SpecializationBuilder::new(db, &env, &constraints, context); + let str = KnownClass::Str.to_instance(db, &env); + + builder.add_type_mapping(typevar, str, TypeVarVariance::Covariant); + let ty = UnionType::from_two_elements(db, &env, str, Type::int_literal(0)); + let relation = ConstraintSet::constrain_typevar(db, &env, &constraints, typevar, ty, ty); + builder.record_constraint_set(relation); + builder.project_for_legacy_fallback(&ConstraintSetAnalysis::BudgetExceeded); + builder.add_type_mapping(typevar, str, TypeVarVariance::Covariant); + assert!(matches!(builder.types, LegacyTypeMappings::BudgetExceeded)); + assert!(!builder.pending.is_never_satisfied(db, &env)); + + let mut choices = 0; + let types = builder.solve_hash_map_with(context, &mut |_, _| { + choices += 1; + Some(str) + }); + assert_eq!(choices, 0); + assert_eq!( + types, + FxHashMap::from_iter([(typevar.identity(db), Type::unknown())]) + ); + } + #[test] fn satisfiable_constraint_analysis_discards_rejected_paths() { let db = setup_db(); diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index 56ac59446c3b1c..6bd326490bd417 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -879,12 +879,9 @@ fn specialize_narrowing_target_from_intersection<'db>( combined_constraints.intersect(db, &constraints, base_constraint); } - let solutions = combined_constraints.solutions( - db, - env, - &constraints, - generic_context.inferable_typevars(db), - ); + let solutions = combined_constraints + .solutions(db, env, generic_context.inferable_typevars(db)) + .ok()?; let specialized_class = specialize_generic_class_from_solutions(db, env, target_class, solutions)?; Some(Type::instance(db, env, specialized_class)) diff --git a/crates/ty_python_semantic/src/types/signatures.rs b/crates/ty_python_semantic/src/types/signatures.rs index cae449ab9dc505..98821a5688b997 100644 --- a/crates/ty_python_semantic/src/types/signatures.rs +++ b/crates/ty_python_semantic/src/types/signatures.rs @@ -1296,15 +1296,15 @@ impl<'db> Signature<'db> { let when = constraints.load(db, env, receiver_constraints); let inferable = self.inferable_typevars(db); - match when.solutions(db, env, &constraints, inferable) { - Solutions::Unsatisfiable => return None, - Solutions::Unconstrained => return Some(self.clone()), + match when.solutions(db, env, inferable) { + Ok(Solutions::Unsatisfiable) => return None, + Ok(Solutions::Unconstrained) | Err(_) => return Some(self.clone()), // Each receiver path can leave a different type variable unconstrained. Preserve the // original relation instead of combining those independent solutions. - Solutions::Constrained(solutions) if solutions.as_slice().len() > 1 => { + Ok(Solutions::Constrained(solutions)) if solutions.as_slice().len() > 1 => { return Some(self.clone()); } - Solutions::Constrained(_) => {} + Ok(Solutions::Constrained(_)) => {} } let Some(generic_context) = self.generic_context else {