From 2ce391d3171c5fb6b57b3b073bf1a00a10a28b7e Mon Sep 17 00:00:00 2001
From: Igor Yakovlev <igor.yakovlev@jetbrains.com>
Date: Wed, 18 Aug 2021 17:04:51 +0200
Subject: [PATCH] [FIR IDE] Fix unensured access to the return type

---
 .../fir/components/KtFirExpressionTypeProvider.kt    |  2 +-
 .../kotlin/idea/frontend/api/fir/utils/firUtils.kt   | 12 +++++++++---
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt
index 3bf85967177..bf6ce7c08b8 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt
@@ -41,7 +41,7 @@ internal class KtFirExpressionTypeProvider(
     override fun getKtExpressionType(expression: KtExpression): KtType? = withValidityAssertion {
         when (val fir = expression.unwrap().getOrBuildFir(firResolveState)) {
             is FirExpression -> fir.typeRef.coneType.asKtType()
-            is FirNamedReference -> fir.getReferencedElementType().asKtType()
+            is FirNamedReference -> fir.getReferencedElementType(firResolveState).asKtType()
             is FirStatement -> with(analysisSession) { builtinTypes.UNIT }
             is FirTypeRef, is FirImport, is FirPackageDirective, is FirLabel -> null
             else -> error("Unexpected ${fir?.let { it::class }} for ${expression::class} with text `${expression.text}`")
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/firUtils.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/firUtils.kt
index 157a06dacc7..172c4ef657c 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/firUtils.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/firUtils.kt
@@ -18,6 +18,9 @@ import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtConstantValue
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSimpleConstantValue
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtUnsupportedConstantValue
@@ -37,15 +40,18 @@ internal fun KtExpression.unwrap(): KtExpression {
     } ?: this
 }
 
-internal fun FirNamedReference.getReferencedElementType(): ConeKotlinType {
+internal fun FirNamedReference.getReferencedElementType(resolveState: FirModuleResolveState): ConeKotlinType {
     val symbols = when (this) {
         is FirResolvedNamedReference -> listOf(resolvedSymbol)
         is FirErrorNamedReference -> FirReferenceResolveHelper.getFirSymbolsByErrorNamedReference(this)
         else -> error("Unexpected ${this::class}")
     }
     val firCallableDeclaration = symbols.singleOrNull()?.fir as? FirCallableDeclaration
-    return firCallableDeclaration?.returnTypeRef?.coneType
-        ?: ConeClassErrorType(ConeUnresolvedNameError(name))
+        ?: return ConeClassErrorType(ConeUnresolvedNameError(name))
+
+    return firCallableDeclaration.withFirDeclaration(ResolveType.CallableReturnType, resolveState) {
+        it.returnTypeRef.coneType
+    }
 }
 
 internal fun mapAnnotationParameters(annotationCall: FirAnnotationCall, session: FirSession): Map<String, FirExpression> {