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

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/ConeInferenceContext.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.declarations.utils.modality
 import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
@@ -466,7 +467,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         return this.isSomeFunctionalType(session)
     }
 
-    private fun ConeKotlinType.isTypeOrSubtypeOf(predicate: (ConeKotlinType) -> Boolean): Boolean {
+    fun ConeKotlinType.isTypeOrSubtypeOf(predicate: (ConeKotlinType) -> Boolean): Boolean {
         return predicate(this) || DFS.dfsFromNode(
             this,
             {
@@ -489,13 +490,6 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         )
     }
 
-    override fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean {
-        require(this is ConeKotlinType)
-        return isTypeOrSubtypeOf {
-            (it.lowerBoundIfFlexible() as ConeKotlinType).isSuspendOrKSuspendFunctionType(session)
-        }
-    }
-
     override fun KotlinTypeMarker.isExtensionFunctionType(): Boolean {
         require(this is ConeKotlinType)
         return (this.lowerBoundIfFlexible() as? ConeKotlinType)?.isExtensionFunctionType(session) == true
@@ -536,24 +530,17 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         }
     }
 
-    override fun getFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
-        val classId = if (isSuspend) {
-            StandardClassIds.SuspendFunctionN(parametersNumber)
-        } else {
-            StandardClassIds.FunctionN(parametersNumber)
-        }
-        return session.symbolProvider.getClassLikeSymbolByClassId(classId)?.toLookupTag()
-            ?: error("Can't find Function type")
+    override fun KotlinTypeMarker.functionalTypeKind(): FunctionalTypeKind? {
+        require(this is ConeKotlinType)
+        return this.functionalTypeKind(session)
     }
 
-    override fun getKFunctionTypeConstructor(parametersNumber: Int, isSuspend: Boolean): TypeConstructorMarker {
-        val classId = if (isSuspend) {
-            StandardClassIds.KSuspendFunctionN(parametersNumber)
-        } else {
-            StandardClassIds.KFunctionN(parametersNumber)
-        }
-        return session.symbolProvider.getClassLikeSymbolByClassId(classId)?.toLookupTag()
-            ?: error("Can't find KFunction type")
+    override fun getNonReflectFunctionTypeConstructor(parametersNumber: Int, kind: FunctionalTypeKind): TypeConstructorMarker {
+        return kind.nonReflectKind().numberedClassId(parametersNumber).toLookupTag()
+    }
+
+    override fun getReflectFunctionTypeConstructor(parametersNumber: Int, kind: FunctionalTypeKind): TypeConstructorMarker {
+        return kind.reflectKind().numberedClassId(parametersNumber).toLookupTag()
     }
 
     override fun createTypeWithAlternativeForIntersectionResult(

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/FirFunctionalTypeKindServiceImpl.kt

@@ -8,11 +8,15 @@ package org.jetbrains.kotlin.fir.types
 import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKindExtractor
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
 import org.jetbrains.kotlin.fir.extensions.FirFunctionalTypeKindExtension
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.functionalTypeKindExtensions
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 
 class FirFunctionalTypeKindServiceImpl(session: FirSession) : FirFunctionalTypeKindService() {
+    private val nonReflectKindsFromExtensions = mutableListOf<FunctionalTypeKind>()
+
     override val extractor: FunctionalTypeKindExtractor = run {
         val kinds = buildList {
             add(FunctionalTypeKind.Function)
@@ -26,6 +30,7 @@ class FirFunctionalTypeKindServiceImpl(session: FirSession) : FirFunctionalTypeK
                     require(reflectKind.nonReflectKind() == nonReflectKind)
                     add(nonReflectKind)
                     add(reflectKind)
+                    nonReflectKindsFromExtensions += nonReflectKind
                 }
             }
 
@@ -38,6 +43,31 @@ class FirFunctionalTypeKindServiceImpl(session: FirSession) : FirFunctionalTypeK
                 "There are clashing functional type kinds: $allNames"
             }
         }
+
         FunctionalTypeKindExtractor(kinds)
     }
+
+    override fun extractSingleSpecialKindForFunction(functionSymbol: FirNamedFunctionSymbol): FunctionalTypeKind? {
+        if (nonReflectKindsFromExtensions.isEmpty()) {
+            return FunctionalTypeKind.SuspendFunction.takeIf { functionSymbol.isSuspend }
+        }
+        return extractAllSpecialKindsForFunction(functionSymbol).singleOrNull()
+    }
+
+    override fun extractAllSpecialKindsForFunction(functionSymbol: FirNamedFunctionSymbol): List<FunctionalTypeKind> {
+        return buildList {
+            if (functionSymbol.isSuspend) {
+                add(FunctionalTypeKind.SuspendFunction)
+            }
+            if (nonReflectKindsFromExtensions.isNotEmpty()) {
+                for (annotationClassId in functionSymbol.resolvedAnnotationClassIds) {
+                    for (kind in nonReflectKindsFromExtensions) {
+                        if (kind.annotationOnInvokeClassId == annotationClassId) {
+                            add(kind)
+                        }
+                    }
+                }
+            }
+        }
+    }
 }

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/FunctionalTypeUtils.kt

