From 8c2603f2d21ebfc32733c17c0cb61426cb9c0caa Mon Sep 17 00:00:00 2001 From: Douglas Creager Date: Fri, 7 Nov 2025 15:30:19 -0500 Subject: [PATCH] store inferable --- .../ty_python_semantic/src/types/call/bind.rs | 18 ++- .../ty_python_semantic/src/types/generics.rs | 130 +++++++++++------- .../src/types/signatures.rs | 23 ++-- 3 files changed, 101 insertions(+), 70 deletions(-) diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index c017ee68d4..cfd60ad829 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/crates/ty_python_semantic/src/types/call/bind.rs @@ -9,7 +9,6 @@ use std::fmt; use itertools::{Either, Itertools}; use ruff_db::parsed::parsed_module; use ruff_python_ast::name::Name; -use rustc_hash::FxHashSet; use smallvec::{SmallVec, smallvec, smallvec_inline}; use super::{Argument, CallArguments, CallError, CallErrorKind, InferContext, Signature, Type}; @@ -1217,21 +1216,20 @@ impl<'db> Bindings<'db> { let extract_inferable = |instance: &NominalInstanceType<'db>| { if instance.has_known_class(db, KnownClass::NoneType) { // Caller explicitly passed None, so no typevars are inferable. - return Some(FxHashSet::default()); + return Some(InferableTypeVars::None); } - instance - .tuple_spec(db)? - .fixed_elements() - .map(|ty| { + Some(InferableTypeVars::from_bound_typevars( + db, + instance.tuple_spec(db)?.fixed_elements().filter_map(|ty| { ty.as_typevar() .map(|bound_typevar| bound_typevar.identity(db)) - }) - .collect() + }), + )) }; let inferable = match overload.parameter_types() { // Caller did not provide argument, so no typevars are inferable. - [None] => FxHashSet::default(), + [None] => InferableTypeVars::None, [Some(Type::NominalInstance(instance))] => { match extract_inferable(instance) { Some(inferable) => inferable, @@ -1243,7 +1241,7 @@ impl<'db> Bindings<'db> { let result = tracked .constraints(db) - .satisfied_by_all_typevars(db, InferableTypeVars::One(&inferable)); + .satisfied_by_all_typevars(db, inferable); overload.set_return_type(Type::BooleanLiteral(result)); } diff --git a/crates/ty_python_semantic/src/types/generics.rs b/crates/ty_python_semantic/src/types/generics.rs index 31cfcbf6a1..10fe7bfa54 100644 --- a/crates/ty_python_semantic/src/types/generics.rs +++ b/crates/ty_python_semantic/src/types/generics.rs @@ -1,10 +1,12 @@ use std::cell::RefCell; use std::collections::hash_map::Entry; use std::fmt::Display; +use std::marker::PhantomData; use itertools::Itertools; use ruff_python_ast as ast; -use rustc_hash::{FxHashMap, FxHashSet}; +use rustc_hash::FxHashMap; +use smallvec::{SmallVec, smallvec}; use crate::semantic_index::definition::Definition; use crate::semantic_index::scope::{FileScopeId, NodeWithScopeKind, ScopeId}; @@ -121,13 +123,18 @@ pub(crate) fn typing_self<'db>( #[derive(Clone, Copy, Debug)] pub(crate) enum InferableTypeVars<'a, 'db> { None, - One(&'a FxHashSet>), - Two( - &'a InferableTypeVars<'a, 'db>, - &'a InferableTypeVars<'a, 'db>, - ), + One(InferableTypeVarsInner<'db>, PhantomData<&'a ()>), } +#[salsa::tracked(debug, heap_size=ruff_memory_usage::heap_size)] +pub(crate) struct InferableTypeVarsInner<'db> { + // The _set_ of typevars that are inferable. This will always be sorted and deduped. + inferable: SmallVec<[BoundTypeVarIdentity<'db>; 4]>, +} + +// The Salsa heap is tracked separately. +impl get_size2::GetSize for InferableTypeVarsInner<'_> {} + impl<'db> BoundTypeVarInstance<'db> { pub(crate) fn is_inferable( self, @@ -136,51 +143,78 @@ impl<'db> BoundTypeVarInstance<'db> { ) -> bool { match inferable { InferableTypeVars::None => false, - InferableTypeVars::One(typevars) => typevars.contains(&self.identity(db)), - InferableTypeVars::Two(left, right) => { - self.is_inferable(db, *left) || self.is_inferable(db, *right) - } + InferableTypeVars::One(inner, _) => inner + .inferable(db) + .binary_search(&self.identity(db)) + .is_ok(), } } } -impl<'a, 'db> InferableTypeVars<'a, 'db> { - pub(crate) fn merge(&'a self, other: Option<&'a InferableTypeVars<'a, 'db>>) -> Self { - match other { - Some(other) => InferableTypeVars::Two(self, other), - None => *self, +impl<'db> InferableTypeVars<'_, 'db> { + pub(crate) fn from_bound_typevars( + db: &'db dyn Db, + bound_typevars: impl IntoIterator>, + ) -> Self { + InferableTypeVars::One( + InferableTypeVarsInner::from_bound_typevars(db, bound_typevars), + PhantomData, + ) + } + + pub(crate) fn merge(self, db: &'db dyn Db, other: Self) -> Self { + match (self, other) { + (InferableTypeVars::None, InferableTypeVars::None) => self, + (InferableTypeVars::One(_, _), InferableTypeVars::None) => self, + (InferableTypeVars::None, InferableTypeVars::One(_, _)) => other, + (InferableTypeVars::One(self_inner, _), InferableTypeVars::One(other_inner, _)) => { + InferableTypeVars::One(self_inner.merge(db, other_inner), PhantomData) + } } } // Keep this around for debugging purposes #[expect(dead_code)] - pub(crate) fn display(&self, db: &'db dyn Db) -> impl Display { - fn find_typevars<'db>( - result: &mut FxHashSet>, - inferable: &InferableTypeVars<'_, 'db>, - ) { - match inferable { - InferableTypeVars::None => {} - InferableTypeVars::One(typevars) => result.extend(typevars.iter().copied()), - InferableTypeVars::Two(left, right) => { - find_typevars(result, left); - find_typevars(result, right); - } - } - } - - let mut typevars = FxHashSet::default(); - find_typevars(&mut typevars, self); + pub(crate) fn display(self, db: &'db dyn Db) -> impl Display { + let inferable = match self { + InferableTypeVars::One(inner, _) => inner.inferable(db), + InferableTypeVars::None => return String::from("[]"), + }; format!( "[{}]", - typevars - .into_iter() + inferable + .iter() .map(|identity| identity.display(db)) .format(", ") ) } } +#[salsa::tracked] +impl<'db> InferableTypeVarsInner<'db> { + fn from_bound_typevars( + db: &'db dyn Db, + bound_typevars: impl IntoIterator>, + ) -> Self { + let mut inferable: SmallVec<_> = bound_typevars.into_iter().collect(); + inferable.sort_unstable(); + inferable.dedup(); + InferableTypeVarsInner::new(db, inferable) + } + + #[salsa::tracked] + fn merge(self, db: &'db dyn Db, other: Self) -> Self { + // The input typevar vecs are already sorted, so we can merge/dedup them instead of + // having to do an expensive sort. + let self_inferable = self.inferable(db); + let self_typevars = self_inferable.iter().copied(); + let other_inferable = other.inferable(db); + let other_typevars = other_inferable.iter().copied(); + let inferable = self_typevars.merge(other_typevars).dedup().collect(); + InferableTypeVarsInner::new(db, inferable) + } +} + #[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, get_size2::GetSize)] pub struct GenericContextTypeVar<'db> { bound_typevar: BoundTypeVarInstance<'db>, @@ -288,7 +322,7 @@ impl<'db> GenericContext<'db> { pub(crate) fn inferable_typevars(self, db: &'db dyn Db) -> InferableTypeVars<'db, 'db> { #[derive(Default)] struct CollectTypeVars<'db> { - typevars: RefCell>>, + typevars: RefCell; 4]>>, recursion_guard: TypeCollector<'db>, } @@ -302,9 +336,7 @@ impl<'db> GenericContext<'db> { db: &'db dyn Db, bound_typevar: BoundTypeVarInstance<'db>, ) { - self.typevars - .borrow_mut() - .insert(bound_typevar.identity(db)); + self.typevars.borrow_mut().push(bound_typevar.identity(db)); walk_bound_type_var_type(db, bound_typevar, self); } @@ -314,25 +346,27 @@ impl<'db> GenericContext<'db> { } #[salsa::tracked( - returns(ref), cycle_initial=inferable_typevars_cycle_initial, heap_size=ruff_memory_usage::heap_size, )] fn inferable_typevars_inner<'db>( db: &'db dyn Db, generic_context: GenericContext<'db>, - ) -> FxHashSet> { + ) -> InferableTypeVarsInner<'db> { let visitor = CollectTypeVars::default(); for bound_typevar in generic_context.variables(db) { visitor.visit_bound_type_var_type(db, bound_typevar); } - visitor.typevars.into_inner() + let mut inferable = visitor.typevars.into_inner(); + inferable.sort_unstable(); + inferable.dedup(); + InferableTypeVarsInner::new(db, inferable) } - // This ensures that salsa caches the FxHashSet, not the InferableTypeVars that wraps it. - // (That way InferableTypeVars can contain references, and doesn't need to impl - // salsa::Update.) - InferableTypeVars::One(inferable_typevars_inner(db, self)) + // This ensures that salsa caches the InferableTypeVarsInner, not the InferableTypeVars + // that wraps it. (That way InferableTypeVars can contain a reference, and doesn't need to + // impl salsa::Update.) + InferableTypeVars::One(inferable_typevars_inner(db, self), PhantomData) } pub(crate) fn variables( @@ -618,11 +652,11 @@ impl<'db> GenericContext<'db> { } fn inferable_typevars_cycle_initial<'db>( - _db: &'db dyn Db, + db: &'db dyn Db, _id: salsa::Id, _self: GenericContext<'db>, -) -> FxHashSet> { - FxHashSet::default() +) -> InferableTypeVarsInner<'db> { + InferableTypeVarsInner::new(db, smallvec![]) } #[derive(Debug, Copy, Clone, PartialEq, Eq)] diff --git a/crates/ty_python_semantic/src/types/signatures.rs b/crates/ty_python_semantic/src/types/signatures.rs index 774a2bf5b7..cf103adecb 100644 --- a/crates/ty_python_semantic/src/types/signatures.rs +++ b/crates/ty_python_semantic/src/types/signatures.rs @@ -628,6 +628,13 @@ impl<'db> Signature<'db> { } } + fn inferable_typevars(&self, db: &'db dyn Db) -> InferableTypeVars<'db, 'db> { + match self.generic_context { + Some(generic_context) => generic_context.inferable_typevars(db), + None => InferableTypeVars::None, + } + } + /// Return `true` if `self` has exactly the same set of possible static materializations as /// `other` (if `self` represents the same set of possible sets of possible runtime objects as /// `other`). @@ -640,12 +647,8 @@ impl<'db> Signature<'db> { ) -> ConstraintSet<'db> { // The typevars in self and other should also be considered inferable when checking whether // two signatures are equivalent. - let self_inferable = - (self.generic_context).map(|generic_context| generic_context.inferable_typevars(db)); - let other_inferable = - (other.generic_context).map(|generic_context| generic_context.inferable_typevars(db)); - let inferable = inferable.merge(self_inferable.as_ref()); - let inferable = inferable.merge(other_inferable.as_ref()); + let inferable = inferable.merge(db, self.inferable_typevars(db)); + let inferable = inferable.merge(db, other.inferable_typevars(db)); let mut result = ConstraintSet::from(true); let mut check_types = |self_type: Option>, other_type: Option>| { @@ -799,12 +802,8 @@ impl<'db> Signature<'db> { // The typevars in self and other should also be considered inferable when checking whether // two signatures are equivalent. - let self_inferable = - (self.generic_context).map(|generic_context| generic_context.inferable_typevars(db)); - let other_inferable = - (other.generic_context).map(|generic_context| generic_context.inferable_typevars(db)); - let inferable = inferable.merge(self_inferable.as_ref()); - let inferable = inferable.merge(other_inferable.as_ref()); + let inferable = inferable.merge(db, self.inferable_typevars(db)); + let inferable = inferable.merge(db, other.inferable_typevars(db)); let mut result = ConstraintSet::from(true); let mut check_types = |type1: Option>, type2: Option>| {