[FIR] Add ability to extract functional type kind from annotation on functional literal type

analysis/analysis-api-fir/analysis-api-fir-generator/src/org/jetbrains/kotlin/analysis/api/fir/generator/HLDiagnosticConverter.kt

@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.config.ApiVersion
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -365,6 +366,7 @@ internal object FirToKtConversionCreator {
         ApiVersion::class,
         CallableId::class,
         ClassKind::class,
+        FunctionalTypeKind::class,
     )
 
     private val KType.kClass: KClass<*>

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDataClassConverters.kt

@@ -1814,6 +1814,15 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.AMBIGUOUS_FUNCTIONAL_TYPE_KIND) { firDiagnostic ->
+        AmbiguousFunctionalTypeKindImpl(
+            firDiagnostic.a.map { functionalTypeKind ->
+                functionalTypeKind
+            },
+            firDiagnostic as KtPsiDiagnostic,
+            token,
+        )
+    }
     add(FirErrors.NO_CONTEXT_RECEIVER) { firDiagnostic ->
         NoContextReceiverImpl(
             firSymbolBuilder.typeBuilder.buildKtType(firDiagnostic.a),

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtTypeParameterSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtVariableLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtVariableSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.config.ApiVersion
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -1292,6 +1293,11 @@ sealed class KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         abstract val candidates: List<KtSymbol>
     }
 
+    abstract class AmbiguousFunctionalTypeKind : KtFirDiagnostic<PsiElement>() {
+        override val diagnosticClass get() = AmbiguousFunctionalTypeKind::class
+        abstract val kinds: List<FunctionalTypeKind>
+    }
+
     abstract class NoContextReceiver : KtFirDiagnostic<KtElement>() {
         override val diagnosticClass get() = NoContextReceiver::class
         abstract val contextReceiverRepresentation: KtType

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtTypeParameterSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtVariableLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtVariableSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.config.ApiVersion
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -1552,6 +1553,12 @@ internal class NextAmbiguityImpl(
     override val token: KtLifetimeToken,
 ) : KtFirDiagnostic.NextAmbiguity(), KtAbstractFirDiagnostic<PsiElement>
 
+internal class AmbiguousFunctionalTypeKindImpl(
+    override val kinds: List<FunctionalTypeKind>,
+    override val firDiagnostic: KtPsiDiagnostic,
+    override val token: KtLifetimeToken,
+) : KtFirDiagnostic.AmbiguousFunctionalTypeKind(), KtAbstractFirDiagnostic<PsiElement>
+
 internal class NoContextReceiverImpl(
     override val contextReceiverRepresentation: KtType,
     override val firDiagnostic: KtPsiDiagnostic,

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/FirDiagnosticsList.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.checkers.generator.diagnostics
 import com.intellij.psi.PsiElement
 import com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.config.ApiVersion
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -567,6 +568,9 @@ object DIAGNOSTICS_LIST : DiagnosticList("FirErrors") {
         val NEXT_AMBIGUITY by error<PsiElement>(PositioningStrategy.REFERENCE_BY_QUALIFIED) {
             parameter<Collection<FirBasedSymbol<*>>>("candidates")
         }
+        val AMBIGUOUS_FUNCTIONAL_TYPE_KIND by error<PsiElement> {
+            parameter<Collection<FunctionalTypeKind>>("kinds")
+        }
     }
 
     val CONTEXT_RECEIVERS_RESOLUTION by object : DiagnosticGroup("Context receivers resolution") {

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.analysis.diagnostics
 import com.intellij.psi.PsiElement
 import com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.config.ApiVersion
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageFeature.ForbidExposingTypesInPrimaryConstructorProperties
@@ -369,6 +370,7 @@ object FirErrors {
     val ITERATOR_AMBIGUITY by error1<PsiElement, Collection<FirBasedSymbol<*>>>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
     val HAS_NEXT_FUNCTION_AMBIGUITY by error1<PsiElement, Collection<FirBasedSymbol<*>>>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
     val NEXT_AMBIGUITY by error1<PsiElement, Collection<FirBasedSymbol<*>>>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
+    val AMBIGUOUS_FUNCTIONAL_TYPE_KIND by error1<PsiElement, Collection<FunctionalTypeKind>>()
 
     // Context receivers resolution
     val NO_CONTEXT_RECEIVER by error1<KtElement, ConeKotlinType>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticRenderers.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.analysis.diagnostics
 
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.diagnostics.KtDiagnosticRenderers
 import org.jetbrains.kotlin.diagnostics.WhenMissingCase
@@ -142,6 +143,12 @@ object FirDiagnosticRenderers {
         }
     }
 
+    val FUNCTIONAL_TYPE_KIND = Renderer { kind: FunctionalTypeKind ->
+        kind.prefixForTypeRender ?: kind.classNamePrefix
+    }
+
+    val FUNCTIONAL_TYPE_KINDS = KtDiagnosticRenderers.COLLECTION(FUNCTIONAL_TYPE_KIND)
+
     @Suppress("FunctionName")
     fun <T> COLLECTION(renderer: ContextIndependentParameterRenderer<T>): ContextIndependentParameterRenderer<Collection<T>> {
         return Renderer { list ->

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrorsDefaultMessages.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.AMBI
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.DECLARATION_NAME
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.FIR
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.FQ_NAMES_IN_TYPES
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.FUNCTIONAL_TYPE_KINDS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.MODULE_DATA
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.NAME_OF_CONTAINING_DECLARATION_OR_FILE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticRenderers.RENDER_CLASS_OR_OBJECT
@@ -60,6 +61,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_ALTERED
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_ANONYMOUS_TYPE_INFERRED
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_EXPECTS
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_FUNCTIONAL_TYPE_KIND
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_SUPER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ANNOTATION_ARGUMENT_KCLASS_LITERAL_OF_TYPE_PARAMETER_ERROR
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ANNOTATION_ARGUMENT_MUST_BE_CONST
@@ -1076,6 +1078,7 @@ object FirErrorsDefaultMessages : BaseDiagnosticRendererFactory() {
 
         map.put(NEXT_MISSING, "Method ''next()'' cannot be called on ''iterator()''")
         map.put(NEXT_AMBIGUITY, "Method ''next()'' is ambiguous for this expression: {0}", SYMBOLS)
+        map.put(AMBIGUOUS_FUNCTIONAL_TYPE_KIND, "Multiple functional type conversions are not allowed for single type. Detected type conversions: {0}", FUNCTIONAL_TYPE_KINDS)
         map.put(NEXT_NONE_APPLICABLE, "None of the ''next()'' functions is applicable for ''iterator()'' of type ''{0}''", SYMBOLS)
 
         map.put(NO_CONTEXT_RECEIVER, "No required context receiver found: {0}", RENDER_TYPE)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/coneDiagnosticToFirDiagnostic.kt

@@ -132,6 +132,7 @@ private fun ConeDiagnostic.toKtDiagnostic(
     is ConeAmbiguouslyResolvedAnnotationFromPlugin -> {
         FirErrors.PLUGIN_ANNOTATION_AMBIGUITY.createOn(source, typeFromCompilerPhase, typeFromTypesPhase)
     }
+    is ConeAmbiguousFunctionalTypeKinds -> FirErrors.AMBIGUOUS_FUNCTIONAL_TYPE_KIND.createOn(source, kinds)
     else -> throw IllegalArgumentException("Unsupported diagnostic type: ${this.javaClass}")
 }
 

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

@@ -8,13 +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.toAnnotationClassId
 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
+import org.jetbrains.kotlin.name.ClassId
 
-class FirFunctionalTypeKindServiceImpl(session: FirSession) : FirFunctionalTypeKindService() {
+class FirFunctionalTypeKindServiceImpl(private val session: FirSession) : FirFunctionalTypeKindService() {
     private val nonReflectKindsFromExtensions = mutableListOf<FunctionalTypeKind>()
 
     override val extractor: FunctionalTypeKindExtractor = run {
@@ -55,12 +57,24 @@ class FirFunctionalTypeKindServiceImpl(session: FirSession) : FirFunctionalTypeK
     }
 
     override fun extractAllSpecialKindsForFunction(functionSymbol: FirNamedFunctionSymbol): List<FunctionalTypeKind> {
+        return extractSpecialKindsImpl(functionSymbol, { isSuspend }, { resolvedAnnotationClassIds })
+    }
+
+    override fun extractAllSpecialKindsForFunctionalTypeRef(typeRef: FirFunctionTypeRef): List<FunctionalTypeKind> {
+        return extractSpecialKindsImpl(typeRef, { isSuspend }, { annotations.mapNotNull { it.toAnnotationClassId(session) } })
+    }
+
+    private inline fun <T> extractSpecialKindsImpl(
+        source: T,
+        isSuspend: T.() -> Boolean,
+        annotations: T.() -> List<ClassId>
+    ): List<FunctionalTypeKind> {
         return buildList {
-            if (functionSymbol.isSuspend) {
+            if (source.isSuspend()) {
                 add(FunctionalTypeKind.SuspendFunction)
             }
             if (nonReflectKindsFromExtensions.isNotEmpty()) {
-                for (annotationClassId in functionSymbol.resolvedAnnotationClassIds) {
+                for (annotationClassId in source.annotations()) {
                     for (kind in nonReflectKindsFromExtensions) {
                         if (kind.annotationOnInvokeClassId == annotationClassId) {
                             add(kind)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirTypeResolverImpl.kt

@@ -6,15 +6,13 @@
 package org.jetbrains.kotlin.fir.resolve.providers.impl
 
 import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirOuterClassTypeParameterRef
 import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
-import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
+import org.jetbrains.kotlin.fir.diagnostics.*
-import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
-import org.jetbrains.kotlin.fir.diagnostics.ConeUnexpectedTypeArgumentsError
-import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.AbstractCallInfo
 import org.jetbrains.kotlin.fir.resolve.calls.AbstractCandidate
@@ -41,23 +39,6 @@ import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled
 
 @ThreadSafeMutableState
 class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
-
-    private val symbolProvider by lazy {
-        session.symbolProvider
-    }
-
-    private data class ClassIdInSession(val session: FirSession, val id: ClassId)
-
-    private val implicitBuiltinTypeSymbols = mutableMapOf<ClassIdInSession, FirClassLikeSymbol<*>>()
-
-    // TODO: get rid of session used here, and may be also of the cache above (see KT-30275)
-    private fun resolveBuiltInQualified(id: ClassId, session: FirSession): FirClassLikeSymbol<*> {
-        val nameInSession = ClassIdInSession(session, id)
-        return implicitBuiltinTypeSymbols.getOrPut(nameInSession) {
-            symbolProvider.getClassLikeSymbolByClassId(id)!!
-        }
-    }
-
     private fun resolveSymbol(
         symbol: FirBasedSymbol<*>,
         qualifier: List<FirQualifierPart>,
@@ -464,18 +445,25 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
         }
     }
 
-    private fun createFunctionalType(typeRef: FirFunctionTypeRef): ConeClassLikeType {
+    private fun createFunctionalType(typeRef: FirFunctionTypeRef): Pair<ConeClassLikeTypeImpl, ConeDiagnostic?> {
         val parameters =
             typeRef.contextReceiverTypeRefs.map { it.coneType } +
                     listOfNotNull(typeRef.receiverTypeRef?.coneType) +
                     typeRef.parameters.map { it.returnTypeRef.coneType.withParameterNameAnnotation(it, session) } +
                     listOf(typeRef.returnTypeRef.coneType)
-        val classId = if (typeRef.isSuspend) {
+        val functionalKinds = session.functionalTypeService.extractAllSpecialKindsForFunctionalTypeRef(typeRef)
-            StandardClassIds.SuspendFunctionN(typeRef.parametersCount)
+        var diagnostic: ConeDiagnostic? = null
-        } else {
+        val kind = when (functionalKinds.size) {
-            StandardClassIds.FunctionN(typeRef.parametersCount)
+            0 -> FunctionalTypeKind.Function
+            1 -> functionalKinds.single()
+            else -> {
+                diagnostic = ConeAmbiguousFunctionalTypeKinds(functionalKinds)
+                FunctionalTypeKind.Function
+            }
         }
 
+        val classId = kind.numberedClassId(typeRef.parametersCount)
+
         val attributes = typeRef.annotations.computeTypeAttributes(
             session,
             predefined = buildList {
@@ -488,17 +476,12 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
                 }
             }
         )
-        val symbol = resolveBuiltInQualified(classId, session)
         return ConeClassLikeTypeImpl(
-            symbol.toLookupTag().also {
+            classId.toLookupTag(),
-                if (it is ConeClassLikeLookupTagImpl) {
-                    it.bindSymbolToLookupTag(session, symbol)
-                }
-            },
             parameters.toTypedArray(),
             typeRef.isMarkedNullable,
             attributes
-        )
+        ) to diagnostic
     }
 
     override fun resolveType(
@@ -522,7 +505,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
                     isOperandOfIsOperator,
                 ) to (result as? TypeResolutionResult.Resolved)?.typeCandidate?.diagnostic
             }
-            is FirFunctionTypeRef -> createFunctionalType(typeRef) to null
+            is FirFunctionTypeRef -> createFunctionalType(typeRef)
             is FirDynamicTypeRef -> ConeDynamicType.create(session) to null
             is FirIntersectionTypeRef -> {
                 val leftType = typeRef.leftType.coneType

compiler/fir/tree/src/org/jetbrains/kotlin/fir/diagnostics/ConeSimpleDiagnostic.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.diagnostics
 
 import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKind
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.name.Name
@@ -51,6 +52,11 @@ object ConeDanglingModifierOnTopLevel : ConeDiagnostic {
         get() = "Top level declaration expected"
 }
 
+class ConeAmbiguousFunctionalTypeKinds(val kinds: List<FunctionalTypeKind>) : ConeDiagnostic {
+    override val reason: String
+        get() = "There are multiple functional kinds for functional type ref"
+}
+
 enum class DiagnosticKind {
     Syntax,
     ExpressionExpected,

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirFunctionalTypeKindService.kt

@@ -26,6 +26,8 @@ abstract class FirFunctionalTypeKindService : FirSessionComponent {
 
     abstract fun extractSingleSpecialKindForFunction(functionSymbol: FirNamedFunctionSymbol): FunctionalTypeKind?
     abstract fun extractAllSpecialKindsForFunction(functionSymbol: FirNamedFunctionSymbol): List<FunctionalTypeKind>
+
+    abstract fun extractAllSpecialKindsForFunctionalTypeRef(typeRef: FirFunctionTypeRef): List<FunctionalTypeKind>
 }
 
 val FirSession.functionalTypeService: FirFunctionalTypeKindService by FirSession.sessionComponentAccessor()