@@ -6,8 +6,10 @@
 package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
+import org.jetbrains.kotlin.builtins.functions.isRegularFunction
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.originalForSubstitutionOverride
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -82,14 +84,12 @@ fun ConeClassLikeLookupTag.isSomeFunctionalType(session: FirSession): Boolean {
 
 // Function, KFunction
 private fun ConeKotlinType.isSimpleFunctionalType(session: FirSession, errorOnNotFunctionalType: Boolean): Boolean {
-    return isFunctionalTypeWithPredicate(session, errorOnNotFunctionalType) {
-        it == FunctionalTypeKind.Function || it == FunctionalTypeKind.KFunction
-    }
+    return isFunctionalTypeWithPredicate(session, errorOnNotFunctionalType) { it.isRegularFunction }
 }
 
 // SuspendFunction, [Custom]Function, KSuspendFunction, K[Custom]Function
-private fun ConeKotlinType.isNotSimpleFunctionalType(session: FirSession): Boolean {
-    return !isSimpleFunctionalType(session, errorOnNotFunctionalType = false)
+fun ConeKotlinType.isNotSimpleFunctionalType(session: FirSession): Boolean {
+    return isFunctionalTypeWithPredicate(session, errorOnNotFunctionalType = false) { !it.isRegularFunction }
 }
 
 // ---------------------------------------------- functional type conversions ----------------------------------------------
@@ -106,7 +106,7 @@ fun ConeKotlinType.customFunctionalTypeToSimpleFunctionalType(session: FirSessio
     } else {
         FunctionalTypeKind.Function
     }
-    return createFunctionalTypeWithNewKind(newKind)
+    return createFunctionalTypeWithNewKind(session, newKind)
 }
 
 /*
@@ -117,12 +117,17 @@ fun ConeKotlinType.customFunctionalTypeToSimpleFunctionalType(session: FirSessio
 fun ConeKotlinType.reflectFunctionalTypeToNonReflectFunctionalType(session: FirSession): ConeClassLikeType {
     val kind = functionalTypeKind(session)
     require(kind != null && kind.isReflectType)
-    return createFunctionalTypeWithNewKind(kind.nonReflectKind())
+    return createFunctionalTypeWithNewKind(session, kind.nonReflectKind())
 }
 
-private fun ConeKotlinType.createFunctionalTypeWithNewKind(kind: FunctionalTypeKind): ConeClassLikeType {
-    val functionalTypeId = ClassId(kind.packageFqName, kind.numberedClassName(typeArguments.size - 1))
-    return functionalTypeId.toLookupTag().constructClassType(typeArguments, isNullable = false, attributes = attributes)
+private fun ConeKotlinType.createFunctionalTypeWithNewKind(session: FirSession, kind: FunctionalTypeKind): ConeClassLikeType {
+    val expandedType = fullyExpandedType(session)
+    val functionalTypeId = ClassId(kind.packageFqName, kind.numberedClassName(expandedType.typeArguments.size - 1))
+    return functionalTypeId.toLookupTag().constructClassType(
+        expandedType.typeArguments,
+        isNullable = false,
+        attributes = expandedType.attributes
+    )
 }
 
 // ---------------------------------------------- functional type subtyping ----------------------------------------------
@@ -139,8 +144,8 @@ private fun ConeKotlinType.findSubtypeOfSimpleFunctionalTypeImpl(
 ): ConeKotlinType? {
     return when (this) {
         is ConeClassLikeType -> {
-            // Expect the argument type is a simple functional type.
             when {
+                // Expect the argument type is a simple functional type.
                 isNotSimpleFunctionalType(session) -> null
                 isSubtypeOfFunctionalType(session, expectedFunctionalType) -> this
                 else -> null
@@ -247,3 +252,11 @@ private fun ConeTypeProjection.typeOrDefault(default: ConeKotlinType): ConeKotli
     is ConeKotlinTypeProjection -> type
     is ConeStarProjection -> default
 }
+
+// ----------------- TODO fir utils
+
+fun FirFunction.specialFunctionalTypeKind(session: FirSession): FunctionalTypeKind? {
+    return (symbol as? FirNamedFunctionSymbol)?.let {
+        session.functionalTypeService.extractSingleSpecialKindForFunction(it)
+    }
+}