Kapt: Support correctErrorTypes in suspend functions (KT-27711)

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/coroutineCodegenUtil.kt

@@ -122,6 +122,8 @@ val INITIAL_DESCRIPTOR_FOR_SUSPEND_FUNCTION = object : CallableDescriptor.UserDa
 @JvmField
 val INITIAL_SUSPEND_DESCRIPTOR_FOR_DO_RESUME = object : CallableDescriptor.UserDataKey<FunctionDescriptor> {}
 
+val CONTINUATION_PARAMETER_NAME = Name.identifier("continuation")
+
 // Resolved calls to suspension function contain descriptors as they visible within coroutines:
 // E.g. `fun <V> await(f: CompletableFuture<V>): V` instead of `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`
 // See `createJvmSuspendFunctionView` and it's usages for clarification
@@ -258,7 +260,7 @@ fun <D : FunctionDescriptor> getOrCreateJvmSuspendFunctionView(
         original = null,
         index = function.valueParameters.size,
         annotations = Annotations.EMPTY,
-        name = Name.identifier("continuation"),
+        name = CONTINUATION_PARAMETER_NAME,
         // Add j.l.Object to invoke(), because that is the type of parameters we have in FunctionN+1
         outType = if (function.containingDeclaration.safeAs<ClassDescriptor>()?.defaultType?.isBuiltinFunctionalType == true)
             function.builtIns.nullableAnyType

plugins/kapt3/kapt3-compiler/src/org/jetbrains/kotlin/kapt3/stubs/ClassFileToSourceStubConverter.kt

@@ -25,6 +25,9 @@ import com.sun.tools.javac.tree.TreeMaker
 import com.sun.tools.javac.tree.TreeScanner
 import kotlinx.kapt.KaptIgnored
 import org.jetbrains.kotlin.base.kapt3.KaptFlag
+import org.jetbrains.kotlin.codegen.coroutines.CONTINUATION_PARAMETER_NAME
+import org.jetbrains.kotlin.codegen.needsExperimentalCoroutinesWrapper
+import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
@@ -46,6 +49,7 @@ import org.jetbrains.kotlin.load.java.sources.JavaSourceElement
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.DefaultValueArgument
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedValueArgument
@@ -743,7 +747,8 @@ class ClassFileToSourceStubConverter(val kaptContext: KaptContextForStubGenerati
     private fun extractMethodSignatureTypes(
         descriptor: CallableDescriptor,
         exceptionTypes: JavacList<JCExpression>,
-        jcReturnType: JCExpression?, method: MethodNode,
+        jcReturnType: JCExpression?,
+        method: MethodNode,
         parameters: JavacList<JCVariableDecl>,
         valueParametersFromDescriptor: List<ValueParameterDescriptor>
     ): Pair<SignatureParser.MethodGenericSignature, JCExpression?> {
@@ -760,12 +765,23 @@ class ClassFileToSourceStubConverter(val kaptContext: KaptContextForStubGenerati
                                     },
                                     ifNonError = { lazyType() })
                 } else if (descriptor is FunctionDescriptor && valueParametersFromDescriptor.size == parameters.size) {
-                    getNonErrorType(valueParametersFromDescriptor[index].type, METHOD_PARAMETER_TYPE,
+                    val parameterDescriptor = valueParametersFromDescriptor[index]
-                                    ktTypeProvider = {
+                    val sourceElement = kaptContext.origins[method]?.element as? KtFunction
-                                        val sourceElement = kaptContext.origins[method]?.element as? KtFunction
+
-                                        sourceElement?.valueParameters?.getOrNull(index)?.typeReference
+                    getNonErrorType(
-                                    },
+                        parameterDescriptor.type, METHOD_PARAMETER_TYPE,
-                                    ifNonError = { lazyType() })
+                        ktTypeProvider = {
+                            if (sourceElement == null) return@getNonErrorType null
+
+                            if (sourceElement.hasDeclaredReturnType() && isContinuationParameter(parameterDescriptor)) {
+                                val continuationTypeFqName = getContinuationTypeFqName(descriptor)
+                                val functionReturnType = sourceElement.typeReference!!.text
+                                KtPsiFactory(kaptContext.project).createType("$continuationTypeFqName<$functionReturnType>")
+                            } else {
+                                sourceElement.valueParameters.getOrNull(index)?.typeReference
+                            }
+                        },
+                        ifNonError = { lazyType() })
                 } else {
                     lazyType()
                 }
@@ -788,7 +804,26 @@ class ClassFileToSourceStubConverter(val kaptContext: KaptContextForStubGenerati
         return Pair(genericSignature, returnType)
     }
 
-    private inline fun <T : JCExpression?> getNonErrorType(
+    private fun isContinuationParameter(descriptor: ValueParameterDescriptor): Boolean {
+        val containingCallable = descriptor.containingDeclaration
+
+        return containingCallable.valueParameters.lastOrNull() == descriptor
+            && descriptor.name == CONTINUATION_PARAMETER_NAME
+            && descriptor.source == SourceElement.NO_SOURCE
+            && descriptor.type.constructor.declarationDescriptor?.fqNameSafe == getContinuationTypeFqName(containingCallable)
+    }
+
+    private fun getContinuationTypeFqName(descriptor: CallableDescriptor): FqName {
+        val areCoroutinesReleased = !descriptor.needsExperimentalCoroutinesWrapper()
+                && kaptContext.generationState.languageVersionSettings.supportsFeature(LanguageFeature.ReleaseCoroutines)
+
+        return when (areCoroutinesReleased) {
+            true -> DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME_RELEASE
+            false -> DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME_EXPERIMENTAL
+        }
+    }
+
+    private fun <T : JCExpression?> getNonErrorType(
         type: KotlinType?,
         kind: ErrorTypeCorrector.TypeKind,
         ktTypeProvider: () -> KtTypeReference?,

plugins/kapt3/kapt3-compiler/testData/converter/suspendErrorTypes.kt

@@ -1,4 +1,5 @@
 // CORRECT_ERROR_TYPES
+// NO_VALIDATION
 // WITH_RUNTIME
 
 @file:Suppress("UNRESOLVED_REFERENCE")

plugins/kapt3/kapt3-compiler/testData/converter/suspendErrorTypes.txt

@@ -5,13 +5,13 @@ public final class Foo {
 
     @org.jetbrains.annotations.Nullable()
     public final java.lang.Object a(@org.jetbrains.annotations.NotNull()
-    kotlin.coroutines.Continuation<? super error.NonExistentClass> p0) {
+    kotlin.coroutines.Continuation<ABC> p0) {
         return null;
     }
 
     @org.jetbrains.annotations.Nullable()
     public final java.lang.Object b(@org.jetbrains.annotations.NotNull()
-    kotlin.coroutines.Continuation<? super kotlin.Result<? extends error.NonExistentClass>> p0) {
+    kotlin.coroutines.Continuation<Result<ABC>> p0) {
         return null;
     }