FIR: extend suspend conversion to subtype of functional type

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

@@ -14,12 +14,14 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.createFunctionalType
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredCallableSymbols
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -70,7 +72,7 @@ internal class AdapterGenerator(
         function: IrFunction
     ): Boolean =
         needSuspendConversion(type, function) || needCoercionToUnit(type, function) ||
-            needVarargSpread(callableReferenceAccess, type, function)
+                needVarargSpread(callableReferenceAccess, type, function)
 
     /**
      * For example,
@@ -136,7 +138,6 @@ internal class AdapterGenerator(
         return coneType.toIrType() as IrSimpleType
     }
 
-    // TODO: refactor to share some logic with suspend conversion on arguments
     internal fun generateAdaptedCallableReference(
         callableReferenceAccess: FirCallableReferenceAccess,
         explicitReceiverExpression: IrExpression?,
@@ -398,32 +399,35 @@ internal class AdapterGenerator(
         }
     }
 
-    // TODO: refactor to share some logic with suspend conversion for callable reference
+    /**
+     * For example,
+     * fun consumer(f: suspend () -> Unit) = ...
+     * fun nonSuspendFunction = { ... }
+     * fun useSite(...) = { ... consumer(nonSuspendFunction) ... }
+     *
+     * At the use site, instead of the argument, we can put the suspend lambda as an adapter.
+     *
+     * Instead of functions, a class with an overridden invoke can be used too:
+     * class Foo {
+     *   override fun invoke() = ...
+     * }
+     * fun useSite(...) = { ... consumer(Foo()) ... }
+     */
     internal fun IrExpression.applySuspendConversionIfNeeded(
         argument: FirExpression,
         parameter: FirValueParameter?
     ): IrExpression {
+        // TODO: should refer to LanguageVersionSettings.SuspendConversion
         if (this is IrBlock && origin == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE) {
             return this
         }
-        if (parameter == null || !needSuspendConversion(argument, parameter)) {
+        if (parameter?.returnTypeRef?.coneType?.isSuspendFunctionType(session) != true) {
             return this
         }
 
         val suspendConvertedType = parameter.returnTypeRef.toIrType() as IrSimpleType
         val returnType = suspendConvertedType.arguments.last().typeOrNull!!
-        val invokeSymbol =
-            (argument.typeRef.coneType as? ConeClassLikeType)?.lookupTag?.classId
-                ?.let { classId ->
-                    session.firSymbolProvider
-                        .getClassDeclaredCallableSymbols(classId, Name.identifier("invoke"))
-                        .filterIsInstance<FirFunctionSymbol<*>>()
-                        .find { firFunctionSymbol ->
-                            firFunctionSymbol.fir.valueParameters.size == suspendConvertedType.arguments.size - 1
-                        }?.let { firFunctionSymbol ->
-                            declarationStorage.getIrFunctionSymbol(firFunctionSymbol) as? IrSimpleFunctionSymbol
-                        }
-                } ?: return this
+        val invokeSymbol = findInvokeSymbol(suspendConvertedType, argument) ?: return this
         return argument.convertWithOffsets { startOffset, endOffset ->
             val irAdapterFunction = createAdapterFunctionForArgument(startOffset, endOffset, suspendConvertedType)
             // TODO: Should be able to reuse `this` if that is an immutable IrGetValue
@@ -448,19 +452,41 @@ internal class AdapterGenerator(
         }
     }
 
-    /**
-     * For example,
-     * fun consumer(f: suspend () -> Unit) = ...
-     * fun nonSuspendFunction = { ... }
-     * fun useSite(...) = { ... consumer(nonSuspendFunction) ... }
-     *
-     * At the use site, instead of the argument, we can put the suspend lambda as an adapter.
-     */
-    private fun needSuspendConversion(argument: FirExpression, parameter: FirValueParameter): Boolean =
-        // TODO: should refer to LanguageVersionSettings.SuspendConversion
-        parameter.returnTypeRef.coneType.isSuspendFunctionType(session) &&
-                argument.typeRef.coneType.isBuiltinFunctionalType(session) &&
-                !argument.typeRef.coneType.isSuspendFunctionType(session)
+    private fun findInvokeSymbol(type: IrSimpleType, argument: FirExpression): IrSimpleFunctionSymbol? {
+        val argumentType = argument.typeRef.coneType
+        // Either the argument type is a non-suspend functional type
+        if (argumentType.isBuiltinFunctionalType(session) && !argumentType.isSuspendFunctionType(session)) {
+            return (argument.typeRef.coneType as? ConeClassLikeType)?.lookupTag?.classId
+                ?.let { classId ->
+                    session.firSymbolProvider
+                        .getClassDeclaredCallableSymbols(classId, Name.identifier("invoke"))
+                        .filterIsInstance<FirFunctionSymbol<*>>()
+                        .find { firFunctionSymbol ->
+                            firFunctionSymbol.fir.valueParameters.size == type.arguments.size - 1
+                        }?.let { firFunctionSymbol ->
+                            declarationStorage.getIrFunctionSymbol(firFunctionSymbol) as? IrSimpleFunctionSymbol
+                        }
+                }
+        }
+
+        // Or is a class-like type whose corresponding class is a user-defined and has an overridden `invoke`.
+        if (argumentType is ConeClassLikeType) {
+            val klass =
+                (session.firSymbolProvider.getClassLikeSymbolByFqName(argumentType.lookupTag.classId) as? FirRegularClassSymbol)?.fir
+                    ?: return null
+            if (klass.origin != FirDeclarationOrigin.Source) {
+                return null
+            }
+            return klass.declarations
+                .filterIsInstance<FirSimpleFunction>()
+                .singleOrNull { it.name == Name.identifier("invoke") && it.valueParameters.size == type.arguments.size - 1 }
+                ?.let { invokeFunction ->
+                    declarationStorage.getIrFunctionSymbol(invokeFunction.symbol) as? IrSimpleFunctionSymbol
+                }
+        }
+
+        return null
+    }
 
     private fun createAdapterFunctionForArgument(
         startOffset: Int,

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

@@ -6,12 +6,12 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
-import org.jetbrains.kotlin.fir.declarations.FirFunction
-import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.createFunctionalType
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isLocalClassOrAnonymousObject
+import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.inference.*
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
@@ -19,8 +19,11 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.transformers.ensureResolvedTypeDeclaration
 import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.name.Name
 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
@@ -224,20 +227,12 @@ fun Candidate.resolvePlainArgumentType(
             capturedType
 
     // If the argument is of functional type and the expected type is a suspend function type, we need to do "suspend conversion."
-    // TODO: should refer to LanguageVersionSettings.SuspendConversion
-    if (expectedType?.isSuspendFunctionType(session) == true &&
-        argumentTypeForApplicabilityCheck.isBuiltinFunctionalType(session) &&
-        !argumentTypeForApplicabilityCheck.isSuspendFunctionType(session)
-    ) {
-        val typeParameters = argumentTypeForApplicabilityCheck.typeArguments.map { it as ConeKotlinType }
-        argumentTypeForApplicabilityCheck =
-            createFunctionalType(
-                typeParameters.dropLast(1), null, typeParameters.last(),
-                isSuspend = true,
-                isKFunctionType = argumentTypeForApplicabilityCheck.isKFunctionType(session)
-            )
-        substitutor.substituteOrSelf(argumentTypeForApplicabilityCheck)
-        usesSuspendConversion = true
+    if (expectedType != null) {
+        argumentTypeWithSuspendConversion(session, expectedType, argumentTypeForApplicabilityCheck)?.let {
+            argumentTypeForApplicabilityCheck = it
+            substitutor.substituteOrSelf(argumentTypeForApplicabilityCheck)
+            usesSuspendConversion = true
+        }
     }
 
     checkApplicabilityForArgumentType(
@@ -245,6 +240,50 @@ fun Candidate.resolvePlainArgumentType(
     )
 }
 
+private fun argumentTypeWithSuspendConversion(
+    session: FirSession, expectedType: ConeKotlinType, argumentType: ConeKotlinType
+): ConeKotlinType? {
+    // TODO: should refer to LanguageVersionSettings.SuspendConversion
+    // Expect the expected type to be a suspend functional type.
+    if (!expectedType.isSuspendFunctionType(session)) {
+        return null
+    }
+
+    // Either the argument type is a non-suspend functional type
+    if (argumentType.isBuiltinFunctionalType(session) && !argumentType.isSuspendFunctionType(session)) {
+        val typeParameters = argumentType.typeArguments.map { it as ConeKotlinType }
+        return createFunctionalType(
+            typeParameters.dropLast(1), null, typeParameters.last(),
+            isSuspend = true,
+            isKFunctionType = argumentType.isKFunctionType(session)
+        )
+    }
+
+    // Or is a class-like type whose corresponding class is a user-defined and has an overridden `invoke`.
+    if (argumentType is ConeClassLikeType) {
+        val klass =
+            (session.firSymbolProvider.getClassLikeSymbolByFqName(argumentType.lookupTag.classId) as? FirRegularClassSymbol)?.fir
+                ?: return null
+        if (klass.origin != FirDeclarationOrigin.Source) {
+            return null
+        }
+        return klass.declarations
+            .filterIsInstance<FirSimpleFunction>()
+            .singleOrNull { it.name == Name.identifier("invoke") && it.valueParameters.size == expectedType.typeArguments.size - 1 }
+            ?.let { invokeFunction ->
+                createFunctionalType(
+                    invokeFunction.valueParameters.map { it.returnTypeRef.coneType },
+                    null,
+                    invokeFunction.returnTypeRef.coneType,
+                    isSuspend = true,
+                    isKFunctionType = false
+                )
+            }
+    }
+
+    return null
+}
+
 fun Candidate.prepareCapturedType(argumentType: ConeKotlinType, context: ResolutionContext): ConeKotlinType {
     return captureTypeFromExpressionOrNull(argumentType, context) ?: argumentType
 }

compiler/testData/codegen/box/coroutines/suspendConversion/subtypeOfFunctionalTypeToSuspendConversion.kt

@@ -3,7 +3,6 @@
 // WITH_COROUTINES
 // IGNORE_BACKEND: JVM, NATIVE, JS, JS_IR
 // IGNORE_BACKEND: JS_IR_ES6
-// IGNORE_BACKEND_FIR: JVM_IR
 // IGNORE_LIGHT_ANALYSIS
 
 import helpers.*