[FIR] Add receiver for lambda in position of default argument

#KT-36905 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -542,7 +542,10 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
                 )
                 af = af.transformValueParameters(ImplicitToErrorTypeTransformer, null)
                 val bodyExpectedType = returnTypeRefFromResolvedAtom ?: expectedTypeRef
-                af = transformFunction(af, withExpectedType(bodyExpectedType)).single as FirAnonymousFunction
+                val labelName = af.label?.name?.let { Name.identifier(it) }
+                withLabelAndReceiverType(labelName, af, af.receiverTypeRef?.coneTypeSafe()) {
+                    af = transformFunction(af, withExpectedType(bodyExpectedType)).single as FirAnonymousFunction
+                }
                 // To separate function and separate commit
                 val writer = FirCallCompletionResultsWriterTransformer(
                     session,
@@ -601,49 +604,58 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
     private inline fun <T> withLabelAndReceiverType(
         labelName: Name?,
         owner: FirDeclaration,
-        type: ConeKotlinType,
+        type: ConeKotlinType?,
         block: () -> T
     ): T {
-        val implicitCompanionValues = mutableListOf<ImplicitReceiverValue<*>>()
+        val implicitCompanionValues: List<ImplicitReceiverValue<*>>
-        val implicitReceiverValue = when (owner) {
+        if (type != null) {
-            is FirClass<*> -> {
+            implicitCompanionValues = mutableListOf()
-                // Questionable: performance
+            val implicitReceiverValue = when (owner) {
-                (owner as? FirRegularClass)?.companionObject?.let { companion ->
+                is FirClass<*> -> {
-                    implicitCompanionValues += ImplicitDispatchReceiverValue(
+                    // Questionable: performance
-                        companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION
+                    (owner as? FirRegularClass)?.companionObject?.let { companion ->
-                    )
-                }
-                lookupSuperTypes(owner, lookupInterfaces = false, deep = true, useSiteSession = session).mapNotNull {
-                    val superClass = (it as? ConeClassLikeType)?.lookupTag?.toSymbol(session)?.fir as? FirRegularClass
-                    superClass?.companionObject?.let { companion ->
                         implicitCompanionValues += ImplicitDispatchReceiverValue(
-                            companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION_FROM_SUPERTYPE
+                            companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION
                         )
                     }
+                    lookupSuperTypes(owner, lookupInterfaces = false, deep = true, useSiteSession = session).mapNotNull {
+                        val superClass = (it as? ConeClassLikeType)?.lookupTag?.toSymbol(session)?.fir as? FirRegularClass
+                        superClass?.companionObject?.let { companion ->
+                            implicitCompanionValues += ImplicitDispatchReceiverValue(
+                                companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION_FROM_SUPERTYPE
+                            )
+                        }
+                    }
+                    // ---
+                    ImplicitDispatchReceiverValue(owner.symbol, type, session, scopeSession)
+                }
+                is FirFunction<*> -> {
+                    ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+                }
+                is FirVariable<*> -> {
+                    ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+                }
+                else -> {
+                    throw IllegalArgumentException("Incorrect label & receiver owner: ${owner.javaClass}")
                 }
-                // ---
-                ImplicitDispatchReceiverValue(owner.symbol, type, session, scopeSession)
             }
-            is FirFunction<*> -> {
+            for (implicitCompanionValue in implicitCompanionValues.asReversed()) {
-                ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+                implicitReceiverStack.add(null, implicitCompanionValue)
             }
-            is FirVariable<*> -> {
+            implicitReceiverStack.add(labelName, implicitReceiverValue)
-                ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+        } else {
-            }
+            implicitCompanionValues = emptyList()
-            else -> {
+        }
-                throw IllegalArgumentException("Incorrect label & receiver owner: ${owner.javaClass}")
+        try {
+            return block()
+        } finally {
+            if (type != null) {
+                implicitReceiverStack.pop(labelName)
+                for (implicitCompanionValue in implicitCompanionValues) {
+                    implicitReceiverStack.pop(null)
+                }
             }
         }
-        for (implicitCompanionValue in implicitCompanionValues.asReversed()) {
-            implicitReceiverStack.add(null, implicitCompanionValue)
-        }
-        implicitReceiverStack.add(labelName, implicitReceiverValue)
-        val result = block()
-        implicitReceiverStack.pop(labelName)
-        for (implicitCompanionValue in implicitCompanionValues) {
-            implicitReceiverStack.pop(null)
-        }
-        return result
     }
 
     private fun storeVariableReturnType(variable: FirVariable<*>) {

compiler/fir/resolve/testData/resolve/arguments/extensionLambdaInDefaultArgument.kt

@@ -0,0 +1,5 @@
+fun test(
+    val f: String.() -> Int = { length }
+): Int {
+    return "".f()
+}

compiler/fir/resolve/testData/resolve/arguments/extensionLambdaInDefaultArgument.txt

@@ -1,6 +1,6 @@
 FILE: extensionLambdaInDefaultArgument.kt
-    public final fun test(f: R|kotlin/String.() -> kotlin/Int| = fun R|kotlin/String|.<anonymous>(): R|ERROR CLASS: Unresolved name: length| {
+    public final fun test(f: R|kotlin/String.() -> kotlin/Int| = fun R|kotlin/String|.<anonymous>(): R|kotlin/Int| {
-        ^ <Unresolved name: length>#
+        ^ this@R|special/anonymous|.R|kotlin/String.length|
     }
     ): R|kotlin/Int| {
         ^test R|<local>/f|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Int|>|(String())

compiler/fir/resolve/testData/resolve/problems/extensionLambdaInDefaultArgument.kt

@@ -1,5 +0,0 @@
-fun test(
-    val f: String.() -> Int = { <!UNRESOLVED_REFERENCE!>length<!> }
-): Int {
-    return "".f()
-}

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -380,6 +380,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/arguments/defaultFromOverrides.kt");
         }
 
+        @TestMetadata("extensionLambdaInDefaultArgument.kt")
+        public void testExtensionLambdaInDefaultArgument() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/arguments/extensionLambdaInDefaultArgument.kt");
+        }
+
         @TestMetadata("fieldPlusAssign.kt")
         public void testFieldPlusAssign() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/arguments/fieldPlusAssign.kt");
@@ -1352,11 +1357,6 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/problems/definitelyNotNullAndOriginalType.kt");
         }
 
-        @TestMetadata("extensionLambdaInDefaultArgument.kt")
-        public void testExtensionLambdaInDefaultArgument() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/problems/extensionLambdaInDefaultArgument.kt");
-        }
-
         @TestMetadata("javaAccessorConversion.kt")
         public void testJavaAccessorConversion() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/problems/javaAccessorConversion.kt");

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -380,6 +380,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/resolve/testData/resolve/arguments/defaultFromOverrides.kt");
         }
 
+        @TestMetadata("extensionLambdaInDefaultArgument.kt")
+        public void testExtensionLambdaInDefaultArgument() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/arguments/extensionLambdaInDefaultArgument.kt");
+        }
+
         @TestMetadata("fieldPlusAssign.kt")
         public void testFieldPlusAssign() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/arguments/fieldPlusAssign.kt");
@@ -1352,11 +1357,6 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/resolve/testData/resolve/problems/definitelyNotNullAndOriginalType.kt");
         }
 
-        @TestMetadata("extensionLambdaInDefaultArgument.kt")
-        public void testExtensionLambdaInDefaultArgument() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/problems/extensionLambdaInDefaultArgument.kt");
-        }
-
         @TestMetadata("javaAccessorConversion.kt")
         public void testJavaAccessorConversion() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/problems/javaAccessorConversion.kt");

compiler/testData/diagnostics/tests/functionAsExpression/ScopeCheck.fir.kt

@@ -9,7 +9,7 @@ fun test(param: String) {
     val local_val = 4
     val bar = fun B.(fun_param: Int) {
         param.length
-        <!UNRESOLVED_REFERENCE!>b_fun<!>()
+        b_fun()
         val inner_bar = local_val + fun_param
 
         <!UNRESOLVED_REFERENCE!>bar<!>

compiler/testData/diagnostics/tests/labels/kt591.fir.kt

@@ -4,7 +4,7 @@ fun test() {
     val a: (Int?).() -> Unit = a@{
         if (this != null) {
             val b: String.() -> Unit = {
-                this@a.<!UNRESOLVED_REFERENCE!>times<!>(5) // a@ Unresolved
+                this@a.times(5) // a@ Unresolved
             }
         }
     }

compiler/testData/diagnostics/tests/regressions/kt306.fir.kt

@@ -2,15 +2,15 @@
 
 fun test() {
     (fun Foo.() {
-        <!UNRESOLVED_REFERENCE!>bar<!>()
+        bar()
         (fun Barr.() {
-            this.<!UNRESOLVED_REFERENCE!>bar<!>()
+            this.bar()
-            <!UNRESOLVED_REFERENCE!>bar<!>()
+            bar()
         })
     })
     (fun Barr.() {
-        this.<!UNRESOLVED_REFERENCE!>bar<!>()
+        this.bar()
-        <!UNRESOLVED_REFERENCE!>bar<!>()
+        bar()
     })
 }
 

compiler/testData/diagnostics/tests/thisAndSuper/thisInFunctionLiterals.fir.kt

@@ -5,14 +5,14 @@ class A(val a:Int) {
     fun Byte.xx() : Double.() -> Any {
       checkSubtype<Byte>(this)
       val a: Double.() -> Unit = {
-        <!INAPPLICABLE_CANDIDATE!>checkSubtype<!><Double>(this)
+        checkSubtype<Double>(this)
         checkSubtype<Byte>(this@xx)
         checkSubtype<B>(this@B)
         checkSubtype<A>(this@A)
       }
-      val b: Double.() -> Unit = a@{ <!INAPPLICABLE_CANDIDATE!>checkSubtype<!><Double>(this@a) + checkSubtype<Byte>(this@xx) }
+      val b: Double.() -> Unit = a@{ checkSubtype<Double>(this@a) + checkSubtype<Byte>(this@xx) }
       val c = a@{ -> this@a + checkSubtype<Byte>(this@xx) }
-      return (a@{<!INAPPLICABLE_CANDIDATE!>checkSubtype<!><Double>(this@a) + checkSubtype<Byte>(this@xx)})
+      return (a@{checkSubtype<Double>(this@a) + checkSubtype<Byte>(this@xx)})
     }
   }
 }