From 9ca8d392403ccd0c33d0732474c3a6b4403a15e6 Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Mon, 30 Aug 2021 11:33:30 -0700
Subject: [PATCH] FIR IDE: use KtVariableWithInvokeFunctionCall for extension
 `invoke`

---
 .../jetbrains/kotlin/analysis/api/fir/FirUtils.kt    |  2 +-
 .../analysis/api/fir/components/KtFirCallResolver.kt |  3 ++-
 .../api/fir/components/ResolveCallTestGenerated.java | 12 +++++++++---
 .../resolveCall/variableAsFunctionLikeCall.kt        |  2 +-
 .../resolveCall/variableAsFunctionLikeCall.txt       |  3 ++-
 .../resolveCall/variableWithExtensionInvoke.kt       |  7 +++++++
 .../resolveCall/variableWithExtensionInvoke.txt      |  5 +++++
 ...ableWithInvoke.kt => variableWithMemberInvoke.kt} |  0
 ...leWithInvoke.txt => variableWithMemberInvoke.txt} |  0
 9 files changed, 27 insertions(+), 7 deletions(-)
 create mode 100644 analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.kt
 create mode 100644 analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.txt
 rename analysis/analysis-api/testData/analysisSession/resolveCall/{variableWithInvoke.kt => variableWithMemberInvoke.kt} (100%)
 rename analysis/analysis-api/testData/analysisSession/resolveCall/{variableWithInvoke.txt => variableWithMemberInvoke.txt} (100%)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirUtils.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirUtils.kt
index 3a521fc1467..be47a590023 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirUtils.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirUtils.kt
@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 fun FirFunctionCall.isImplicitFunctionCall(): Boolean {
-    if (dispatchReceiver !is FirQualifiedAccessExpression) return false
+    if (extensionReceiver !is FirQualifiedAccessExpression && dispatchReceiver !is FirQualifiedAccessExpression) return false
     return calleeReference.getCandidateSymbols().any(FirBasedSymbol<*>::isInvokeFunction)
 }
 
diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt
index eb4b9a3cfab..95e11331ffb 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt
@@ -181,7 +181,8 @@ internal class KtFirCallResolver(
         return when {
             firCall.isImplicitFunctionCall() -> {
                 val target = with(FirReferenceResolveHelper) {
-                    val calleeReference = (firCall.dispatchReceiver as FirQualifiedAccessExpression).calleeReference
+                    val calleeReference = ((firCall.extensionReceiver as? FirQualifiedAccessExpression)
+                        ?: firCall.dispatchReceiver as FirQualifiedAccessExpression).calleeReference
                     calleeReference.toTargetSymbol(session, firSymbolBuilder).singleOrNull()
                 }
                 when (target) {
diff --git a/analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/components/ResolveCallTestGenerated.java b/analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/components/ResolveCallTestGenerated.java
index 0081ce23a25..8270b11665d 100644
--- a/analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/components/ResolveCallTestGenerated.java
+++ b/analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/components/ResolveCallTestGenerated.java
@@ -259,8 +259,14 @@ public class ResolveCallTestGenerated extends AbstractResolveCallTest {
     }
 
     @Test
-    @TestMetadata("variableWithInvoke.kt")
-    public void testVariableWithInvoke() throws Exception {
-        runTest("analysis/analysis-api/testData/analysisSession/resolveCall/variableWithInvoke.kt");
+    @TestMetadata("variableWithExtensionInvoke.kt")
+    public void testVariableWithExtensionInvoke() throws Exception {
+        runTest("analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.kt");
+    }
+
+    @Test
+    @TestMetadata("variableWithMemberInvoke.kt")
+    public void testVariableWithMemberInvoke() throws Exception {
+        runTest("analysis/analysis-api/testData/analysisSession/resolveCall/variableWithMemberInvoke.kt");
     }
 }
diff --git a/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.kt b/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.kt
index 638eec4c231..78124268ff5 100644
--- a/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.kt
+++ b/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.kt
@@ -1,5 +1,5 @@
 operator fun Int.invoke(): String {}
 
-fun call(x: kotlin.int) {
+fun call(x: kotlin.Int) {
     <expr>x()</expr>
 }
diff --git a/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.txt b/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.txt
index ca526504d8e..a29b36a748b 100644
--- a/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.txt
+++ b/analysis/analysis-api/testData/analysisSession/resolveCall/variableAsFunctionLikeCall.txt
@@ -1,4 +1,5 @@
-KtFunctionCall:
+KtVariableWithInvokeFunctionCall:
+target = x: kotlin.Int
 argumentMapping = {  }
 targetFunction = /invoke(<extension receiver>: kotlin.Int): kotlin.String
 substitutor = <empty substitutor>
\ No newline at end of file
diff --git a/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.kt b/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.kt
new file mode 100644
index 00000000000..0ecf6d7d8f7
--- /dev/null
+++ b/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.kt
@@ -0,0 +1,7 @@
+interface Foo
+
+operator fun <T> Foo.invoke(t: T) {}
+
+fun test(f: Foo) {
+    <expr>f("")</expr>
+}
diff --git a/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.txt b/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.txt
new file mode 100644
index 00000000000..9ebfb7ed0b6
--- /dev/null
+++ b/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithExtensionInvoke.txt
@@ -0,0 +1,5 @@
+KtVariableWithInvokeFunctionCall:
+target = f: Foo
+argumentMapping = { "" -> (t: kotlin.String) }
+targetFunction = /invoke(<extension receiver>: Foo, t: kotlin.String): kotlin.Unit
+substitutor = <map substitutor: {FirTypeParameterSymbol T=kotlin/String}>
\ No newline at end of file
diff --git a/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithInvoke.kt b/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithMemberInvoke.kt
similarity index 100%
rename from analysis/analysis-api/testData/analysisSession/resolveCall/variableWithInvoke.kt
rename to analysis/analysis-api/testData/analysisSession/resolveCall/variableWithMemberInvoke.kt
diff --git a/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithInvoke.txt b/analysis/analysis-api/testData/analysisSession/resolveCall/variableWithMemberInvoke.txt
similarity index 100%
rename from analysis/analysis-api/testData/analysisSession/resolveCall/variableWithInvoke.txt
rename to analysis/analysis-api/testData/analysisSession/resolveCall/variableWithMemberInvoke.txt