Skip to content

[ty] Apply type mappings to functions eagerly #20596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sharkdp merged 4 commits into from
Sep 29, 2025
Merged

[ty] Apply type mappings to functions eagerly #20596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6a449de
don't apply type mappings lazily
dcreager Sep 26, 2025
2ceff49
walk_function_type: visit updated signatures as well
sharkdp Sep 29, 2025
cec807d
remove superfluous lifetimes
sharkdp Sep 29, 2025
5fb5cc6
Add doc comments for updated_* fields, make them non-public
sharkdp Sep 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 5 additions & 83 deletions crates/ty_python_semantic/src/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,6 @@ use crate::types::function::{
};
use crate::types::generics::{
GenericContext, PartialSpecialization, Specialization, bind_typevar, walk_generic_context,
walk_partial_specialization, walk_specialization,
};
pub use crate::types::ide_support::{
CallSignatureDetails, Member, MemberWithDefinition, all_members, call_signature_details,
Expand Down Expand Up @@ -6109,7 +6108,7 @@ impl<'db> Type<'db> {
}

Type::FunctionLiteral(function) => {
let function = Type::FunctionLiteral(function.with_type_mapping(db, type_mapping));
let function = Type::FunctionLiteral(function.apply_type_mapping_impl(db, type_mapping, visitor));

match type_mapping {
TypeMapping::PromoteLiterals => function.literal_promotion_type(db)
Expand All @@ -6120,8 +6119,8 @@ impl<'db> Type<'db> {

Type::BoundMethod(method) => Type::BoundMethod(BoundMethodType::new(
db,
method.function(db).with_type_mapping(db, type_mapping),
method.self_instance(db).apply_type_mapping(db, type_mapping),
method.function(db).apply_type_mapping_impl(db, type_mapping, visitor),
method.self_instance(db).apply_type_mapping_impl(db, type_mapping, visitor),
)),

Type::NominalInstance(instance) =>
Expand All @@ -6140,13 +6139,13 @@ impl<'db> Type<'db> {

Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderGet(function)) => {
Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderGet(
function.with_type_mapping(db, type_mapping),
function.apply_type_mapping_impl(db, type_mapping, visitor),
))
}

Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderCall(function)) => {
Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderCall(
function.with_type_mapping(db, type_mapping),
function.apply_type_mapping_impl(db, type_mapping, visitor),
))
}

Expand Down Expand Up @@ -6782,84 +6781,7 @@ pub enum TypeMapping<'a, 'db> {
Materialize(MaterializationKind),
}

fn walk_type_mapping<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
db: &'db dyn Db,
mapping: &TypeMapping<'_, 'db>,
visitor: &V,
) {
match mapping {
TypeMapping::Specialization(specialization) => {
walk_specialization(db, *specialization, visitor);
}
TypeMapping::PartialSpecialization(specialization) => {
walk_partial_specialization(db, specialization, visitor);
}
TypeMapping::BindSelf(self_type) => {
visitor.visit_type(db, *self_type);
}
TypeMapping::ReplaceSelf { new_upper_bound } => {
visitor.visit_type(db, *new_upper_bound);
}
TypeMapping::PromoteLiterals
| TypeMapping::BindLegacyTypevars(_)
| TypeMapping::MarkTypeVarsInferable(_)
| TypeMapping::Materialize(_) => {}
}
}

impl<'db> TypeMapping<'_, 'db> {
fn to_owned(&self) -> TypeMapping<'db, 'db> {
match self {
TypeMapping::Specialization(specialization) => {
TypeMapping::Specialization(*specialization)
}
TypeMapping::PartialSpecialization(partial) => {
TypeMapping::PartialSpecialization(partial.to_owned())
}
TypeMapping::PromoteLiterals => TypeMapping::PromoteLiterals,
TypeMapping::BindLegacyTypevars(binding_context) => {
TypeMapping::BindLegacyTypevars(*binding_context)
}
TypeMapping::BindSelf(self_type) => TypeMapping::BindSelf(*self_type),
TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
new_upper_bound: *new_upper_bound,
},
TypeMapping::MarkTypeVarsInferable(binding_context) => {
TypeMapping::MarkTypeVarsInferable(*binding_context)
}
TypeMapping::Materialize(materialization_kind) => {
TypeMapping::Materialize(*materialization_kind)
}
}
}

fn normalized_impl(&self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
match self {
TypeMapping::Specialization(specialization) => {
TypeMapping::Specialization(specialization.normalized_impl(db, visitor))
}
TypeMapping::PartialSpecialization(partial) => {
TypeMapping::PartialSpecialization(partial.normalized_impl(db, visitor))
}
TypeMapping::PromoteLiterals => TypeMapping::PromoteLiterals,
TypeMapping::BindLegacyTypevars(binding_context) => {
TypeMapping::BindLegacyTypevars(*binding_context)
}
TypeMapping::BindSelf(self_type) => {
TypeMapping::BindSelf(self_type.normalized_impl(db, visitor))
}
TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
new_upper_bound: new_upper_bound.normalized_impl(db, visitor),
},
TypeMapping::MarkTypeVarsInferable(binding_context) => {
TypeMapping::MarkTypeVarsInferable(*binding_context)
}
TypeMapping::Materialize(materialization_kind) => {
TypeMapping::Materialize(*materialization_kind)
}
}
}

