From 5498f40ce76d32858066b520066b7b49e919f493 Mon Sep 17 00:00:00 2001
From: Igor Yakovlev <igor.yakovlev@jetbrains.com>
Date: Tue, 17 Aug 2021 12:52:50 +0200
Subject: [PATCH] [FIR IDE] Fix reference resolve for functions call

---
 .../expressions/binaryExpressionOperator.txt     |  4 ++--
 .../expressions/unraryExpressionOperator.txt     |  4 ++--
 .../idea/references/FirReferenceResolveHelper.kt | 16 ++++++++++++++--
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/binaryExpressionOperator.txt b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/binaryExpressionOperator.txt
index 2b4a5d4cd7b..2a258ddeebb 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/binaryExpressionOperator.txt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/binaryExpressionOperator.txt
@@ -1,5 +1,5 @@
 KT element: KtOperationReferenceExpression
-FIR element: FirFunctionCallImpl
+FIR element: FirResolvedNamedReferenceImpl
 
 FIR element rendered:
-Int(1).R|kotlin/Int.plus|(Int(1))
\ No newline at end of file
+R|kotlin/Int.plus|
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/unraryExpressionOperator.txt b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/unraryExpressionOperator.txt
index f0b6d124c7b..df9f066bf14 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/unraryExpressionOperator.txt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/testdata/getOrBuildFir/expressions/unraryExpressionOperator.txt
@@ -1,5 +1,5 @@
 KT element: KtOperationReferenceExpression
-FIR element: FirFunctionCallImpl
+FIR element: FirResolvedNamedReferenceImpl
 
 FIR element rendered:
-R|<local>/x|.R|kotlin/Boolean.not|()
\ No newline at end of file
+R|kotlin/Boolean.not|
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt
index f0a475116d8..1d640e5dd3b 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt
@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.getCandidateSymbols
 import org.jetbrains.kotlin.idea.fir.isImplicitFunctionCall
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
@@ -189,7 +190,7 @@ internal object FirReferenceResolveHelper {
             is FirReturnExpression -> getSymbolsByReturnExpression(expression, fir, symbolBuilder)
             is FirErrorNamedReference -> getSymbolsByErrorNamedReference(fir, symbolBuilder)
             is FirVariableAssignment -> getSymbolsByVariableAssignment(fir, session, symbolBuilder)
-            is FirResolvedNamedReference -> getSymbolByResolvedNameReference(fir, session, symbolBuilder)
+            is FirResolvedNamedReference -> getSymbolByResolvedNameReference(fir, expression, analysisSession, session, symbolBuilder)
             is FirResolvable -> getSymbolsByResolvable(fir, expression, session, symbolBuilder)
             is FirNamedArgumentExpression -> getSymbolsByNameArgumentExpression(expression, analysisSession, symbolBuilder)
             else -> handleUnknownFirElement(expression, analysisSession, session, symbolBuilder)
@@ -206,9 +207,20 @@ internal object FirReferenceResolveHelper {
 
     private fun getSymbolByResolvedNameReference(
         fir: FirResolvedNamedReference,
+        expression: KtSimpleNameExpression,
+        analysisSession: KtFirAnalysisSession,
         session: FirSession,
         symbolBuilder: KtSymbolByFirBuilder
-    ): Collection<KtSymbol> = fir.toTargetSymbol(session, symbolBuilder)
+    ): Collection<KtSymbol> {
+        val parentAsCall = expression.parent as? KtCallExpression
+        if (parentAsCall != null) {
+            val firResolvable = parentAsCall.getOrBuildFirSafe<FirResolvable>(analysisSession.firResolveState)
+            if (firResolvable != null) {
+                return getSymbolsByResolvable(firResolvable, expression, session, symbolBuilder)
+            }
+        }
+        return fir.toTargetSymbol(session, symbolBuilder)
+    }
 
     private fun KtSimpleNameExpression.isSyntheticOperatorReference() = when (this) {
         is KtOperationReferenceExpression -> operationSignTokenType in syntheticTokenTypes