[K2] [KN-55977] Fix suspend type serialization

Merge-request: KT-MR-8328
Merged-by: Vladimir Sukharev <Vladimir.Sukharev@jetbrains.com>

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/types/KtFirFunctionalType.kt

@@ -49,7 +49,7 @@ internal class KtFirFunctionalType(
 
     override val nullability: KtTypeNullability get() = withValidityAssertion { coneType.nullability.asKtNullability() }
 
-    override val isSuspend: Boolean get() = withValidityAssertion { coneType.isSuspendFunctionType(builder.rootSession) }
+    override val isSuspend: Boolean get() = withValidityAssertion { coneType.isSuspendOrKSuspendFunctionType(builder.rootSession) }
     override val arity: Int
         get() = withValidityAssertion {
             if (coneType.isExtensionFunctionType) coneType.typeArguments.size - 2

compiler/fir/checkers/checkers.jvm/src/org/jetbrains/kotlin/fir/analysis/jvm/checkers/declaration/FirSynchronizedAnnotationChecker.kt

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.isAbstract
 import org.jetbrains.kotlin.fir.declarations.utils.isInline
 import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
 import org.jetbrains.kotlin.fir.types.coneType
-import org.jetbrains.kotlin.fir.types.isSuspendFunctionType
+import org.jetbrains.kotlin.fir.types.isSuspendOrKSuspendFunctionType
 import org.jetbrains.kotlin.name.JvmNames.SYNCHRONIZED_ANNOTATION_CLASS_ID
 
 object FirSynchronizedAnnotationChecker : FirFunctionChecker() {
@@ -33,7 +33,7 @@ object FirSynchronizedAnnotationChecker : FirFunctionChecker() {
             return
         }
         if (declaration.isSuspend ||
-            (declaration as? FirAnonymousFunction)?.typeRef?.coneType?.isSuspendFunctionType(session) == true
+            (declaration as? FirAnonymousFunction)?.typeRef?.coneType?.isSuspendOrKSuspendFunctionType(session) == true
         ) {
             reporter.reportOn(annotation.source, FirJvmErrors.SYNCHRONIZED_ON_SUSPEND, context)
             return

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirInlineDeclarationChecker.kt

@@ -355,7 +355,7 @@ abstract class FirInlineDeclarationChecker : FirFunctionChecker() {
         for (param in function.valueParameters) {
             val coneType = param.returnTypeRef.coneType
             val isFunctionalType = coneType.isFunctionalType(context.session)
-            val isSuspendFunctionalType = coneType.isSuspendFunctionType(context.session)
+            val isSuspendFunctionalType = coneType.isSuspendOrKSuspendFunctionType(context.session)
             val defaultValue = param.defaultValue
 
             if (!(isFunctionalType || isSuspendFunctionalType) && (param.isNoinline || param.isCrossinline)) {
@@ -438,7 +438,7 @@ abstract class FirInlineDeclarationChecker : FirFunctionChecker() {
             function.valueParameters.any { param ->
                 val type = param.returnTypeRef.coneType
                 !param.isNoinline && !type.isNullable
-                        && (type.isFunctionalType(session) || type.isSuspendFunctionType(session))
+                        && (type.isFunctionalType(session) || type.isSuspendOrKSuspendFunctionType(session))
             }
         if (hasInlinableParameters) return
         if (function.isInlineOnly(session)) return

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirSuspendCallChecker.kt

@@ -130,7 +130,7 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
     private fun findEnclosingSuspendFunction(context: CheckerContext): FirFunction? {
         return context.containingDeclarations.lastOrNull {
             when (it) {
-                is FirAnonymousFunction -> it.typeRef.coneType.isSuspendFunctionType(context.session)
+                is FirAnonymousFunction -> it.typeRef.coneType.isSuspendOrKSuspendFunctionType(context.session)
                 is FirSimpleFunction -> it.isSuspend
                 else -> false
             }
@@ -218,7 +218,7 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() {
         calledDeclarationSymbol: FirCallableSymbol<*>
     ): Triple<FirExpression?, FirExpression?, ConeKotlinType?> {
         if (this is FirImplicitInvokeCall &&
-            dispatchReceiver != FirNoReceiverExpression && dispatchReceiver.typeRef.coneType.isSuspendFunctionType(session)
+            dispatchReceiver != FirNoReceiverExpression && dispatchReceiver.typeRef.coneType.isSuspendOrKSuspendFunctionType(session)
         ) {
             val variableForInvoke = dispatchReceiver
             val variableForInvokeType = variableForInvoke.typeRef.coneType

compiler/fir/fir2ir/jvm-backend/src/org/jetbrains/kotlin/fir/backend/jvm/FirJvmTypeMapper.kt

@@ -266,7 +266,7 @@ class ConeTypeSystemCommonBackendContextForTypeMapping(
 
     override fun SimpleTypeMarker.isSuspendFunction(): Boolean {
         require(this is ConeSimpleKotlinType)
-        return isSuspendFunctionType(session)
+        return isSuspendOrKSuspendFunctionType(session)
     }
 
     override fun SimpleTypeMarker.isKClass(): Boolean {

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrDeclarationStorage.kt

@@ -550,7 +550,7 @@ class Fir2IrDeclarationStorage(
             ?: if (isLambda) SpecialNames.ANONYMOUS else Name.special("<no name provided>")
         val visibility = simpleFunction?.visibility ?: Visibilities.Local
         val isSuspend =
-            if (isLambda) ((function as FirAnonymousFunction).typeRef as? FirResolvedTypeRef)?.type?.isSuspendFunctionType(session) == true
+            if (isLambda) ((function as FirAnonymousFunction).typeRef as? FirResolvedTypeRef)?.type?.isSuspendOrKSuspendFunctionType(session) == true
             else simpleFunction?.isSuspend == true
         val created = function.convertWithOffsets { startOffset, endOffset ->
             val result = declareIrSimpleFunction(signature, simpleFunction?.containerSource) { symbol ->

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

@@ -511,7 +511,7 @@ internal class AdapterGenerator(
             return this
         }
         // Expect the expected type to be a suspend functional type.
-        if (!parameterType.isSuspendFunctionType(session)) {
+        if (!parameterType.isSuspendOrKSuspendFunctionType(session)) {
             return this
         }
         val expectedFunctionalType = parameterType.suspendFunctionTypeToFunctionType(session)

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -18385,6 +18385,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/fir/SuspendExtension.kt");
         }
 
+        @Test
+        @TestMetadata("SuspendFunctionReference.kt")
+        public void testSuspendFunctionReference() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/SuspendFunctionReference.kt");
+        }
+
         @Test
         @TestMetadata("unqualifiedEnum.kt")
         public void testUnqualifiedEnum() throws Exception {

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

@@ -492,7 +492,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
     override fun KotlinTypeMarker.isSuspendFunctionTypeOrSubtype(): Boolean {
         require(this is ConeKotlinType)
         return isTypeOrSubtypeOf {
-            (it.lowerBoundIfFlexible() as ConeKotlinType).isSuspendFunctionType(session)
+            (it.lowerBoundIfFlexible() as ConeKotlinType).isSuspendOrKSuspendFunctionType(session)
         }
     }
 

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

@@ -85,10 +85,14 @@ fun ConeKotlinType.isFunctionalOrSuspendFunctionalType(session: FirSession): Boo
 }
 
 // SuspendFunction, KSuspendFunction
-fun ConeKotlinType.isSuspendFunctionType(session: FirSession): Boolean {
+fun ConeKotlinType.isSuspendOrKSuspendFunctionType(session: FirSession): Boolean {
     return isFunctionalType(session) { it.isSuspendType }
 }
 
+fun ConeKotlinType.isSuspendFunctionType(session: FirSession): Boolean {
+    return isFunctionalType(session) { it == FunctionClassKind.SuspendFunction }
+}
+
 // KFunction, KSuspendFunction
 fun ConeKotlinType.isKFunctionType(session: FirSession): Boolean {
     return isFunctionalType(session) { it.isReflectType }
@@ -97,14 +101,14 @@ fun ConeKotlinType.isKFunctionType(session: FirSession): Boolean {
 fun ConeKotlinType.kFunctionTypeToFunctionType(session: FirSession): ConeClassLikeType {
     require(this.isKFunctionType(session))
     val kind =
-        if (isSuspendFunctionType(session)) FunctionClassKind.SuspendFunction
+        if (isSuspendOrKSuspendFunctionType(session)) FunctionClassKind.SuspendFunction
         else FunctionClassKind.Function
     val functionalTypeId = ClassId(kind.packageFqName, kind.numberedClassName(typeArguments.size - 1))
     return ConeClassLikeTypeImpl(ConeClassLikeLookupTagImpl(functionalTypeId), typeArguments, isNullable = false)
 }
 
 fun ConeKotlinType.suspendFunctionTypeToFunctionType(session: FirSession): ConeClassLikeType {
-    require(this.isSuspendFunctionType(session))
+    require(this.isSuspendOrKSuspendFunctionType(session))
     val kind =
         if (isKFunctionType(session)) FunctionClassKind.KFunction
         else FunctionClassKind.Function
@@ -113,7 +117,7 @@ fun ConeKotlinType.suspendFunctionTypeToFunctionType(session: FirSession): ConeC
 }
 
 fun ConeKotlinType.suspendFunctionTypeToFunctionTypeWithContinuation(session: FirSession, continuationClassId: ClassId): ConeClassLikeType {
-    require(this.isSuspendFunctionType(session))
+    require(this.isSuspendOrKSuspendFunctionType(session))
     val kind =
         if (isKFunctionType(session)) FunctionClassKind.KFunction
         else FunctionClassKind.Function
@@ -138,24 +142,24 @@ fun ConeKotlinType.isSubtypeOfFunctionalType(session: FirSession, expectedFuncti
 }
 
 fun ConeKotlinType.findSubtypeOfNonSuspendFunctionalType(session: FirSession, expectedFunctionalType: ConeClassLikeType): ConeKotlinType? {
-    require(expectedFunctionalType.isBuiltinFunctionalType(session) && !expectedFunctionalType.isSuspendFunctionType(session))
+    require(expectedFunctionalType.isBuiltinFunctionalType(session) && !expectedFunctionalType.isSuspendOrKSuspendFunctionType(session))
     return when (this) {
         is ConeClassLikeType -> {
             // Expect the argument type is not a suspend functional type.
-            if (isSuspendFunctionType(session) || !isSubtypeOfFunctionalType(session, expectedFunctionalType))
+            if (isSuspendOrKSuspendFunctionType(session) || !isSubtypeOfFunctionalType(session, expectedFunctionalType))
                 null
             else
                 this
         }
         is ConeIntersectionType -> {
-            if (intersectedTypes.any { it.isSuspendFunctionType(session) })
+            if (intersectedTypes.any { it.isSuspendOrKSuspendFunctionType(session) })
                 null
             else
                 intersectedTypes.find { it.findSubtypeOfNonSuspendFunctionalType(session, expectedFunctionalType) != null }
         }
         is ConeTypeParameterType -> {
             val bounds = lookupTag.typeParameterSymbol.resolvedBounds.map { it.coneType }
-            if (bounds.any { it.isSuspendFunctionType(session) })
+            if (bounds.any { it.isSuspendOrKSuspendFunctionType(session) })
                 null
             else
                 bounds.find { it.findSubtypeOfNonSuspendFunctionalType(session, expectedFunctionalType) != null }

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

@@ -69,7 +69,7 @@ fun FirAnonymousFunction.shouldReturnUnit(returnStatements: Collection<FirExpres
     isLambda && returnStatements.any { it is FirUnitExpression }
 
 fun FirAnonymousFunction.isExplicitlySuspend(session: FirSession): Boolean =
-    typeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendFunctionType(session) == true
+    typeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendOrKSuspendFunctionType(session) == true
 
 fun FirAnonymousFunction.addReturnToLastStatementIfNeeded() {
     // If this lambda's resolved, expected return type is Unit, we don't need an explicit return statement.

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

@@ -288,7 +288,7 @@ private fun argumentTypeWithSuspendConversion(
     // TODO: should refer to LanguageVersionSettings.SuspendConversion
 
     // Expect the expected type to be a suspend functional type.
-    if (!expectedType.isSuspendFunctionType(session)) {
+    if (!expectedType.isSuspendOrKSuspendFunctionType(session)) {
         return null
     }
 

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

@@ -253,7 +253,7 @@ private fun BodyResolveComponents.getCallableReferenceAdaptation(
     else
         mappedArguments
 
-    val suspendConversionStrategy = if ((function as? FirSimpleFunction)?.isSuspend != true && expectedType.isSuspendFunctionType(session))
+    val suspendConversionStrategy = if ((function as? FirSimpleFunction)?.isSuspend != true && expectedType.isSuspendOrKSuspendFunctionType(session))
         SuspendConversionStrategy.SUSPEND_CONVERSION
     else
         SuspendConversionStrategy.NO_CONVERSION

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

@@ -100,7 +100,7 @@ fun extractLambdaInfoFromFunctionalType(
     return ResolvedLambdaAtom(
         argument,
         expectedType,
-        expectedType.isSuspendFunctionType(session),
+        expectedType.isSuspendOrKSuspendFunctionType(session),
         receiverType,
         contextReceivers,
         parameters,

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

@@ -105,7 +105,7 @@ private fun extractLambdaInfo(
     session: FirSession,
     candidate: Candidate?
 ): ResolvedLambdaAtom {
-    val isSuspend = expectedType?.isSuspendFunctionType(session) ?: false
+    val isSuspend = expectedType?.isSuspendOrKSuspendFunctionType(session) ?: false
 
     val isFunctionSupertype =
         expectedType != null && expectedType.lowerBoundIfFlexible()

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

@@ -640,7 +640,7 @@ class FirCallCompletionResultsWriterTransformer(
         }
 
         if (needUpdateLambdaType) {
-            val isSuspend = expectedType?.isSuspendFunctionType(session) ?: result.isExplicitlySuspend(session)
+            val isSuspend = expectedType?.isSuspendOrKSuspendFunctionType(session) ?: result.isExplicitlySuspend(session)
             result.replaceTypeRef(result.constructFunctionalTypeRef(isSuspend))
             session.lookupTracker?.recordTypeResolveAsLookup(result.typeRef, result.source, context.file.source)
         }

compiler/testData/codegen/box/fir/SuspendFunctionReference.kt

@@ -0,0 +1,20 @@
+// TARGET_BACKEND: JVM
+// MODULE: lib
+// FILE: A.kt
+
+suspend fun foo(): String = "OK"
+fun fooref() = ::foo
+
+// MODULE: main(lib)
+// FILE: B.kt
+
+fun box(): String {
+    val expectedRefNameJVM = "function foo (Kotlin reflection is not available)"
+    val expectedRefNameNative = "suspend function foo"
+    val actualRefName = fooref().toString()
+
+    if (actualRefName == expectedRefNameJVM) return "OK"
+    if (actualRefName == expectedRefNameNative) return "OK"
+    return actualRefName
+}
+

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -17562,6 +17562,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
         public void testSuspendExtension() throws Exception {
             runTest("compiler/testData/codegen/box/fir/SuspendExtension.kt");
         }
+
+        @Test
+        @TestMetadata("SuspendFunctionReference.kt")
+        public void testSuspendFunctionReference() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/SuspendFunctionReference.kt");
+        }
     }
 
     @Nested

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -18385,6 +18385,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/fir/SuspendExtension.kt");
         }
 
+        @Test
+        @TestMetadata("SuspendFunctionReference.kt")
+        public void testSuspendFunctionReference() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/SuspendFunctionReference.kt");
+        }
+
         @Test
         @TestMetadata("unqualifiedEnum.kt")
         public void testUnqualifiedEnum() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -14552,6 +14552,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
         public void testNotFoundClasses() throws Exception {
             runTest("compiler/testData/codegen/box/fir/notFoundClasses.kt");
         }
+
+        @TestMetadata("SuspendFunctionReference.kt")
+        public void testSuspendFunctionReference() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/SuspendFunctionReference.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/fullJdk")