From e4781d850833b9761f94f81eb2208dfb51fdda20 Mon Sep 17 00:00:00 2001
From: Nikolay Lunyak <nikolay.lunyak@jetbrains.com>
Date: Wed, 22 Nov 2023 14:16:03 +0200
Subject: [PATCH] [FIR] Fix FP UNRESOLVED_REFERENCE on implicit invoke() after
 a safe call

`invoke()` function of functional types is always defined
in such a way that its explicit receiver comes as the first
value parameter. Extension-function types contain the very
same `invoke` function, they just additionally have the
`@ExtensionFunctionType` annotation.

In the related test there is such an `invoke` function and
in the presented call its explicit receiver is the anonymous
function `fun Int.() = 1` itself. So the safe-checked `1`

^KT-60056
---
 .../kotlin/fir/builder/ConversionUtils.kt         | 15 +++++++++++++--
 .../tests/FunctionCalleeExpressions.fir.kt        |  2 +-
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt
index 787a3f58ad1..9776ed7d034 100644
--- a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt
+++ b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt
@@ -612,8 +612,19 @@ fun FirQualifiedAccessExpression.createSafeCall(receiver: FirExpression, source:
         }
         this.source = receiver.source?.fakeElement(KtFakeSourceElementKind.CheckedSafeCallSubject)
     }
-
-    replaceExplicitReceiver(checkedSafeCallSubject)
+    // If an `invoke` function from a functional type expects a
+    // receiver, it's still defined as the first value parameter.
+    // A construction like `1.(fun Int.() = 1)()` means we're calling
+    // `Function1<Int, Unit>.invoke(Int)`.
+    if (this is FirImplicitInvokeCall) {
+        val newArguments = buildArgumentList {
+            arguments.add(checkedSafeCallSubject)
+            arguments.addAll(this@createSafeCall.arguments)
+        }
+        replaceArgumentList(newArguments)
+    } else {
+        replaceExplicitReceiver(checkedSafeCallSubject)
+    }
     return buildSafeCallExpression {
         this.receiver = receiver
         @OptIn(FirContractViolation::class)
diff --git a/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt b/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt
index 2d462b598b7..21088a65ee2 100644
--- a/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt
+++ b/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt
@@ -81,6 +81,6 @@ fun test() {
     val i : Int? = null
     i.<!UNSAFE_IMPLICIT_INVOKE_CALL!>(fun Int.() = 1)<!>();
     <!INAPPLICABLE_CANDIDATE!>{}<!><Int>()
-    1<!UNNECESSARY_SAFE_CALL!>?.<!>(<!UNRESOLVED_REFERENCE!>fun Int.() = 1<!>)()
+    1<!UNNECESSARY_SAFE_CALL!>?.<!>(fun Int.() = 1)()
     1.<!NO_RECEIVER_ALLOWED!>{}<!>()
 }