/// Update the generic context of a [`Signature`] according to the current type mapping
pub(crate) fn update_signature_generic_context(
&self,
Expand Down
146 changes: 85 additions & 61 deletions crates/ty_python_semantic/src/types/function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,11 @@ use crate::types::narrow::ClassInfoConstraintFunction;
use crate::types::signatures::{CallableSignature, Signature};
use crate::types::visitor::any_over_type;
use crate::types::{
BoundMethodType, BoundTypeVarInstance, CallableType, ClassBase, ClassLiteral, ClassType,
DeprecatedInstance, DynamicType, FindLegacyTypeVarsVisitor, HasRelationToVisitor,
IsEquivalentVisitor, KnownClass, KnownInstanceType, NormalizedVisitor, SpecialFormType,
TrackedConstraintSet, Truthiness, Type, TypeMapping, TypeRelation, UnionBuilder, all_members,
binding_type, todo_type, walk_type_mapping,
ApplyTypeMappingVisitor, BoundMethodType, BoundTypeVarInstance, CallableType, ClassBase,
ClassLiteral, ClassType, DeprecatedInstance, DynamicType, FindLegacyTypeVarsVisitor,
HasRelationToVisitor, IsEquivalentVisitor, KnownClass, KnownInstanceType, NormalizedVisitor,
SpecialFormType, TrackedConstraintSet, Truthiness, Type, TypeMapping, TypeRelation,
UnionBuilder, all_members, binding_type, todo_type, walk_signature,
};
use crate::{Db, FxOrderSet, ModuleName, resolve_module};

Expand Down Expand Up @@ -623,33 +623,24 @@ impl<'db> FunctionLiteral<'db> {
/// calling query is not in the same file as this function is defined in, then this will create
/// a cross-module dependency directly on the full AST which will lead to cache
/// over-invalidation.
fn signature<'a>(
self,
db: &'db dyn Db,
type_mappings: &'a [TypeMapping<'a, 'db>],
) -> CallableSignature<'db>
where
'db: 'a,
{
fn signature(self, db: &'db dyn Db) -> CallableSignature<'db> {
// We only include an implementation (i.e. a definition not decorated with `@overload`) if
// it's the only definition.
let inherited_generic_context = self.inherited_generic_context(db);
let (overloads, implementation) = self.overloads_and_implementation(db);
if let Some(implementation) = implementation {
if overloads.is_empty() {
return CallableSignature::single(type_mappings.iter().fold(
return CallableSignature::single(
implementation.signature(db, inherited_generic_context),
|sig, mapping| sig.apply_type_mapping(db, mapping),
));
);
}
}

CallableSignature::from_overloads(overloads.iter().map(|overload| {
type_mappings.iter().fold(
overload.signature(db, inherited_generic_context),
|sig, mapping| sig.apply_type_mapping(db, mapping),
)
}))
CallableSignature::from_overloads(
overloads
.iter()
.map(|overload| overload.signature(db, inherited_generic_context)),
)
}

/// Typed externally-visible signature of the last overload or implementation of this function.
Expand All @@ -660,20 +651,10 @@ impl<'db> FunctionLiteral<'db> {
/// calling query is not in the same file as this function is defined in, then this will create
/// a cross-module dependency directly on the full AST which will lead to cache
/// over-invalidation.
fn last_definition_signature<'a>(
self,
db: &'db dyn Db,
type_mappings: &'a [TypeMapping<'a, 'db>],
) -> Signature<'db>
where
'db: 'a,
{
fn last_definition_signature(self, db: &'db dyn Db) -> Signature<'db> {
let inherited_generic_context = self.inherited_generic_context(db);
type_mappings.iter().fold(
self.last_definition(db)
.signature(db, inherited_generic_context),
|sig, mapping| sig.apply_type_mapping(db, mapping),
)
self.last_definition(db)
.signature(db, inherited_generic_context)
}

fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
Expand All @@ -691,11 +672,19 @@ impl<'db> FunctionLiteral<'db> {
pub struct FunctionType<'db> {
pub(crate) literal: FunctionLiteral<'db>,
Comment on lines 672 to 673
Copy link
Member Author

@dcreager dcreager Sep 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We might also consider consolidating FunctionType and FunctionLiteral into a single type now. (FunctionLiteral would have the extra fields that hold the updated signatures.)


/// Type mappings that should be applied to the function's parameter and return types. This
/// might include specializations of enclosing generic contexts (e.g. for non-generic methods
/// of a specialized generic class).
#[returns(deref)]
type_mappings: Box<[TypeMapping<'db, 'db>]>,
/// Contains a potentially modified signature for this function literal, in case certain operations
/// (like type mappings) have been applied to it.
///
/// See also: [`FunctionLiteral::updated_signature`].
#[returns(as_ref)]
updated_signature: Option<CallableSignature<'db>>,

/// Contains a potentially modified signature for the last overload or the implementation of this
/// function literal, in case certain operations (like type mappings) have been applied to it.
///
/// See also: [`FunctionLiteral::last_definition_signature`].
#[returns(as_ref)]
updated_last_definition_signature: Option<Signature<'db>>,
}

// The Salsa heap is tracked separately.
Expand All @@ -707,8 +696,13 @@ pub(super) fn walk_function_type<'db, V: super::visitor::TypeVisitor<'db> + ?Siz
visitor: &V,
) {
walk_function_literal(db, function.literal(db), visitor);
for mapping in function.type_mappings(db) {
walk_type_mapping(db, mapping, visitor);
if let Some(callable_signature) = function.updated_signature(db) {
for signature in &callable_signature.overloads {
walk_signature(db, signature, visitor);
}
}
if let Some(signature) = function.updated_last_definition_signature(db) {
walk_signature(db, signature, visitor);
}
}

Expand All @@ -722,21 +716,41 @@ impl<'db> FunctionType<'db> {
let literal = self
.literal(db)
.with_inherited_generic_context(db, inherited_generic_context);
Self::new(db, literal, self.type_mappings(db))
let updated_signature = self.updated_signature(db).map(|signature| {
signature.with_inherited_generic_context(Some(inherited_generic_context))
});
let updated_last_definition_signature =
self.updated_last_definition_signature(db).map(|signature| {
signature
.clone()
.with_inherited_generic_context(Some(inherited_generic_context))
});
Self::new(
db,
literal,
updated_signature,
updated_last_definition_signature,
)
}

pub(crate) fn with_type_mapping<'a>(
pub(crate) fn apply_type_mapping_impl<'a>(
self,
db: &'db dyn Db,
type_mapping: &TypeMapping<'a, 'db>,
visitor: &ApplyTypeMappingVisitor<'db>,
) -> Self {
let type_mappings: Box<[_]> = self
.type_mappings(db)
.iter()
.cloned()
.chain(std::iter::once(type_mapping.to_owned()))
.collect();
Self::new(db, self.literal(db), type_mappings)
let updated_signature =
self.signature(db)
.apply_type_mapping_impl(db, type_mapping, visitor);
let updated_last_definition_signature = self
.last_definition_signature(db)
.apply_type_mapping_impl(db, type_mapping, visitor);
Self::new(
db,
self.literal(db),
Some(updated_signature),
Some(updated_last_definition_signature),
)
}

pub(crate) fn with_dataclass_transformer_params(
Expand All @@ -752,7 +766,7 @@ impl<'db> FunctionType<'db> {
.with_dataclass_transformer_params(db, params);
let literal =
FunctionLiteral::new(db, last_definition, literal.inherited_generic_context(db));
Self::new(db, literal, self.type_mappings(db))
Self::new(db, literal, None, None)
}

/// Returns the [`File`] in which this function is defined.
Expand Down Expand Up @@ -907,7 +921,9 @@ impl<'db> FunctionType<'db> {
/// would depend on the function's AST and rerun for every change in that file.
#[salsa::tracked(returns(ref), cycle_fn=signature_cycle_recover, cycle_initial=signature_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
pub(crate) fn signature(self, db: &'db dyn Db) -> CallableSignature<'db> {
self.literal(db).signature(db, self.type_mappings(db))
self.updated_signature(db)
.cloned()
.unwrap_or_else(|| self.literal(db).signature(db))
}

/// Typed externally-visible signature of the last overload or implementation of this function.
Expand All @@ -926,8 +942,9 @@ impl<'db> FunctionType<'db> {
heap_size=ruff_memory_usage::heap_size,
)]
pub(crate) fn last_definition_signature(self, db: &'db dyn Db) -> Signature<'db> {
self.literal(db)
.last_definition_signature(db, self.type_mappings(db))
self.updated_last_definition_signature(db)
.cloned()
.unwrap_or_else(|| self.literal(db).last_definition_signature(db))
}

/// Convert the `FunctionType` into a [`CallableType`].
Expand Down Expand Up @@ -1017,12 +1034,19 @@ impl<'db> FunctionType<'db> {
}

pub(crate) fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
let mappings: Box<_> = self
.type_mappings(db)
.iter()
.map(|mapping| mapping.normalized_impl(db, visitor))
.collect();
Self::new(db, self.literal(db).normalized_impl(db, visitor), mappings)
let literal = self.literal(db).normalized_impl(db, visitor);
let updated_signature = self
.updated_signature(db)
.map(|signature| signature.normalized_impl(db, visitor));
let updated_last_definition_signature = self
.updated_last_definition_signature(db)
.map(|signature| signature.normalized_impl(db, visitor));
Self::new(
db,
literal,
updated_signature,
updated_last_definition_signature,
)
}
}

Expand Down
Loading
Loading