From c24ad0ba51888df0aff8a72ee09a386049c760dd Mon Sep 17 00:00:00 2001
From: Stanislav Erokhin <Stanislav.Erokhin@jetbrains.com>
Date: Tue, 19 Jan 2021 18:34:19 +0100
Subject: [PATCH] Use KtFunctionLiteral as key for KtAnonymousFunctionSymbol
 retrieving

It seems logical to do, because KtFunctionLiteral is KtDeclaration
and KtLambdaExpression is not. Also FirAnonymousFunction use it as
realPsi reference
---
 .../idea/frontend/api/symbols/KtSymbolProvider.kt     | 11 +++--------
 .../frontend/api/fir/symbols/KtFirSymbolProvider.kt   |  2 +-
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt
index a463c1a1c8f..527d4296f32 100644
--- a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt
+++ b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt
@@ -20,18 +20,13 @@ abstract class KtSymbolProvider : KtAnalysisSessionComponent() {
         is KtTypeParameter -> getTypeParameterSymbol(psi)
         is KtTypeAlias -> getTypeAliasSymbol(psi)
         is KtEnumEntry -> getEnumEntrySymbol(psi)
-        is KtLambdaExpression -> getAnonymousFunctionSymbol(psi)
+        is KtFunctionLiteral -> getAnonymousFunctionSymbol(psi)
         is KtProperty -> getVariableSymbol(psi)
         is KtClassOrObject -> {
             val literalExpression = (psi as? KtObjectDeclaration)?.parent as? KtObjectLiteralExpression
             literalExpression?.let(::getAnonymousObjectSymbol) ?: getClassOrObjectSymbol(psi)
         }
         is KtPropertyAccessor -> getPropertyAccessorSymbol(psi)
-        is KtFunctionLiteral -> {
-            val literalExpression = psi.parent as? KtLambdaExpression
-            check(literalExpression != null) { "Unexpected parent for KtFunctionLiteral" }
-            getAnonymousFunctionSymbol(literalExpression)
-        }
         else -> error("Cannot build symbol for ${psi::class}")
     }
 
@@ -43,7 +38,7 @@ abstract class KtSymbolProvider : KtAnalysisSessionComponent() {
     abstract fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol
     abstract fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol
     abstract fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol
-    abstract fun getAnonymousFunctionSymbol(psi: KtLambdaExpression): KtAnonymousFunctionSymbol
+    abstract fun getAnonymousFunctionSymbol(psi: KtFunctionLiteral): KtAnonymousFunctionSymbol
     abstract fun getVariableSymbol(psi: KtProperty): KtVariableSymbol
     abstract fun getAnonymousObjectSymbol(psi: KtObjectLiteralExpression): KtAnonymousObjectSymbol
     abstract fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol
@@ -86,7 +81,7 @@ interface KtSymbolProviderMixIn : KtAnalysisSessionMixIn {
     fun KtNamedFunction.getAnonymousFunctionSymbol(): KtAnonymousFunctionSymbol =
         analysisSession.symbolProvider.getAnonymousFunctionSymbol(this)
 
-    fun KtLambdaExpression.getAnonymousFunctionSymbol(): KtAnonymousFunctionSymbol =
+    fun KtFunctionLiteral.getAnonymousFunctionSymbol(): KtAnonymousFunctionSymbol =
         analysisSession.symbolProvider.getAnonymousFunctionSymbol(this)
 
     fun KtProperty.getVariableSymbol(): KtVariableSymbol =
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt
index d69851d5b28..62adb5e4149 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt
@@ -85,7 +85,7 @@ internal class KtFirSymbolProvider(
         firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType(resolveState))
     }
 
-    override fun getAnonymousFunctionSymbol(psi: KtLambdaExpression): KtAnonymousFunctionSymbol = withValidityAssertion {
+    override fun getAnonymousFunctionSymbol(psi: KtFunctionLiteral): KtAnonymousFunctionSymbol = withValidityAssertion {
         psi.withFirDeclarationOfType<FirAnonymousFunction, KtAnonymousFunctionSymbol>(resolveState) {
             firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(it)
         }