speedup DateType benchmark

Author: AlexWaygood

crates/ty_python_semantic/resources/mdtest/protocols.md

@@ -418,7 +418,6 @@ reveal_protocol_interface(SupportsIndex)
 # error: [revealed-type] "Revealed protocol interface: `{"__abs__": MethodMember(`(self) -> Unknown`)}`"
 reveal_protocol_interface(SupportsAbs)
 
-
 # error: [invalid-argument-type] "Invalid argument to `reveal_protocol_interface`: Only protocol classes can be passed to `reveal_protocol_interface`"
 reveal_protocol_interface(int)
 # error: [invalid-argument-type] "Argument to function `reveal_protocol_interface` is incorrect: Expected `type`, found `Literal["foo"]`"

crates/ty_python_semantic/src/types/protocol_class.rs

@@ -7,6 +7,7 @@ use ruff_python_ast::name::Name;
 
 use super::TypeVarVariance;
 use crate::semantic_index::place_table;
+use crate::types::generics::Specialization;
 use crate::{
     Db, FxOrderSet,
     place::{Boundness, Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
@@ -85,6 +86,7 @@ impl<'db> Deref for ProtocolClass<'db> {
 pub(super) struct ProtocolInterface<'db> {
     #[returns(ref)]
     inner: BTreeMap<Name, ProtocolMemberData<'db>>,
+    specialization: Option<Specialization<'db>>,
 }
 
 impl get_size2::GetSize for ProtocolInterface<'_> {}
@@ -128,11 +130,11 @@ impl<'db> ProtocolInterface<'db> {
                 )
             })
             .collect();
-        Self::new(db, members)
+        Self::new(db, members, None)
     }
 
     fn empty(db: &'db dyn Db) -> Self {
-        Self::new(db, BTreeMap::default())
+        Self::new(db, BTreeMap::default(), None)
     }
 
     pub(super) fn members<'a>(
@@ -142,17 +144,22 @@ impl<'db> ProtocolInterface<'db> {
     where
         'db: 'a,
     {
-        self.inner(db).iter().map(|(name, data)| ProtocolMember {
-            name,
-            kind: data.kind,
-            qualifiers: data.qualifiers,
-        })
+        let specialization = self.specialization(db);
+        self.inner(db)
+            .iter()
+            .map(move |(name, data)| ProtocolMember {
+                name,
+                kind: data.kind.apply_optional_specialization(db, specialization),
+                qualifiers: data.qualifiers,
+            })
     }
 
     fn member_by_name<'a>(self, db: &'db dyn Db, name: &'a str) -> Option<ProtocolMember<'a, 'db>> {
         self.inner(db).get(name).map(|data| ProtocolMember {
             name,
-            kind: data.kind,
+            kind: data
+                .kind
+                .apply_optional_specialization(db, self.specialization(db)),
             qualifiers: data.qualifiers,
         })
     }
@@ -186,6 +193,8 @@ impl<'db> ProtocolInterface<'db> {
                 .iter()
                 .map(|(name, data)| (name.clone(), data.normalized_impl(db, visitor)))
                 .collect::<BTreeMap<_, _>>(),
+            self.specialization(db)
+                .map(|s| s.normalized_impl(db, visitor)),
         )
     }
 
@@ -196,6 +205,7 @@ impl<'db> ProtocolInterface<'db> {
                 .iter()
                 .map(|(name, data)| (name.clone(), data.materialize(db, variance)))
                 .collect::<BTreeMap<_, _>>(),
+            self.specialization(db).map(|s| s.materialize(db, variance)),
         )
     }
 
@@ -204,17 +214,24 @@ impl<'db> ProtocolInterface<'db> {
         db: &'db dyn Db,
         type_mapping: &TypeMapping<'a, 'db>,
     ) -> Self {
+        let mut visitor = TypeTransformer::default();
+
         Self::new(
             db,
             self.inner(db)
                 .iter()
                 .map(|(name, data)| {
                     (
                         name.clone(),
-                        data.apply_type_mapping(db, type_mapping).normalized(db),
+                        data.apply_type_mapping(db, type_mapping)
+                            .normalized_impl(db, &mut visitor),
                     )
                 })
                 .collect::<BTreeMap<_, _>>(),
+            self.specialization(db).map(|s| {
+                s.apply_type_mapping(db, type_mapping)
+                    .normalized_impl(db, &mut visitor)
+            }),
         )
     }
 
