[red-knot] Refactor KnownFunction::takes_expression_arguments()

crates/red_knot_python_semantic/src/types.rs

@@ -3428,37 +3428,77 @@ impl KnownFunction {
         }
     }
 
-    /// Returns a `u32` bitmask specifying whether or not
-    /// arguments given to a particular function
-    /// should be interpreted as type expressions or value expressions.
-    ///
-    /// The argument is treated as a type expression
-    /// when the corresponding bit is `1`.
-    /// The least-significant (right-most) bit corresponds to
-    /// the argument at the index 0 and so on.
-    ///
-    /// For example, `assert_type()` has the bitmask value of `0b10`.
-    /// This means the second argument is a type expression and the first a value expression.
-    const fn takes_type_expression_arguments(self) -> u32 {
-        const ALL_VALUES: u32 = 0b0;
-        const SINGLE_TYPE: u32 = 0b1;
-        const TYPE_TYPE: u32 = 0b11;
-        const VALUE_TYPE: u32 = 0b10;
+    /// Return the [`ParameterExpectations`] for this function.
+    const fn takes_type_expression_arguments(self) -> ParameterExpectations {
+        match self {
+            Self::IsFullyStatic | Self::IsSingleton | Self::IsSingleValued => {
+                ParameterExpectations::SINGLE_TYPE_EXPRESSION
+            }
+
+            Self::IsEquivalentTo
+            | Self::IsSubtypeOf
+            | Self::IsAssignableTo
+            | Self::IsDisjointFrom => ParameterExpectations::TWO_TYPE_EXPRESSIONS,
+
+            Self::AssertType => ParameterExpectations::VALUE_EXPRESSION_AND_TYPE_EXPRESSION,
+
+            Self::ConstraintFunction(_)
+            | Self::Len
+            | Self::Final
+            | Self::NoTypeCheck
+            | Self::RevealType
+            | Self::StaticAssert => ParameterExpectations::AllValueExpressions,
+        }
+    }
+}
 
+/// Describes whether the parameters in a function expect value expressions or type expressions.
+///
+/// Whether a specific parameter in the function expects a type expression can be queried
+/// using [`ParameterExpectations::expectation_at_index`].
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
+enum ParameterExpectations {
+    /// All parameters in the function expect value expressions
+    #[default]
+    AllValueExpressions,
+    /// The function is special-cased by the type system: one or more parameters expect type expressions
+    /// rather than value expressions.
+    SpecialCased(&'static [ParameterExpectation]),
+}
+
+impl ParameterExpectations {
+    const SINGLE_TYPE_EXPRESSION: Self = Self::SpecialCased(&[ParameterExpectation::Type]);
+
+    const TWO_TYPE_EXPRESSIONS: Self =
+        Self::SpecialCased(&[ParameterExpectation::Type, ParameterExpectation::Type]);
+
+    const VALUE_EXPRESSION_AND_TYPE_EXPRESSION: Self =
+        Self::SpecialCased(&[ParameterExpectation::Value, ParameterExpectation::Type]);
+
+    /// Query whether the parameter at `parameter_index` expects a value expression or a type expression
+    fn expectation_at_index(self, parameter_index: usize) -> ParameterExpectation {
         match self {
-            KnownFunction::IsEquivalentTo => TYPE_TYPE,
-            KnownFunction::IsSubtypeOf => TYPE_TYPE,
-            KnownFunction::IsAssignableTo => TYPE_TYPE,
-            KnownFunction::IsDisjointFrom => TYPE_TYPE,
-            KnownFunction::IsFullyStatic => SINGLE_TYPE,
-            KnownFunction::IsSingleton => SINGLE_TYPE,
-            KnownFunction::IsSingleValued => SINGLE_TYPE,
-            KnownFunction::AssertType => VALUE_TYPE,
-            _ => ALL_VALUES,
+            Self::AllValueExpressions => ParameterExpectation::Value,
+            Self::SpecialCased(expectations) => expectations
+                .get(parameter_index)
+                .copied()
+                .unwrap_or_default(),
         }
     }
 }
 
+/// Whether a single parameter in a given function expects a value expression or a [type expression]
+///
+/// [type expression]: https://typing.readthedocs.io/en/latest/spec/annotations.html#type-and-annotation-expressions
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
+enum ParameterExpectation {
+    /// The parameter expects a value expression
+    #[default]
+    Value,
+    /// The parameter expects a type expression
+    Type,
+}
+
 #[salsa::interned]
 pub struct ModuleLiteralType<'db> {
     /// The file in which this module was imported.

crates/red_knot_python_semantic/src/types/infer.rs

@@ -83,6 +83,7 @@ use super::slots::check_class_slots;
 use super::string_annotation::{
     parse_string_annotation, BYTE_STRING_TYPE_ANNOTATION, FSTRING_TYPE_ANNOTATION,
 };
+use super::{ParameterExpectation, ParameterExpectations};
 
 /// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope.
 /// Use when checking a scope, or needing to provide a type for an arbitrary expression in the
@@ -956,7 +957,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         self.infer_type_parameters(type_params);
 
         if let Some(arguments) = class.arguments.as_deref() {
-            self.infer_arguments(arguments, 0b0);
+            self.infer_arguments(arguments, ParameterExpectations::default());
         }
     }
 
@@ -2601,18 +2602,15 @@ impl<'db> TypeInferenceBuilder<'db> {
     fn infer_arguments<'a>(
         &mut self,
         arguments: &'a ast::Arguments,
-        infer_as_type_expressions: u32,
+        parameter_expectations: ParameterExpectations,
     ) -> CallArguments<'a, 'db> {
         arguments
             .arguments_source_order()
             .enumerate()
             .map(|(index, arg_or_keyword)| {
-                let infer_argument_type = if index < u32::BITS as usize
-                    && infer_as_type_expressions & (1 << index) != 0
-                {
-                    Self::infer_type_expression
-                } else {
-                    Self::infer_expression
+                let infer_argument_type = match parameter_expectations.expectation_at_index(index) {
+                    ParameterExpectation::Type => Self::infer_type_expression,
+                    ParameterExpectation::Value => Self::infer_expression,
                 };
 
                 match arg_or_keyword {
@@ -3161,7 +3159,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             .into_function_literal()
             .and_then(|f| f.known(self.db()))
             .map(KnownFunction::takes_type_expression_arguments)
-            .unwrap_or(0b0);
+            .unwrap_or_default();
 
         let call_arguments = self.infer_arguments(arguments, infer_arguments_as_type_expressions);
         function_type