[FIR] Extract functional type kind from annotation on lambda literal

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

@@ -13,7 +13,8 @@ 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.fir.symbols.impl.FirAnonymousFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.name.ClassId
 
 class FirFunctionalTypeKindServiceImpl(private val session: FirSession) : FirFunctionalTypeKindService() {
@@ -49,15 +50,25 @@ class FirFunctionalTypeKindServiceImpl(private val session: FirSession) : FirFun
         FunctionalTypeKindExtractor(kinds)
     }
 
-    override fun extractSingleSpecialKindForFunction(functionSymbol: FirNamedFunctionSymbol): FunctionalTypeKind? {
+    override fun extractSingleSpecialKindForFunction(functionSymbol: FirFunctionSymbol<*>): FunctionalTypeKind? {
         if (nonReflectKindsFromExtensions.isEmpty()) {
             return FunctionalTypeKind.SuspendFunction.takeIf { functionSymbol.isSuspend }
         }
+
         return extractAllSpecialKindsForFunction(functionSymbol).singleOrNull()
     }
 
-    override fun extractAllSpecialKindsForFunction(functionSymbol: FirNamedFunctionSymbol): List<FunctionalTypeKind> {
-        return extractSpecialKindsImpl(functionSymbol, { isSuspend }, { resolvedAnnotationClassIds })
+    override fun extractAllSpecialKindsForFunction(functionSymbol: FirFunctionSymbol<*>): List<FunctionalTypeKind> {
+        return extractSpecialKindsImpl(
+            functionSymbol,
+            { isSuspend },
+            {
+                when (functionSymbol) {
+                    is FirAnonymousFunctionSymbol -> functionSymbol.annotations.mapNotNull { it.toAnnotationClassId(session) }
+                    else -> resolvedAnnotationClassIds
+                }
+            }
+        )
     }
 
     override fun extractAllSpecialKindsForFunctionalTypeRef(typeRef: FirFunctionTypeRef): List<FunctionalTypeKind> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -18,10 +18,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.canNarrowDownGetterType
 import org.jetbrains.kotlin.fir.declarations.utils.expandedConeType
 import org.jetbrains.kotlin.fir.declarations.utils.isFinal
 import org.jetbrains.kotlin.fir.declarations.utils.modality
-import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
-import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
-import org.jetbrains.kotlin.fir.diagnostics.ConeStubDiagnostic
-import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
+import org.jetbrains.kotlin.fir.diagnostics.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
@@ -119,10 +116,30 @@ fun FirFunction.constructFunctionalType(kind: FunctionalTypeKind? = null): ConeL
 /**
  * [kind] == null means that [FunctionalTypeKind.Function] will be used
  */
-fun FirFunction.constructFunctionalTypeRef(kind: FunctionalTypeKind? = null): FirResolvedTypeRef {
-    return buildResolvedTypeRef {
-        source = this@constructFunctionalTypeRef.source?.fakeElement(KtFakeSourceElementKind.ImplicitTypeRef)
-        type = constructFunctionalType(kind)
+fun FirAnonymousFunction.constructFunctionalTypeRef(session: FirSession, kind: FunctionalTypeKind? = null): FirResolvedTypeRef {
+    var diagnostic: ConeDiagnostic? = null
+    val kinds = session.functionalTypeService.extractAllSpecialKindsForFunction(symbol)
+    val kindFromDeclaration = when(kinds.size) {
+        0 -> null
+        1 -> kinds.single()
+        else -> {
+            diagnostic = ConeAmbiguousFunctionalTypeKinds(kinds)
+            FunctionalTypeKind.Function
+        }
+    }
+    val type = constructFunctionalType(kindFromDeclaration ?: kind)
+    val source = this@constructFunctionalTypeRef.source?.fakeElement(KtFakeSourceElementKind.ImplicitTypeRef)
+    return if (diagnostic == null) {
+        buildResolvedTypeRef {
+            this.source = source
+            this.type = type
+        }
+    } else {
+        buildErrorTypeRef {
+            this.source = source
+            this.type = type
+            this.diagnostic = diagnostic
+        }
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt

@@ -37,6 +37,7 @@ fun extractLambdaInfoFromFunctionalType(
         )
     }
     val expectedFunctionalKind = expectedType.functionalTypeKind(session) ?: return null
+    val actualFunctionalKind = session.functionalTypeService.extractSingleSpecialKindForFunction(argument.symbol)
 
     val singleStatement = argument.body?.statements?.singleOrNull() as? FirReturnExpression
     if (argument.returnType == null && singleStatement != null &&
@@ -100,7 +101,7 @@ fun extractLambdaInfoFromFunctionalType(
     return ResolvedLambdaAtom(
         argument,
         expectedType,
-        expectedFunctionalKind,
+        actualFunctionalKind ?: expectedFunctionalKind,
         receiverType,
         contextReceivers,
         parameters,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -619,7 +619,7 @@ class FirCallCompletionResultsWriterTransformer(
         if (needUpdateLambdaType) {
             val kind = expectedType?.functionalTypeKind(session)
                 ?: result.typeRef.coneTypeSafe<ConeClassLikeType>()?.functionalTypeKind(session)
-            result.replaceTypeRef(result.constructFunctionalTypeRef(kind))
+            result.replaceTypeRef(result.constructFunctionalTypeRef(session, kind))
             session.lookupTracker?.recordTypeResolveAsLookup(result.typeRef, result.source, context.file.source)
         }
         // Have to delay this until the type is written to avoid adding a return if the type is Unit.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -714,6 +714,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
         data: ResolutionMode
     ): FirStatement = whileAnalysing(session, anonymousFunction) {
         // Either ContextDependent, ContextIndependent or WithExpectedType could be here
+        anonymousFunction.transformAnnotations(transformer, ResolutionMode.ContextIndependent)
         if (data !is ResolutionMode.LambdaResolution) {
             anonymousFunction.transformReturnTypeRef(transformer, ResolutionMode.ContextIndependent)
             anonymousFunction.transformReceiverParameter(transformer, ResolutionMode.ContextIndependent)
@@ -805,7 +806,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
             session.lookupTracker?.recordTypeResolveAsLookup(lambda.returnTypeRef, lambda.source, context.file.source)
         }
 
-        lambda.replaceTypeRef(lambda.constructFunctionalTypeRef(resolvedLambdaAtom?.expectedFunctionalTypeKind))
+        lambda.replaceTypeRef(lambda.constructFunctionalTypeRef(session, resolvedLambdaAtom?.expectedFunctionalTypeKind))
         session.lookupTracker?.recordTypeResolveAsLookup(lambda.typeRef, lambda.source, context.file.source)
         lambda.addReturnToLastStatementIfNeeded()
         return lambda

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

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.builtins.functions.FunctionalTypeKindExtractor
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.name.FqName
 
 abstract class FirFunctionalTypeKindService : FirSessionComponent {
@@ -24,9 +23,8 @@ abstract class FirFunctionalTypeKindService : FirSessionComponent {
         return extractor.hasKindWithSpecificPackage(packageFqName)
     }
 
-    abstract fun extractSingleSpecialKindForFunction(functionSymbol: FirNamedFunctionSymbol): FunctionalTypeKind?
-    abstract fun extractAllSpecialKindsForFunction(functionSymbol: FirNamedFunctionSymbol): List<FunctionalTypeKind>
-
+    abstract fun extractSingleSpecialKindForFunction(functionSymbol: FirFunctionSymbol<*>): FunctionalTypeKind?
+    abstract fun extractAllSpecialKindsForFunction(functionSymbol: FirFunctionSymbol<*>): List<FunctionalTypeKind>
     abstract fun extractAllSpecialKindsForFunctionalTypeRef(typeRef: FirFunctionTypeRef): List<FunctionalTypeKind>
 }