FIR: if necessary, use explicit receiver as argument in invoke collector

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CandidateCollector.kt

@@ -5,11 +5,14 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls
 
+import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirQualifiedAccessExpressionImpl
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.transformQualifiedAccessUsingSmartcastInfo
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.firUnsafe
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.types.isExtensionFunctionType
+import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 open class CandidateCollector(
@@ -80,16 +83,21 @@ class InvokeReceiverCandidateCollector(
         val applicability = super.consumeCandidate(group, candidate)
 
         if (applicability >= CandidateApplicability.SYNTHETIC_RESOLVED) {
-
             val session = components.session
+            val symbol = candidate.symbol as FirCallableSymbol<*>
+            val extensionReceiverExpression = candidate.extensionReceiverExpression()
+            val useExtensionReceiverAsArgument =
+                symbol.fir.receiverTypeRef == null &&
+                        candidate.explicitReceiverKind == ExplicitReceiverKind.EXTENSION_RECEIVER &&
+                        symbol.fir.returnTypeRef.isExtensionFunctionType()
             val explicitReceiver = FirQualifiedAccessExpressionImpl(null).apply {
                 calleeReference = FirNamedReferenceWithCandidate(
                     null,
-                    (candidate.symbol as FirCallableSymbol<*>).callableId.callableName,
+                    symbol.callableId.callableName,
                     candidate
                 )
                 dispatchReceiver = candidate.dispatchReceiverExpression()
-                extensionReceiver = candidate.extensionReceiverExpression()
+                extensionReceiver = extensionReceiverExpression.takeIf { !useExtensionReceiverAsArgument } ?: FirNoReceiverExpression
                 typeRef = towerResolver.typeCalculator.tryCalculateReturnType(candidate.symbol.firUnsafe())
             }.let {
                 components.transformQualifiedAccessUsingSmartcastInfo(it)
@@ -107,7 +115,10 @@ class InvokeReceiverCandidateCollector(
                 invokeCallInfo.outerCSBuilder,
                 invokeCallInfo.lhs,
                 invokeCallInfo.typeProvider
-            )
+            ).let {
+                if (useExtensionReceiverAsArgument) it.withReceiverAsArgument(extensionReceiverExpression)
+                else it
+            }
 
             invokeConsumer.addConsumerAndProcessAccumulatedData(
                 createSimpleFunctionConsumer(

compiler/fir/resolve/testData/resolve/expresssions/invoke/inBrackets.txt

@@ -1,5 +1,5 @@
 FILE: inBrackets.kt
     public final fun test(e: R|kotlin/Int.() -> kotlin/String|): R|kotlin/Unit| {
-        lval s: R|kotlin/String| = Int(3).R|<local>/e|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/String|>|(Int(3))
+        lval s: R|kotlin/String| = R|<local>/e|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/String|>|(Int(3))
-        lval ss: R|kotlin/String| = Int(3).R|<local>/e|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/String|>|(Int(3))
+        lval ss: R|kotlin/String| = R|<local>/e|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/String|>|(Int(3))
     }

compiler/fir/resolve/testData/resolve/expresssions/invoke/propertyWithExtensionType.txt

@@ -15,10 +15,10 @@ FILE: propertyWithExtensionType.kt
         when () {
             !=(R|<local>/a|.R|/A.x|, Null(null)) ->  {
                 lval b: R|kotlin/String.() -> kotlin/Unit| = R|<local>/a|.R|/A.x|
-                String().R|<local>/b|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(String())
+                R|<local>/b|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(String())
             }
         }
 
         lval c: R|kotlin/String.() -> kotlin/Int| = R|<local>/a|.R|/A.y|
-        lval d: R|kotlin/Int| = String().R|<local>/c|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Int|>|(String())
+        lval d: R|kotlin/Int| = R|<local>/c|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Int|>|(String())
     }

compiler/fir/resolve/testData/resolve/expresssions/lambdaWithReceiver.txt

@@ -4,10 +4,10 @@ FILE: lambdaWithReceiver.kt
 
     }
     public final fun <T> myWith(receiver: R|T|, block: R|T.() -> kotlin/Unit|): R|kotlin/Unit| {
-        R|<local>/receiver|.R|<local>/block|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/receiver|)
+        R|<local>/block|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/receiver|)
     }
     public final fun <T> R|T|.myApply(block: R|T.() -> kotlin/Unit|): R|kotlin/Unit| {
-        this@R|/myApply|.R|<local>/block|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(this@R|/myApply|)
+        R|<local>/block|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(this@R|/myApply|)
     }
     public final fun withA(block: R|A.() -> kotlin/Unit|): R|kotlin/Unit| {
     }

compiler/fir/resolve/testData/resolve/functionTypes.txt

@@ -5,7 +5,7 @@ FILE: functionTypes.kt
     public final fun <T, R> R|kotlin/collections/List<T>|.simpleMap(f: R|(T) -> R|): R|R| {
     }
     public final fun <T> simpleWith(t: R|T|, f: R|T.() -> kotlin/Unit|): R|kotlin/Unit| {
-        ^simpleWith R|<local>/t|.R|<local>/f|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/t|)
+        ^simpleWith R|<local>/f|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/t|)
     }
     public abstract interface KMutableProperty1<T, R> : R|KProperty1<T, R>|, R|KMutableProperty<R>| {
     }

compiler/fir/resolve/testData/resolve/problems/invokeOfLambdaWithReceiver.txt

@@ -2,7 +2,7 @@ FILE: invokeOfLambdaWithReceiver.kt
     public abstract interface A : R|kotlin/Any| {
     }
     public final fun test(a: R|A|, block: R|A.() -> kotlin/Int|): R|kotlin/Unit| {
-        R|<local>/a|.R|<local>/block|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Int|>|(R|<local>/a|)
+        R|<local>/block|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Int|>|(R|<local>/a|)
     }
     public abstract interface B : R|kotlin/Any| {
     }