Use FxOrderSet to store generic context typevars

Author: dcreager

crates/red_knot_python_semantic/src/types/display.rs

@@ -14,7 +14,7 @@ use crate::types::{
     StringLiteralType, SubclassOfInner, Type, TypeVarBoundOrConstraints, TypeVarInstance,
     UnionType, WrapperDescriptorKind,
 };
-use crate::Db;
+use crate::{Db, FxOrderSet};
 use rustc_hash::FxHashMap;
 
 impl<'db> Type<'db> {
@@ -317,7 +317,7 @@ impl<'db> GenericContext<'db> {
 }
 
 pub struct DisplayGenericContext<'db> {
-    typevars: &'db [TypeVarInstance<'db>],
+    typevars: &'db FxOrderSet<TypeVarInstance<'db>>,
     db: &'db dyn Db,
 }
 
@@ -376,7 +376,7 @@ impl<'db> Specialization<'db> {
 }
 
 pub struct DisplaySpecialization<'db> {
-    typevars: &'db [TypeVarInstance<'db>],
+    typevars: &'db FxOrderSet<TypeVarInstance<'db>>,
     types: &'db [Type<'db>],
     db: &'db dyn Db,
     full: bool,

crates/red_knot_python_semantic/src/types/generics.rs

@@ -1,4 +1,3 @@
-use itertools::Itertools;
 use ruff_python_ast as ast;
 use rustc_hash::FxHashMap;
 
@@ -17,7 +16,7 @@ use crate::{Db, FxOrderSet};
 #[salsa::interned(debug)]
 pub struct GenericContext<'db> {
     #[return_ref]
-    pub(crate) variables: Box<[TypeVarInstance<'db>]>,
+    pub(crate) variables: FxOrderSet<TypeVarInstance<'db>>,
 }
 
 impl<'db> GenericContext<'db> {
@@ -27,7 +26,7 @@ impl<'db> GenericContext<'db> {
         index: &'db SemanticIndex<'db>,
         type_params_node: &ast::TypeParams,
     ) -> Self {
-        let variables: Box<[_]> = type_params_node
+        let variables: FxOrderSet<_> = type_params_node
             .iter()
             .filter_map(|type_param| Self::variable_from_type_param(db, index, type_param))
             .collect();
@@ -77,7 +76,6 @@ impl<'db> GenericContext<'db> {
         if variables.is_empty() {
             return None;
         }
-        let variables: Box<[_]> = variables.into_iter().collect();
         Some(Self::new(db, variables))
     }
 
@@ -94,7 +92,6 @@ impl<'db> GenericContext<'db> {
         if variables.is_empty() {
             return None;
         }
-        let variables: Box<[_]> = variables.into_iter().collect();
         Some(Self::new(db, variables))
     }
 
@@ -149,19 +146,17 @@ impl<'db> GenericContext<'db> {
     }
 
     pub(crate) fn is_subset_of(self, db: &'db dyn Db, other: GenericContext<'db>) -> bool {
-        for variable in self.variables(db) {
-            if !other.variables(db).iter().contains(variable) {
-                return false;
-            }
-        }
-        true
+        self.variables(db).is_subset(other.variables(db))
     }
 
+    /// Creates a specialization of this generic context. Panics if the length of `types` does not
+    /// match the number of typevars in the generic context.
     pub(crate) fn specialize(
         self,
         db: &'db dyn Db,
         types: Box<[Type<'db>]>,
     ) -> Specialization<'db> {
+        assert!(self.variables(db).len() == types.len());
         Specialization::new(db, self, types)
     }
 }
@@ -232,12 +227,11 @@ impl<'db> Specialization<'db> {
     /// Returns the type that a typevar is specialized to, or None if the typevar isn't part of
     /// this specialization.
     pub(crate) fn get(self, db: &'db dyn Db, typevar: TypeVarInstance<'db>) -> Option<Type<'db>> {
-        self.generic_context(db)
+        let index = self
+            .generic_context(db)
             .variables(db)
-            .into_iter()
-            .zip(self.types(db))
-            .find(|(var, _)| **var == typevar)
-            .map(|(_, ty)| *ty)
+            .get_index_of(&typevar)?;
+        Some(self.types(db)[index])
     }
 
     pub(crate) fn is_subtype_of(self, db: &'db dyn Db, other: Specialization<'db>) -> bool {

crates/red_knot_python_semantic/src/types/infer.rs

@@ -93,7 +93,7 @@ use crate::types::{
 };
 use crate::unpack::{Unpack, UnpackPosition};
 use crate::util::subscript::{PyIndex, PySlice};
-use crate::Db;
+use crate::{Db, FxOrderSet};
 
 use super::context::{InNoTypeCheck, InferContext};
 use super::diagnostic::{
@@ -6998,7 +6998,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         value_node: &ast::Expr,
         typevars: &[Type<'db>],
     ) -> Type<'db> {
-        let typevars: Option<Box<[_]>> = typevars
+        let typevars: Option<FxOrderSet<_>> = typevars
             .iter()
             .map(|typevar| match typevar {
                 Type::KnownInstance(KnownInstanceType::TypeVar(typevar)) => Some(*typevar),