[FIR] Fix suspend functional type deserialization

compiler/fir/analysis-tests/testData/loadCompiledKotlin/coroutines/Basic.txt

@@ -1,4 +1,4 @@
-public final fun builder(c: @R|kotlin/ExtensionFunctionType|() R|ERROR CLASS: createSuspendFunctionType not supported|): R|kotlin/Unit|
+public final fun builder(c: R|suspend test/Controller.() -> kotlin/Unit|): R|kotlin/Unit|
 
 public final class Controller : R|kotlin/Any| {
     public final suspend fun suspendFun(): R|kotlin/Unit|
@@ -6,3 +6,4 @@ public final class Controller : R|kotlin/Any| {
     public constructor(): R|test/Controller|
 
 }
+

compiler/fir/analysis-tests/testData/loadCompiledKotlin/type/SuspendFunction.txt

@@ -1,8 +1,8 @@
-public final fun test1(): R|ERROR CLASS: createSuspendFunctionType not supported|
+public final fun test1(): R|suspend () -> kotlin/Unit|
 
-public final fun test2(): @R|kotlin/ExtensionFunctionType|() R|ERROR CLASS: createSuspendFunctionType not supported|
+public final fun test2(): R|suspend kotlin/Int.() -> kotlin/Int|
 
-public final fun test3(): R|kotlin/collections/List<ERROR CLASS: createSuspendFunctionType not supported>|
+public final fun test3(): R|kotlin/collections/List<kotlin/coroutines/SuspendFunction0<kotlin/Unit>>|
 
-public final fun test4(): R|ERROR CLASS: createSuspendFunctionType not supported|
+public final fun test4(): R|suspend () -> kotlin/coroutines/SuspendFunction0<kotlin/Unit>|
 

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/receiverResolutionInLambda.kt

@@ -10,9 +10,9 @@ class Case1() {
 
             val  y = this
             //this is Case1 instead of SequenceScope<String>
-            <!UNRESOLVED_REFERENCE!>yield<!>("") // UNRESOLVED_REFERENCE
+            yield("") // UNRESOLVED_REFERENCE
 
-            this.<!UNRESOLVED_REFERENCE!>yield<!>("") //UNRESOLVED_REFERENCE
+            this.yield("") //UNRESOLVED_REFERENCE
 
             this as SequenceScope<String>
 
@@ -27,14 +27,14 @@ fun case2() {
     val x = sequence<String> {
 
         val  y = this
-        <!UNRESOLVED_REFERENCE!>yield<!>("") // UNRESOLVED_REFERENCE
+        yield("") // UNRESOLVED_REFERENCE
 
-        this.<!UNRESOLVED_REFERENCE!>yield<!>("") //UNRESOLVED_REFERENCE
+        this.yield("") //UNRESOLVED_REFERENCE
 
         this as SequenceScope<String>
 
-        <!UNRESOLVED_REFERENCE!>yield<!>("") // UNRESOLVED_REFERENCE
+        yield("") // UNRESOLVED_REFERENCE
 
-        this.<!UNRESOLVED_REFERENCE!>yield<!>("") // UNRESOLVED_REFERENCE
+        this.yield("") // UNRESOLVED_REFERENCE
     }
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/receiverResolutionInLambda.txt

@@ -5,26 +5,26 @@ FILE: receiverResolutionInLambda.kt
         }
 
         public final fun foo(): R|kotlin/Unit| {
-            lval x: R|kotlin/sequences/Sequence<kotlin/String>| = R|kotlin/sequences/sequence|<R|kotlin/String|>(<L> = sequence@fun <anonymous>(): R|kotlin/Unit| {
+            lval x: R|kotlin/sequences/Sequence<kotlin/String>| = R|kotlin/sequences/sequence|<R|kotlin/String|>(<L> = sequence@fun R|kotlin/sequences/SequenceScope<kotlin/String>|.<anonymous>(): R|kotlin/Unit| {
-                lval y: R|Case1| = this@R|/Case1|
+                lval y: R|kotlin/sequences/SequenceScope<kotlin/String>| = this@R|special/anonymous|
-                <Unresolved name: yield>#(String())
+                this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
-                this@R|/Case1|.<Unresolved name: yield>#(String())
+                this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
-                (this@R|/Case1| as R|kotlin/sequences/SequenceScope<kotlin/String>|)
+                (this@R|special/anonymous| as R|kotlin/sequences/SequenceScope<kotlin/String>|)
-                this@R|/Case1|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
+                this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
-                this@R|/Case1|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
+                this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
             }
             )
         }
 
     }
     public final fun case2(): R|kotlin/Unit| {
-        lval x: R|kotlin/sequences/Sequence<kotlin/String>| = R|kotlin/sequences/sequence|<R|kotlin/String|>(<L> = sequence@fun <anonymous>(): <ERROR TYPE REF: Unresolved name: yield> {
+        lval x: R|kotlin/sequences/Sequence<kotlin/String>| = R|kotlin/sequences/sequence|<R|kotlin/String|>(<L> = sequence@fun R|kotlin/sequences/SequenceScope<kotlin/String>|.<anonymous>(): R|kotlin/Unit| {
-            lval y: R|ERROR CLASS: Unresolved this@null| = this#
+            lval y: R|kotlin/sequences/SequenceScope<kotlin/String>| = this@R|special/anonymous|
-            <Unresolved name: yield>#(String())
+            this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
-            this#.<Unresolved name: yield>#(String())
+            this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
-            (this# as R|kotlin/sequences/SequenceScope<kotlin/String>|)
+            (this@R|special/anonymous| as R|kotlin/sequences/SequenceScope<kotlin/String>|)
-            <Unresolved name: yield>#(String())
+            this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
-            ^ this#.<Unresolved name: yield>#(String())
+            this@R|special/anonymous|.R|FakeOverride<kotlin/sequences/SequenceScope.yield: R|kotlin/Unit|>|(String())
         }
         )
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/deserialization/FirTypeDeserializer.kt

@@ -5,10 +5,13 @@
 
 package org.jetbrains.kotlin.fir.deserialization
 
+import org.jetbrains.kotlin.builtins.functions.FunctionClassDescriptor
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirTypeParametersOwner
 import org.jetbrains.kotlin.fir.declarations.addDefaultBoundIfNecessary
 import org.jetbrains.kotlin.fir.declarations.builder.FirTypeParameterBuilder
+import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterLookupTag
@@ -21,6 +24,7 @@ import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.*
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.serialization.deserialization.ProtoEnumFlags
 import org.jetbrains.kotlin.serialization.deserialization.getClassId
 import org.jetbrains.kotlin.serialization.deserialization.getName
@@ -128,8 +132,7 @@ class FirTypeDeserializer(
         val arguments = proto.collectAllArguments().map(this::typeArgument).toTypedArray()
 
         val simpleType = if (Flags.SUSPEND_TYPE.get(proto.flags)) {
-            //createSuspendFunctionType(annotations, constructor, arguments, proto.nullable)
+            createSuspendFunctionType(constructor, arguments, isNullable = proto.nullable)
-            ConeClassErrorType("createSuspendFunctionType not supported")
         } else {
             ConeClassLikeTypeImpl(constructor, arguments, isNullable = proto.nullable)
         }
@@ -140,6 +143,63 @@ class FirTypeDeserializer(
 
     }
 
+    private fun createSuspendFunctionTypeForBasicCase(
+        //annotations: Annotations, TODO?,
+        functionTypeConstructor: ConeClassLikeLookupTag,
+        arguments: Array<ConeTypeProjection>,
+        isNullable: Boolean
+    ): ConeClassLikeType? {
+        fun ConeClassLikeType.isContinuation(): Boolean {
+            if (this.typeArguments.size != 1) return false
+            if (this.lookupTag.classId != CONTINUATION_INTERFACE_CLASS_ID) return false
+            return true
+        }
+
+        val returnType = arguments.lastOrNull()
+        val continuationType = arguments.getOrNull(arguments.lastIndex - 1) as? ConeClassLikeType ?: return null
+
+        if (!continuationType.isContinuation()) return ConeClassLikeTypeImpl(functionTypeConstructor, arguments, isNullable)
+
+        val suspendReturnType = continuationType.typeArguments.single() as ConeKotlinTypeProjection
+
+        val valueParameters = arguments.dropLast(2)
+
+        val kind = FunctionClassDescriptor.Kind.SuspendFunction
+        return ConeClassLikeTypeImpl(
+            ConeClassLikeLookupTagImpl(ClassId(kind.packageFqName, kind.numberedClassName(valueParameters.size))),
+            (valueParameters + suspendReturnType).toTypedArray(),
+            isNullable
+        )
+    }
+
+    private fun createSuspendFunctionType(
+        //annotations: Annotations, TODO?
+        functionTypeConstructor: ConeClassLikeLookupTag,
+        arguments: Array<ConeTypeProjection>,
+        isNullable: Boolean
+    ): ConeClassLikeType {
+        val result =
+            when (functionTypeConstructor.toSymbol(session)!!.firUnsafe<FirTypeParametersOwner>().typeParameters.size - arguments.size) {
+                0 -> createSuspendFunctionTypeForBasicCase(/* annotations, */ functionTypeConstructor, arguments, isNullable)
+//                 This case for types written by eap compiler 1.1
+                1 -> {
+                    val arity = arguments.size - 1
+                    if (arity >= 0) {
+                        val kind = FunctionClassDescriptor.Kind.SuspendFunction
+                        ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(ClassId(kind.packageFqName, kind.numberedClassName(arity))),
+                            arguments,
+                            isNullable
+                        )
+                    } else {
+                        null
+                    }
+                }
+                else -> null
+            }
+        return result ?: ConeClassErrorType("Bad suspend function in metadata with constructor: $functionTypeConstructor")
+    }
+
     private fun typeSymbol(proto: ProtoBuf.Type): ConeClassifierLookupTag? {
 
         return when {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/deserialization/suspendTypeUtils.kt

@@ -0,0 +1,14 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.deserialization
+
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.resolve.DescriptorUtils
+
+
+
+internal val CONTINUATION_INTERFACE_CLASS_ID = ClassId.topLevel(DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME_RELEASE)

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.functions.FunctionClassDescriptor
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
@@ -34,6 +35,7 @@ import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -219,11 +221,17 @@ fun createFunctionalType(
     receiverType: ConeKotlinType?,
     rawReturnType: ConeKotlinType,
     isKFunctionType: Boolean = false,
+    isSuspend: Boolean = false
 ): ConeLookupTagBasedType {
     val receiverAndParameterTypes = listOfNotNull(receiverType) + parameters + listOf(rawReturnType)
 
-    val postfix = "Function${receiverAndParameterTypes.size - 1}"
+    val kind = if (isSuspend) {
-    val functionalTypeId = if (isKFunctionType) StandardClassIds.reflectByName("K$postfix") else StandardClassIds.byName(postfix)
+        if (isKFunctionType) FunctionClassDescriptor.Kind.KSuspendFunction else FunctionClassDescriptor.Kind.SuspendFunction
+    } else {
+        if (isKFunctionType) FunctionClassDescriptor.Kind.KFunction else FunctionClassDescriptor.Kind.Function
+    }
+
+    val functionalTypeId = ClassId(kind.packageFqName, kind.numberedClassName(receiverAndParameterTypes.size - 1))
     return ConeClassLikeTypeImpl(ConeClassLikeLookupTagImpl(functionalTypeId), receiverAndParameterTypes.toTypedArray(), isNullable = false)
 }
 

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.inference
 
+import org.jetbrains.kotlin.builtins.functions.FunctionClassDescriptor
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
@@ -15,7 +16,12 @@ import org.jetbrains.kotlin.fir.types.*
 
 fun ConeKotlinType.isBuiltinFunctionalType(session: FirSession): Boolean {
     return when (this) {
-        is ConeClassLikeType -> fullyExpandedType(session).lookupTag.classId.asString().startsWith("kotlin/Function")
+        is ConeClassLikeType -> {
+            val classId = fullyExpandedType(session).lookupTag.classId
+            val kind =
+                FunctionClassDescriptor.Kind.byClassNamePrefix(classId.packageFqName, classId.relativeClassName.asString()) ?: return false
+            kind == FunctionClassDescriptor.Kind.Function || kind == FunctionClassDescriptor.Kind.SuspendFunction
+        }
         else -> false
     }
 }
@@ -24,7 +30,9 @@ fun ConeKotlinType.isSuspendFunctionType(session: FirSession): Boolean {
     return when (val type = this) {
         is ConeClassLikeType -> {
             val classId = type.fullyExpandedType(session).lookupTag.classId
-            classId.packageFqName.asString() == "kotlin" && classId.relativeClassName.asString().startsWith("SuspendFunction")
+            val kind =
+                FunctionClassDescriptor.Kind.byClassNamePrefix(classId.packageFqName, classId.relativeClassName.asString()) ?: return false
+            kind == FunctionClassDescriptor.Kind.SuspendFunction
         }
         else -> false
     }

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

@@ -33,7 +33,12 @@ fun Candidate.preprocessLambdaArgument(
 
     if (expectedType != null) {
         // TODO: add SAM conversion processing
-        val lambdaType = createFunctionalType(resolvedArgument.parameters, resolvedArgument.receiver, resolvedArgument.returnType)
+        val lambdaType = createFunctionalType(
+            resolvedArgument.parameters,
+            resolvedArgument.receiver,
+            resolvedArgument.returnType,
+            isSuspend = resolvedArgument.isSuspend
+        )
         csBuilder.addSubtypeConstraint(lambdaType, expectedType, SimpleConstraintSystemConstraintPosition)
     }
 

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

@@ -108,7 +108,7 @@ private fun extractInputOutputTypesFromCallableReferenceExpectedType(
     if (expectedType == null) return null
 
     return when {
-        expectedType.isBuiltinFunctionalType(session) || expectedType.isSuspendFunctionType(session) ->
+        expectedType.isBuiltinFunctionalType(session) ->
             extractInputOutputTypesFromFunctionType(expectedType, session)
 
 //        ReflectionTypes.isBaseTypeForNumberedReferenceTypes(expectedType) ->