FIR: extend SAM conversion to subtype of functional type

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/CallAndReferenceGenerator.kt

@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.fir.resolve.calls.isFunctional
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.resolve.inference.isKMutableProperty
+import org.jetbrains.kotlin.fir.resolve.resolveFunctionTypeIfSamInterface
+import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
@@ -618,8 +620,12 @@ class CallAndReferenceGenerator(
         if (expectedType is ConeTypeParameterType || expectedType.isBuiltinFunctionalType(session)) {
             return false
         }
-        // On the other hand, the actual type should be a functional type.
-        return argument.isFunctional(session)
+        // On the other hand, the actual type should be either a functional type or a subtype of a class that has a contributed `invoke`.
+        val lookupTag = parameter.returnTypeRef.coneTypeSafe<ConeClassLikeType>()?.lookupTag
+        val firRegularClass = lookupTag?.toFirRegularClass(session)
+        // TODO: cache SAM interface to resolved function type, as [SamResolution] does?
+        val expectedFunctionType = firRegularClass?.resolveFunctionTypeIfSamInterface(session, scopeSession)
+        return argument.isFunctional(session, scopeSession, expectedFunctionType)
     }
 
     private fun IrExpression.applyAssigningArrayElementsToVarargInNamedForm(

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
@@ -32,6 +33,9 @@ fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): FirClassLikeSym
     }
 }
 
