From 030749eaeb335d1e05c21721b1a9261ce6375b7f Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Tue, 15 Jun 2021 14:24:03 -0700
Subject: [PATCH] FIR IDE: make KtExpressionTypeProvider work for more cases

Currently it doesn't work for the following

* labeled expression
* annotated expressions
* object literal expressions
---
 .../api/fir/components/KtFirExpressionTypeProvider.kt | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

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 9bccd15d79a..864cdb5afdf 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
@@ -35,7 +35,7 @@ internal class KtFirExpressionTypeProvider(
     }
 
     override fun getKtExpressionType(expression: KtExpression): KtType = withValidityAssertion {
-        when (val fir = expression.getOrBuildFir(firResolveState)) {
+        when (val fir = expression.unwrap().getOrBuildFir(firResolveState)) {
             is FirExpression -> fir.typeRef.coneType.asKtType()
             is FirNamedReference -> fir.getReferencedElementType().asKtType()
             is FirStatement -> with(analysisSession) { builtinTypes.UNIT }
@@ -43,6 +43,15 @@ internal class KtFirExpressionTypeProvider(
         }
     }
 
+    private fun KtExpression.unwrap(): KtExpression {
+        return when (this) {
+            is KtLabeledExpression -> baseExpression?.unwrap()
+            is KtAnnotatedExpression -> baseExpression?.unwrap()
+            is KtObjectLiteralExpression -> objectDeclaration
+            else -> null
+        } ?: this
+    }
+
     override fun getExpectedType(expression: PsiElement): KtType? {
         val expectedType = getExpectedTypeByReturnExpression(expression)
             ?: getExpressionTypeByIfOrBooleanCondition(expression)