FIR: Support typeRef-name label for this

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

@@ -36,10 +36,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirWhenSubjectImportingScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames.UNDERSCORE_FOR_UNUSED_VAR
 
@@ -204,8 +201,8 @@ class BodyResolveContext(
     }
 
     @PrivateForInline
-    fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>) {
+    fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>, additionalLabelName: Name? = null) {
-        replaceTowerDataContext(towerDataContext.addReceiver(name, implicitReceiverValue))
+        replaceTowerDataContext(towerDataContext.addReceiver(name, implicitReceiverValue, additionalLabelName))
     }
 
     @PrivateForInline
@@ -225,6 +222,7 @@ class BodyResolveContext(
         owner: FirCallableDeclaration,
         type: ConeKotlinType?,
         holder: SessionHolder,
+        additionalLabelName: Name? = null,
         f: () -> T
     ): T = withTowerDataCleanup {
         replaceTowerDataContext(towerDataContext.addContextReceiverGroup(owner.createContextReceiverValues(holder)))
@@ -236,7 +234,7 @@ class BodyResolveContext(
                 holder.session,
                 holder.scopeSession
             )
-            addReceiver(labelName, receiver)
+            addReceiver(labelName, receiver, additionalLabelName)
         }
 
         f()
@@ -575,13 +573,19 @@ class BodyResolveContext(
                     storeVariable(parameter, holder.session)
                 }
                 val receiverTypeRef = function.receiverTypeRef
-                withLabelAndReceiverType(function.name, function, receiverTypeRef?.coneType, holder, f)
+                val type = receiverTypeRef?.coneType
+                val additionalLabelName = type?.labelName()
+                withLabelAndReceiverType(function.name, function, type, holder, additionalLabelName, f)
             } else {
                 f()
             }
         }
     }
 
+    private fun ConeKotlinType.labelName(): Name? {
+        return (this as? ConeLookupTagBasedType)?.lookupTag?.name
+    }
+
     @OptIn(PrivateForInline::class)
     fun <T> forConstructorBody(
         constructor: FirConstructor,
@@ -714,7 +718,9 @@ class BodyResolveContext(
                 storeBackingField(property, holder.session)
             }
             withContainer(accessor) {
-                withLabelAndReceiverType(property.name, property, receiverTypeRef?.coneType, holder, f)
+                val type = receiverTypeRef?.coneType
+                val additionalLabelName = type?.labelName()
+                withLabelAndReceiverType(property.name, property, type, holder, additionalLabelName, f)
             }
         }
     }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/declarations/ImplicitReceiverUtils.kt

@@ -168,11 +168,11 @@ class FirTowerDataContext private constructor(
         )
     }
 
