From d005dc685f684e4ff657000dfcba034ee8bc1566 Mon Sep 17 00:00:00 2001
From: Igor Yakovlev <igor.yakovlev@jetbrains.com>
Date: Thu, 6 Jun 2019 02:51:52 +0300
Subject: [PATCH] Add nullability support to
 KtUltraLightSuspendContinuationParameter

---
 .../asJava/classes/ultraLightParameters.kt    | 29 +++++++++++--------
 .../kotlin/asJava/classes/ultraLightUtils.kt  |  4 ++-
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightParameters.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightParameters.kt
index ecf69ac8c36..3bd7ed0e924 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightParameters.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightParameters.kt
@@ -14,31 +14,36 @@ import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.asJava.elements.KtLightSimpleModifierList
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.builtins.isSuspendFunctionType
-import org.jetbrains.kotlin.builtins.isSuspendFunctionTypeOrSubtype
-import org.jetbrains.kotlin.psi.KtCallableDeclaration
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_FUNCTION_CONTINUATION_PARAMETER
-import org.jetbrains.kotlin.codegen.kotlinType
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
-import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.*
 
 internal class KtUltraLightSuspendContinuationParameter(
     private val ktFunction: KtFunction,
     support: KtUltraLightSupport,
     method: KtLightMethod
-) : org.jetbrains.kotlin.asJava.elements.LightParameter(SUSPEND_FUNCTION_CONTINUATION_PARAMETER, PsiType.NULL, method, method.language) {
+) : org.jetbrains.kotlin.asJava.elements.LightParameter(SUSPEND_FUNCTION_CONTINUATION_PARAMETER, PsiType.NULL, method, method.language),
+    KtUltraLightElementWithNullabilityAnnotation<KtDeclaration, PsiParameter> {
 
-    private val psiType by lazyPub {
+    override val kotlinTypeForNullabilityAnnotation: KotlinType? get() = ktType
+    override val psiTypeForNullabilityAnnotation: PsiType? get() = psiType
+    override val kotlinOrigin: KtDeclaration? = null
+    override val clsDelegate: PsiParameter
+        get() = throw IllegalStateException("Cls delegate shouldn't be loaded for ultra-light PSI!")
+
+    private val ktType by lazyPub {
         val descriptor = ktFunction.resolve() as? FunctionDescriptor
-        val returnType = descriptor?.returnType ?: return@lazyPub PsiType.NULL
-        val typeFromDescriptor = support.moduleDescriptor.getContinuationOfTypeOrAny(returnType, support.isReleasedCoroutine)
-        typeFromDescriptor.asPsiType(support, TypeMappingMode.DEFAULT, method)
+        val returnType = descriptor?.returnType ?: return@lazyPub null
+        support.moduleDescriptor.getContinuationOfTypeOrAny(returnType, support.isReleasedCoroutine)
     }
 
-    private val lightModifierList by lazyPub { KtLightSimpleModifierList(method, emptySet()) }
+    private val psiType by lazyPub {
+        ktType?.asPsiType(support, TypeMappingMode.DEFAULT, method) ?: PsiType.NULL
+    }
+
+    private val lightModifierList by lazyPub { KtLightSimpleModifierList(this, emptySet()) }
 
     override fun getType(): PsiType = psiType
 
diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightUtils.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightUtils.kt
index 7e601485379..375fc9a501d 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightUtils.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightUtils.kt
@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.annotations.JVM_STATIC_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.annotations.argumentValue
 import org.jetbrains.kotlin.resolve.constants.EnumValue
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeProjectionImpl
 import org.jetbrains.kotlin.types.replace
@@ -129,7 +130,8 @@ internal fun KtDeclaration.getKotlinType(): KotlinType? {
     val descriptor = resolve()
     return when (descriptor) {
         is ValueDescriptor -> descriptor.type
-        is CallableDescriptor -> descriptor.returnType
+        is CallableDescriptor -> if (descriptor is FunctionDescriptor && descriptor.isSuspend)
+            descriptor.module.builtIns.nullableAnyType else descriptor.returnType
         else -> null
     }
 }