[ty] Add Salsa caching to ClassLiteral::fields (#21512)

Author: AlexWaygood

Cargo.lock

@@ -32,7 +32,7 @@ camino = { workspace = true }
 colored = { workspace = true }
 compact_str = { workspace = true }
 drop_bomb = { workspace = true }
-get-size2 = { workspace = true }
+get-size2 = { workspace = true, features = ["indexmap"]}
 indexmap = { workspace = true }
 itertools = { workspace = true }
 ordermap = { workspace = true }

crates/ty/docs/rules.md

@@ -396,7 +396,7 @@ impl<'db> BoundSuperType<'db> {
                 let mut key_builder = UnionBuilder::new(db);
                 let mut value_builder = UnionBuilder::new(db);
                 for (name, field) in td.items(db) {
-                    key_builder = key_builder.add(Type::string_literal(db, &name));
+                    key_builder = key_builder.add(Type::string_literal(db, name));
                     value_builder = value_builder.add(field.declared_ty);
                 }
                 return delegate_to(

crates/ty_python_semantic/Cargo.toml

@@ -3036,8 +3036,8 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         if let Type::TypedDict(typed_dict) = argument_type {
             for (argument_type, parameter_index) in typed_dict
                 .items(self.db)
-                .iter()
-                .map(|(_, field)| field.declared_ty)
+                .values()
+                .map(|field| field.declared_ty)
                 .zip(&self.argument_matches[argument_index].parameters)
             {
                 self.check_argument_type(

crates/ty_python_semantic/src/types/bound_super.rs

@@ -7,7 +7,6 @@ use super::{
     SpecialFormType, SubclassOfType, Truthiness, Type, TypeQualifiers, class_base::ClassBase,
     function::FunctionType, infer_expression_type, infer_unpack_types,
 };
-use crate::FxOrderMap;
 use crate::module_resolver::KnownModule;
 use crate::place::TypeOrigin;
 use crate::semantic_index::definition::{Definition, DefinitionState};
@@ -128,7 +127,7 @@ fn try_metaclass_cycle_initial<'db>(
 }
 
 /// A category of classes with code generation capabilities (with synthesized methods).
-#[derive(Clone, Copy, Debug, PartialEq, salsa::Update, get_size2::GetSize)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
 pub(crate) enum CodeGeneratorKind<'db> {
     /// Classes decorated with `@dataclass` or similar dataclass-like decorators
     DataclassLike(Option<DataclassTransformerParams<'db>>),
@@ -1253,7 +1252,7 @@ impl MethodDecorator {
 }
 
 /// Kind-specific metadata for different types of fields
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
 pub(crate) enum FieldKind<'db> {
     /// `NamedTuple` field metadata
     NamedTuple { default_ty: Option<Type<'db>> },
@@ -1281,7 +1280,7 @@ pub(crate) enum FieldKind<'db> {
 }
 
 /// Metadata regarding a dataclass field/attribute or a `TypedDict` "item" / key-value pair.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
 pub(crate) struct Field<'db> {
     /// The declared type of the field
     pub(crate) declared_ty: Type<'db>,
@@ -2329,7 +2328,8 @@ impl<'db> ClassLiteral<'db> {
                     || kw_only.unwrap_or(has_dataclass_param(DataclassFlags::KW_ONLY));
 
                 // Use the alias name if provided, otherwise use the field name
-                let parameter_name = alias.map(Name::new).unwrap_or(field_name);
+                let parameter_name =
+                    Name::new(alias.map(|alias| &**alias).unwrap_or(&**field_name));
 
                 let mut parameter = if is_kw_only {
                     Parameter::keyword_only(parameter_name)
@@ -2595,7 +2595,7 @@ impl<'db> ClassLiteral<'db> {
             (CodeGeneratorKind::TypedDict, "get") => {
                 let overloads = self
                     .fields(db, specialization, field_policy)
-                    .into_iter()
+                    .iter()
                     .flat_map(|(name, field)| {
                         let key_type =
                             Type::StringLiteral(StringLiteralType::new(db, name.as_str()));
@@ -2824,12 +2824,13 @@ impl<'db> ClassLiteral<'db> {
     /// Returns a list of all annotated attributes defined in this class, or any of its superclasses.
     ///
     /// See [`ClassLiteral::own_fields`] for more details.
+    #[salsa::tracked(returns(ref), heap_size=get_size2::GetSize::get_heap_size)]
     pub(crate) fn fields(
         self,
         db: &'db dyn Db,
         specialization: Option<Specialization<'db>>,
-        field_policy: CodeGeneratorKind,
-    ) -> FxOrderMap<Name, Field<'db>> {
+        field_policy: CodeGeneratorKind<'db>,
+    ) -> FxIndexMap<Name, Field<'db>> {
         if field_policy == CodeGeneratorKind::NamedTuple {
             // NamedTuples do not allow multiple inheritance, so it is sufficient to enumerate the
             // fields of this class only.
@@ -2877,8 +2878,8 @@ impl<'db> ClassLiteral<'db> {
         db: &'db dyn Db,
         specialization: Option<Specialization<'db>>,
         field_policy: CodeGeneratorKind,
-    ) -> FxOrderMap<Name, Field<'db>> {
-        let mut attributes = FxOrderMap::default();
+    ) -> FxIndexMap<Name, Field<'db>> {
+        let mut attributes = FxIndexMap::default();
 
         let class_body_scope = self.body_scope(db);
         let table = place_table(db, class_body_scope);

crates/ty_python_semantic/src/types/call/bind.rs

@@ -26,9 +26,7 @@ use crate::types::{
     ProtocolInstanceType, SpecialFormType, SubclassOfInner, Type, TypeContext, binding_type,
     infer_isolated_expression, protocol_class::ProtocolClass,
 };
-use crate::{
-    Db, DisplaySettings, FxIndexMap, FxOrderMap, Module, ModuleName, Program, declare_lint,
-};
+use crate::{Db, DisplaySettings, FxIndexMap, Module, ModuleName, Program, declare_lint};
 use itertools::Itertools;
 use ruff_db::diagnostic::{Annotation, Diagnostic, Span, SubDiagnostic, SubDiagnosticSeverity};
 use ruff_db::source::source_text;
@@ -3192,7 +3190,7 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
     typed_dict_ty: Type<'db>,
     full_object_ty: Option<Type<'db>>,
     key_ty: Type<'db>,
-    items: &FxOrderMap<Name, Field<'db>>,
+    items: &FxIndexMap<Name, Field<'db>>,
 ) {
     let db = context.db();
     if let Some(builder) = context.report_lint(&INVALID_KEY, key_node) {
@@ -3221,7 +3219,7 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
                         .message(format_args!("TypedDict `{typed_dict_name}`"))
                 });
 
-                let existing_keys = items.iter().map(|(name, _)| name.as_str());
+                let existing_keys = items.keys();
                 if let Some(suggestion) = did_you_mean(existing_keys, key) {
                     if key_node.is_expr_string_literal() {
                         diagnostic

crates/ty_python_semantic/src/types/class.rs

@@ -922,7 +922,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
                 let kw_only_sentinel_fields: Vec<_> = class
                     .fields(self.db(), specialization, field_policy)
-                    .into_iter()
+                    .iter()
                     .filter_map(|(name, field)| {
                         field.is_kw_only_sentinel(self.db()).then_some(name)
                     })
@@ -7954,7 +7954,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                     Type::TypedDict(typed_dict_ty),
                                     None,
                                     key_ty,
-                                    &items,
+                                    items,
                                 );
                                 // Return `Unknown` to prevent the overload system from generating its own error
                                 return Type::unknown();
@@ -11295,7 +11295,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                 value_ty,
                                 None,
                                 slice_ty,
-                                &typed_dict.items(db),
+                                typed_dict.items(db),
                             );
                         } else {
                             if let Some(builder) =

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -13,7 +13,7 @@ use super::diagnostic::{
 };
 use super::{ApplyTypeMappingVisitor, Type, TypeMapping, visitor};
 use crate::types::TypeContext;
-use crate::{Db, FxOrderMap};
+use crate::{Db, FxIndexMap};
 
 use ordermap::OrderSet;
 
@@ -54,7 +54,7 @@ impl<'db> TypedDictType<'db> {
         self.defining_class
     }
 
-    pub(crate) fn items(self, db: &'db dyn Db) -> FxOrderMap<Name, Field<'db>> {
+    pub(crate) fn items(self, db: &'db dyn Db) -> &'db FxIndexMap<Name, Field<'db>> {
         let (class_literal, specialization) = self.defining_class.class_literal(db);
         class_literal.fields(db, specialization, CodeGeneratorKind::TypedDict)
     }
@@ -165,7 +165,7 @@ pub(super) fn validate_typed_dict_key_assignment<'db, 'ast>(
                 Type::TypedDict(typed_dict),
                 full_object_ty,
                 Type::string_literal(db, key),
-                &items,
+                items,
             );
         }