-    fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>): FirTowerDataContext {
+    fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>, additionalLabName: Name? = null): FirTowerDataContext {
         val element = implicitReceiverValue.asTowerDataElement()
         return FirTowerDataContext(
             towerDataElements.add(element),
-            implicitReceiverStack.add(name, implicitReceiverValue),
+            implicitReceiverStack.add(name, implicitReceiverValue, additionalLabName),
             localScopes,
             nonLocalTowerDataElements.add(element)
         )

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/PersistentImplicitReceiverStack.kt

@@ -42,7 +42,7 @@ class PersistentImplicitReceiverStack private constructor(
         val stack = stack.add(value)
         val originalTypes = originalTypes.add(value.originalType)
         val index = stack.size - 1
-        val receiversPerLabel = name?.let { receiversPerLabel.put(it, value) } ?: receiversPerLabel
+        val receiversPerLabel = receiversPerLabel.putIfNameIsNotNull(name, value).putIfNameIsNotNull(aliasLabel, value)
         val indexesPerSymbol = indexesPerSymbol.put(value.boundSymbol, index)
 
         return PersistentImplicitReceiverStack(
@@ -53,6 +53,12 @@ class PersistentImplicitReceiverStack private constructor(
         )
     }
 
+    private fun PersistentSetMultimap<Name, ImplicitReceiverValue<*>>.putIfNameIsNotNull(name: Name?, value: ImplicitReceiverValue<*>) =
+        if (name != null)
+            put(name, value)
+        else
+            this
+
     fun addContextReceiver(value: ContextReceiverValue<*>): PersistentImplicitReceiverStack {
         val labelName = value.labelName ?: return this
 

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/decimateEveryEvenThird.fir.kt

@@ -1,12 +0,0 @@
-// !LANGUAGE: +ContextReceivers
-// WITH_STDLIB
-
-fun List<Int>.decimateEveryEvenThird() = <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>sequence<!> {
-    var counter = 1
-    for (e in <!ITERATOR_AMBIGUITY!>this<!UNRESOLVED_LABEL!>@List<!><!>) {
-        if (e <!NONE_APPLICABLE!>%<!> 2 == 0 && counter % 3 == 0) {
-            yield(e)
-        }
-        counter += 1
-    }
-}

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/decimateEveryEvenThird.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +ContextReceivers
 // WITH_STDLIB
 

compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.fir.kt

@@ -1,72 +0,0 @@
-// !LANGUAGE: +ContextReceivers
-
-class Param
-class C {
-    val c = 42
-}
-class R {
-    val r = 42
-}
-
-context(C)
-fun R.f1(g: context(C) R.(Param) -> Unit) {
-    g(this@C, this<!UNRESOLVED_LABEL!>@R<!>, Param())
-}
-
-context(C)
-fun f2(g: context(C) (Param) -> Unit) {
-    g(this@C, Param())
-}
-
-context(C)
-fun R.f3(g: context(C) R.() -> Unit) {
-    g(this@C, this<!UNRESOLVED_LABEL!>@R<!>)
-}
-
-context(C)
-fun f4(g: context(C) () -> Unit) {
-    g(this@C)
-}
-
-fun test() {
-    val lf1: context(C) R.(Param) -> Unit = { _ ->
-        r
-        c
-    }
-    val lf2: context(C) (Param) -> Unit = { _ ->
-        c
-    }
-    val lf3: context(C) R.() -> Unit = {
-        r
-        c
-    }
-    val lf4: context(C) () -> Unit = {
-        c
-    }
-
-    with(C()) {
-        with(R()) {
-            f1(lf1)
-            f1 { _ ->
-                r
-                c
-            }
-
-            f2(lf2)
-            f2 { _ ->
-                c
-            }
-
-            f3(lf3)
-            f3 {
-                r
-                c
-            }
-
-            f4(lf4)
-            f4 {
-                c
-            }
-        }
-    }
-}

compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +ContextReceivers
 
 class Param

compiler/testData/diagnostics/tests/extensions/contextReceivers/labelsFromClassNameForbidden.fir.kt

@@ -1,11 +1,11 @@
 fun List<Int>.f() {
-    this<!UNRESOLVED_LABEL!>@List<!>.size
+    this@List.size
 }
 
 context(String)
 fun Int.f() {
     this@String.length
-    this<!UNRESOLVED_LABEL!>@Int<!>.toDouble()
+    this@Int.toDouble()
 }
 
 context(String)

compiler/testData/diagnostics/tests/extensions/contextReceivers/noLabelsByClassName.fir.kt

@@ -1,27 +1,27 @@
 fun Int.f() {
-    this<!UNRESOLVED_LABEL!>@Int<!>
+    this@Int
 }
 
 var Int.p: Int
     get() {
-        this<!UNRESOLVED_LABEL!>@Int<!>
+        this@Int
         <!RETURN_NOT_ALLOWED!>return@p<!> 42
     }
     set(value) {
-        this<!UNRESOLVED_LABEL!>@Int<!>
+        this@Int
     }
 
 class X {
     var Int.p: Int
         get() {
-            this<!UNRESOLVED_LABEL!>@Int<!>
+            this@Int
             <!RETURN_NOT_ALLOWED!>return@p<!> 42
         }
         set(value) {
-            this<!UNRESOLVED_LABEL!>@Int<!>
+            this@Int
         }
 
     fun Int.f() {
-        this<!UNRESOLVED_LABEL!>@Int<!>
+        this@Int
     }
 }

compiler/testData/diagnostics/tests/extensions/contextReceivers/thisWithCustomLabel.fir.kt

@@ -17,7 +17,7 @@ context(labelAInt@A<Int>, A<String>, labelB@B) val C.p: Int
         this@A.a.length
         this<!UNRESOLVED_LABEL!>@B<!>
         this@labelB.b
-        this<!UNRESOLVED_LABEL!>@C<!>.c
+        this@C.c
         this@p.c
         this.c
         return 1

compiler/testData/diagnostics/tests/extensions/contextReceivers/thisWithReceiverLabelsFunctions.fir.kt

@@ -23,7 +23,7 @@ context(A<Int>, A<String>, B) fun f() {
 context(A<Int>, A<String>, B) fun C.f() {
     this@A.a.length
     this@B.b
-    this<!UNRESOLVED_LABEL!>@C<!>.c
+    this@C.c
     this@f.c
     this.c
 }

compiler/testData/diagnostics/tests/extensions/contextReceivers/thisWithReceiverLabelsProperties.fir.kt

@@ -37,7 +37,7 @@ context(A<Int>, A<String>, B) val C.p: Int
     get() {
         this@A.a.length
         this@B.b
-        this<!UNRESOLVED_LABEL!>@C<!>.c
+        this@C.c
         this@p.c
         this.c
         return 1

compiler/testData/diagnostics/testsWithStdLib/labelClashesWithContextReceivers.fir.kt

@@ -6,7 +6,7 @@ class Some {
     fun foo() {
         //this@foo
         this@Some
-        this<!UNRESOLVED_LABEL!>@String<!>
+        this@String
     }
 
     context(Some)