[FIR] Generify resolution and inference of arguments of functional types

core/compiler.common/src/org/jetbrains/kotlin/builtins/functions/FunctionalTypeKind.kt

@@ -85,6 +85,8 @@ abstract class FunctionalTypeKind internal constructor(
 
     fun numberedClassName(arity: Int): Name = Name.identifier("$classNamePrefix$arity")
 
+    fun numberedClassId(arity: Int): ClassId = ClassId(packageFqName, numberedClassName(arity))
+
     override fun toString(): String {
         return "$packageFqName.${classNamePrefix}N"
     }
@@ -141,5 +143,7 @@ val FunctionalTypeKind.isBuiltin: Boolean
     }
 
 val FunctionalTypeKind.isSuspendType: Boolean
-    get() = this == FunctionalTypeKind.SuspendFunction || this == FunctionalTypeKind.KSuspendFunction
+    get() = this.nonReflectKind() == FunctionalTypeKind.SuspendFunction
 
+val FunctionalTypeKind.isRegularFunction: Boolean
+    get() = this.nonReflectKind() == FunctionalTypeKind.Function

core/compiler.common/src/org/jetbrains/kotlin/types/model/TypeSystemContext.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.types.model
 
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.resolve.checkers.EmptyIntersectionTypeChecker
 import org.jetbrains.kotlin.resolve.checkers.EmptyIntersectionTypeInfo
 import org.jetbrains.kotlin.types.*
@@ -247,9 +248,11 @@ interface TypeSystemInferenceExtensionContext : TypeSystemContext, TypeSystemBui
 
     fun KotlinTypeMarker.isSignedOrUnsignedNumberType(): Boolean
 
+    // ------------- functional type utils -------------
+
     fun KotlinTypeMarker.isFunctionOrKFunctionWithAnySuspendability(): Boolean
 
-    fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean
+    fun KotlinTypeMarker.functionalTypeKind(): FunctionalTypeKind?
 
     fun KotlinTypeMarker.isExtensionFunctionType(): Boolean
 
@@ -257,12 +260,14 @@ interface TypeSystemInferenceExtensionContext : TypeSystemContext, TypeSystemBui
 
     fun KotlinTypeMarker.getFunctionalTypeFromSupertypes(): KotlinTypeMarker
 
+    fun getNonReflectFunctionTypeConstructor(parametersNumber: Int, kind: FunctionalTypeKind): TypeConstructorMarker
+
+    fun getReflectFunctionTypeConstructor(parametersNumber: Int, kind: FunctionalTypeKind): TypeConstructorMarker
+
+    // -------------------------------------------------
+
     fun StubTypeMarker.getOriginalTypeVariable(): TypeVariableTypeConstructorMarker
 
-    fun getFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker
-
-    fun getKFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker
-
     private fun <T> KotlinTypeMarker.extractTypeOf(to: MutableSet<T>, getIfApplicable: (TypeConstructorMarker) -> T?) {
         for (i in 0 until argumentsCount()) {
             val argument = getArgument(i)

core/descriptors/src/org/jetbrains/kotlin/builtins/functionTypes.kt

@@ -59,17 +59,20 @@ val KotlinType.isBuiltinFunctionalTypeOrSubtype: Boolean
 fun KotlinType.isFunctionTypeOrSubtype(predicate: (KotlinType) -> Boolean): Boolean =
     isTypeOrSubtypeOf { it.isFunctionType && predicate(it) }
 
+val KotlinType.functionalTypeKind: FunctionalTypeKind?
+    get() = constructor.declarationDescriptor?.getFunctionalClassKind()
+
 val KotlinType.isFunctionType: Boolean
-    get() = constructor.declarationDescriptor?.getFunctionalClassKind() == FunctionalTypeKind.Function
+    get() = functionalTypeKind == FunctionalTypeKind.Function
 
 val KotlinType.isKFunctionType: Boolean
-    get() = constructor.declarationDescriptor?.getFunctionalClassKind() == FunctionalTypeKind.KFunction
+    get() = functionalTypeKind == FunctionalTypeKind.KFunction
 
 val KotlinType.isSuspendFunctionType: Boolean
-    get() = constructor.declarationDescriptor?.getFunctionalClassKind() == FunctionalTypeKind.SuspendFunction
+    get() = functionalTypeKind == FunctionalTypeKind.SuspendFunction
 
 val KotlinType.isKSuspendFunctionType: Boolean
-    get() = constructor.declarationDescriptor?.getFunctionalClassKind() == FunctionalTypeKind.KSuspendFunction
+    get() = functionalTypeKind == FunctionalTypeKind.KSuspendFunction
 
 val KotlinType.isFunctionOrSuspendFunctionType: Boolean
     get() = isFunctionType || isSuspendFunctionType

core/descriptors/src/org/jetbrains/kotlin/types/checker/ClassicTypeSystemContext.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.types.checker
 
 import org.jetbrains.kotlin.builtins.*
 import org.jetbrains.kotlin.builtins.StandardNames.FqNames
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
@@ -849,11 +850,6 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
         return this.isFunctionOrKFunctionTypeWithAnySuspendability
     }
 
-    override fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean {
-        require(this is KotlinType, this::errorMessage)
-        return this.isSuspendFunctionTypeOrSubtype
-    }
-
     override fun KotlinTypeMarker.isExtensionFunctionType(): Boolean {
         require(this is KotlinType, this::errorMessage)
         return this.isBuiltinExtensionFunctionalType
@@ -869,12 +865,25 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
         return this.extractFunctionalTypeFromSupertypes()
     }
 
-    override fun getFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
-        return getFunctionDescriptor(builtIns, parametersNumber, isSuspend).typeConstructor
+    override fun KotlinTypeMarker.functionalTypeKind(): FunctionalTypeKind? {
+        require(this is KotlinType)
+        return this.functionalTypeKind
     }
 
-    override fun getKFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
-        return getKFunctionDescriptor(builtIns, parametersNumber, isSuspend).typeConstructor
+    override fun getNonReflectFunctionTypeConstructor(parametersNumber: Int, kind: FunctionalTypeKind): TypeConstructorMarker {
+        return getFunctionDescriptor(
+            builtIns,
+            parametersNumber,
+            isSuspendFunction = kind.nonReflectKind() == FunctionalTypeKind.SuspendFunction
+        ).typeConstructor
+    }
+
+    override fun getReflectFunctionTypeConstructor(parametersNumber: Int, kind: FunctionalTypeKind): TypeConstructorMarker {
+        return getKFunctionDescriptor(
+            builtIns,
+            parametersNumber,
+            isSuspendFunction = kind.reflectKind() == FunctionalTypeKind.KSuspendFunction
+        ).typeConstructor
     }
 
     override fun createSubstitutorForSuperTypes(baseType: KotlinTypeMarker): TypeSubstitutorMarker? {