+fun ConeClassLikeLookupTag.toFirRegularClass(session: FirSession): FirRegularClass? =
+    session.firSymbolProvider.getSymbolByLookupTag(this)?.fir as? FirRegularClass
+
 @OptIn(LookupTagInternals::class)
 fun ConeClassLikeLookupTagImpl.bindSymbolToLookupTag(session: FirSession, symbol: FirClassLikeSymbol<*>?) {
     boundSymbol = OneElementWeakMap(session, symbol)

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

@@ -70,11 +70,7 @@ class FirSamResolverImpl(
     }
 
     private fun getFunctionTypeForPossibleSamType(type: ConeClassLikeType): ConeLookupTagBasedType? {
-        val firRegularClass =
-            firSession.firSymbolProvider
-                .getSymbolByLookupTag(type.lookupTag)
-                ?.fir as? FirRegularClass
-                ?: return null
+        val firRegularClass = type.lookupTag.toFirRegularClass(firSession) ?: return null
 
         val unsubstitutedFunctionType = resolveFunctionTypeIfSamInterface(firRegularClass) ?: return null
 
@@ -216,11 +212,7 @@ class FirSamResolverImpl(
 
     private fun resolveFunctionTypeIfSamInterface(firRegularClass: FirRegularClass): ConeLookupTagBasedType? {
         return resolvedFunctionType.getOrPut(firRegularClass) {
-            if (!firRegularClass.status.isFun) return@getOrPut NULL_STUB
-            val abstractMethod = firRegularClass.getSingleAbstractMethodOrNull(firSession, scopeSession) ?: return@getOrPut NULL_STUB
-            // TODO: val shouldConvertFirstParameterToDescriptor = samWithReceiverResolvers.any { it.shouldConvertFirstSamParameterToReceiver(abstractMethod) }
-
-            abstractMethod.getFunctionTypeForAbstractMethod()
+            firRegularClass.resolveFunctionTypeIfSamInterface(firSession, scopeSession) ?: NULL_STUB
         } as? ConeLookupTagBasedType
     }
 
@@ -230,6 +222,17 @@ class FirSamResolverImpl(
     }
 }
 
+fun FirRegularClass.resolveFunctionTypeIfSamInterface(
+    session: FirSession,
+    scopeSession: ScopeSession
+): ConeLookupTagBasedType? {
+    if (!this.status.isFun) return null
+    val abstractMethod = getSingleAbstractMethodOrNull(session, scopeSession) ?: return null
+    // TODO: val shouldConvertFirstParameterToDescriptor = samWithReceiverResolvers.any { it.shouldConvertFirstSamParameterToReceiver(abstractMethod) }
+
+    return abstractMethod.getFunctionTypeForAbstractMethod()
+}
+
 private fun FirRegularClass.getSingleAbstractMethodOrNull(
     session: FirSession,
     scopeSession: ScopeSession,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.addSubtypeConstraintIfCompatible
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.types.model.CaptureStatus
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 fun Candidate.resolveArgumentExpression(
     csBuilder: ConstraintSystemBuilder,
@@ -336,10 +337,9 @@ internal fun Candidate.resolveArgument(
     sink: CheckerSink,
     context: ResolutionContext
 ) {
-
     argument.resultType.ensureResolvedTypeDeclaration(context.session)
 
-    val expectedType = prepareExpectedType(context.session, argument, parameter, context)
+    val expectedType = prepareExpectedType(context.session, context.bodyResolveComponents.scopeSession, argument, parameter, context)
     resolveArgumentExpression(
         this.system.getBuilder(),
         argument,
@@ -354,18 +354,20 @@ internal fun Candidate.resolveArgument(
 
 private fun Candidate.prepareExpectedType(
     session: FirSession,
+    scopeSession: ScopeSession,
     argument: FirExpression,
     parameter: FirValueParameter?,
     context: ResolutionContext
 ): ConeKotlinType? {
     if (parameter == null) return null
     val basicExpectedType = argument.getExpectedTypeForSAMConversion(parameter/*, LanguageVersionSettings*/)
-    val expectedType = getExpectedTypeWithSAMConversion(session, argument, basicExpectedType, context) ?: basicExpectedType
+    val expectedType = getExpectedTypeWithSAMConversion(session, scopeSession, argument, basicExpectedType, context) ?: basicExpectedType
     return this.substitutor.substituteOrSelf(expectedType)
 }
 
 private fun Candidate.getExpectedTypeWithSAMConversion(
     session: FirSession,
+    scopeSession: ScopeSession,
     argument: FirExpression,
     candidateExpectedType: ConeKotlinType,
     context: ResolutionContext
@@ -375,20 +377,43 @@ private fun Candidate.getExpectedTypeWithSAMConversion(
     val firFunction = symbol.fir as? FirFunction<*> ?: return null
     if (!context.bodyResolveComponents.samResolver.shouldRunSamConversionForFunction(firFunction)) return null
 
-    if (!argument.isFunctional(session)) return null
-
     // TODO: resolvedCall.registerArgumentWithSamConversion(argument, SamConversionDescription(convertedTypeByOriginal, convertedTypeByCandidate!!))
 
-    return context.bodyResolveComponents.samResolver.getFunctionTypeForPossibleSamType(candidateExpectedType).apply {
-        usesSAM = true
+    val expectedFunctionType = context.bodyResolveComponents.samResolver.getFunctionTypeForPossibleSamType(candidateExpectedType)
+    return runIf(argument.isFunctional(session, scopeSession, expectedFunctionType)) {
+        expectedFunctionType.apply {
+            // Even though the `expectedFunctionalType` could be `null`, we should mark the flag to indicate that the argument is a
+            // functional type. That will help avoid ambiguous `invoke` resolutions. See KT-39824
+            usesSAM = true
+        }
     }
 }
 
-fun FirExpression.isFunctional(session: FirSession): Boolean =
+fun FirExpression.isFunctional(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    expectedFunctionType: ConeKotlinType?,
+): Boolean {
     when ((this as? FirWrappedArgumentExpression)?.expression ?: this) {
-        is FirAnonymousFunction, is FirCallableReferenceAccess -> true
-        else -> typeRef.coneTypeSafe<ConeKotlinType>()?.isBuiltinFunctionalType(session) == true
+        is FirAnonymousFunction, is FirCallableReferenceAccess -> return true
+        else -> {
+            // Either a functional type or a subtype of a class that has a contributed `invoke`.
+            val coneType = typeRef.coneTypeSafe<ConeKotlinType>() ?: return false
+            if (coneType.isBuiltinFunctionalType(session)) {
+                return true
+            }
+            val classLikeExpectedFunctionType = expectedFunctionType as? ConeClassLikeType
+            if (classLikeExpectedFunctionType == null || coneType is ConeIntegerLiteralType) {
+                return false
+            }
+            val invokeSymbol =
+                coneType.findContributedInvokeSymbol(
+                    session, scopeSession, classLikeExpectedFunctionType, shouldCalculateReturnTypesOfFakeOverrides = false
+                )
+            return invokeSymbol != null
+        }
     }
+}
 
 fun FirExpression.getExpectedTypeForSAMConversion(
     parameter: FirValueParameter/*, languageVersionSettings: LanguageVersionSettings*/