@@ -237,8 +254,13 @@ impl<'db> ProtocolInterface<'db> {
         impl std::fmt::Display for ProtocolInterfaceDisplay<'_> {
             fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                 f.write_char('{')?;
-                for (i, (name, data)) in self.interface.inner(self.db).iter().enumerate() {
-                    write!(f, "\"{name}\": {data}", data = data.display(self.db))?;
+                for (i, member) in self.interface.members(self.db).enumerate() {
+                    let ProtocolMember {
+                        name,
+                        kind,
+                        qualifiers: _,
+                    } = member;
+                    write!(f, "\"{name}\": {kind}", kind = kind.display(self.db))?;
                     if i < self.interface.inner(self.db).len() - 1 {
                         f.write_str(", ")?;
                     }
@@ -261,10 +283,6 @@ pub(super) struct ProtocolMemberData<'db> {
 }
 
 impl<'db> ProtocolMemberData<'db> {
-    fn normalized(&self, db: &'db dyn Db) -> Self {
-        self.normalized_impl(db, &mut TypeTransformer::default())
-    }
-
     fn normalized_impl(&self, db: &'db dyn Db, visitor: &mut TypeTransformer<'db>) -> Self {
         Self {
             kind: self.kind.normalized_impl(db, visitor),
@@ -293,41 +311,6 @@ impl<'db> ProtocolMemberData<'db> {
             qualifiers: self.qualifiers,
         }
     }
-
-    fn display(&self, db: &'db dyn Db) -> impl std::fmt::Display {
-        struct ProtocolMemberDataDisplay<'db> {
-            db: &'db dyn Db,
-            data: ProtocolMemberKind<'db>,
-        }
-
-        impl std::fmt::Display for ProtocolMemberDataDisplay<'_> {
-            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-                match self.data {
-                    ProtocolMemberKind::Method(callable) => {
-                        write!(f, "MethodMember(`{}`)", callable.display(self.db))
-                    }
-                    ProtocolMemberKind::Property(property) => {
-                        let mut d = f.debug_struct("PropertyMember");
-                        if let Some(getter) = property.getter(self.db) {
-                            d.field("getter", &format_args!("`{}`", &getter.display(self.db)));
-                        }
-                        if let Some(setter) = property.setter(self.db) {
-                            d.field("setter", &format_args!("`{}`", &setter.display(self.db)));
-                        }
-                        d.finish()
-                    }
-                    ProtocolMemberKind::Other(ty) => {
-                        write!(f, "AttributeMember(`{}`)", ty.display(self.db))
-                    }
-                }
-            }
-        }
-
-        ProtocolMemberDataDisplay {
-            db,
-            data: self.kind,
-        }
-    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, salsa::Update, Hash)]
@@ -391,6 +374,59 @@ impl<'db> ProtocolMemberKind<'db> {
             }
         }
     }
+
+    fn apply_optional_specialization(
+        self,
+        db: &'db dyn Db,
+        specialization: Option<Specialization<'db>>,
+    ) -> Self {
+        let Some(specialization) = specialization else {
+            return self;
+        };
+        match self {
+            ProtocolMemberKind::Method(callable) => ProtocolMemberKind::Method(
+                callable.apply_type_mapping(db, &TypeMapping::Specialization(specialization)),
+            ),
+            ProtocolMemberKind::Property(property) => ProtocolMemberKind::Property(
+                property.apply_type_mapping(db, &TypeMapping::Specialization(specialization)),
+            ),
+            ProtocolMemberKind::Other(ty) => {
+                ProtocolMemberKind::Other(ty.apply_specialization(db, specialization))
+            }
+        }
+    }
+
+    fn display(self, db: &'db dyn Db) -> impl std::fmt::Display {
+        struct ProtocolMemberKindDisplay<'db> {
+            db: &'db dyn Db,
+            kind: ProtocolMemberKind<'db>,
+        }
+
+        impl std::fmt::Display for ProtocolMemberKindDisplay<'_> {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                match self.kind {
+                    ProtocolMemberKind::Method(callable) => {
+                        write!(f, "MethodMember(`{}`)", callable.display(self.db))
+                    }
+                    ProtocolMemberKind::Property(property) => {
+                        let mut d = f.debug_struct("PropertyMember");
+                        if let Some(getter) = property.getter(self.db) {
+                            d.field("getter", &format_args!("`{}`", &getter.display(self.db)));
+                        }
+                        if let Some(setter) = property.setter(self.db) {
+                            d.field("setter", &format_args!("`{}`", &setter.display(self.db)));
+                        }
+                        d.finish()
+                    }
+                    ProtocolMemberKind::Other(ty) => {
+                        write!(f, "AttributeMember(`{}`)", ty.display(self.db))
+                    }
+                }
+            }
+        }
+
+        ProtocolMemberKindDisplay { db, kind: self }
+    }
 }
 
 /// A single member of a protocol interface.
@@ -532,13 +568,10 @@ fn cached_protocol_interface<'db>(
 ) -> ProtocolInterface<'db> {
     let mut members = BTreeMap::default();
 
-    for (parent_protocol, specialization) in class
+    for parent_protocol in class
         .iter_mro(db)
         .filter_map(ClassBase::into_class)
-        .filter_map(|class| {
-            let (class, specialization) = class.class_literal(db);
-            Some((class.into_protocol_class(db)?, specialization))
-        })
+        .filter_map(|class| class.class_literal(db).0.into_protocol_class(db))
     {
         let parent_scope = parent_protocol.class_literal(db).0.body_scope(db);
         let use_def_map = use_def_map(db, parent_scope);
@@ -581,8 +614,6 @@ fn cached_protocol_interface<'db>(
                 })
                 .filter(|(name, _, _, _)| !excluded_from_proto_members(name))
                 .map(|(name, ty, qualifiers, bound_on_class)| {
-                    let ty = ty.apply_optional_specialization(db, specialization);
-
                     let kind = match (ty, bound_on_class) {
                         // TODO: if the getter or setter is a function literal, we should
                         // upcast it to a `CallableType` so that two protocols with identical property
@@ -607,7 +638,7 @@ fn cached_protocol_interface<'db>(
         );
     }
 
-    ProtocolInterface::new(db, members)
+    ProtocolInterface::new(db, members, class.class_literal(db).1)
 }
 
 #[allow(clippy::trivially_copy_pass_by_ref)]