K1: Support referencing class context receivers in a form of this@Name

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -123,6 +123,7 @@ public interface BindingContext {
     WritableSlice<KtSuperExpression, Boolean> SUPER_EXPRESSION_FROM_ANY_MIGRATION = Slices.createSimpleSlice();
 
     WritableSlice<KtReferenceExpression, DeclarationDescriptor> REFERENCE_TARGET = new BasicWritableSlice<>(DO_NOTHING);
+    WritableSlice<KtReferenceExpression, ReceiverParameterDescriptor> THIS_REFERENCE_TARGET = new BasicWritableSlice<>(DO_NOTHING);
     // if 'A' really means 'A.Companion' then this slice stores class descriptor for A, REFERENCE_TARGET stores descriptor Companion in this case
     WritableSlice<KtReferenceExpression, ClassifierDescriptorWithTypeParameters> SHORT_REFERENCE_TO_COMPANION_OBJECT =
             new BasicWritableSlice<>(DO_NOTHING);

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/LazyClassDescriptor.java

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.lazy.descriptors;
 
+import com.google.common.collect.HashMultimap;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiNameIdentifierOwner;
 import kotlin.Pair;
@@ -305,7 +306,8 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
             if (classOrObject == null) {
                 return CollectionsKt.emptyList();
             }
-            return classOrObject.getContextReceivers().stream()
+            List<KtContextReceiver> contextReceivers = classOrObject.getContextReceivers();
+            List<ReceiverParameterDescriptor> contextReceiverDescriptors = contextReceivers.stream()
                     .map(KtContextReceiver::typeReference)
                     .filter(Objects::nonNull)
                     .map(typeReference -> {
@@ -316,7 +318,19 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
                                 kotlinType,
                                 Annotations.Companion.getEMPTY()
                         );
-            }).collect(Collectors.toList());
+                    }).collect(Collectors.toList());
+
+            if (c.getLanguageVersionSettings().supportsFeature(LanguageFeature.ContextReceivers)) {
+                HashMultimap<String, ReceiverParameterDescriptor> labelNameToReceiverMap = HashMultimap.create();
+
+                for (int i = 0; i < contextReceivers.size(); i++) {
+                    labelNameToReceiverMap.put(contextReceivers.get(i).name(), contextReceiverDescriptors.get(i));
+                }
+
+                c.getTrace().record(BindingContext.DESCRIPTOR_TO_CONTEXT_RECEIVER_MAP, this, labelNameToReceiverMap);
+            }
+
+            return contextReceiverDescriptors;
         });
     }
 

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/BasicExpressionTypingVisitor.java

@@ -404,7 +404,9 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
                 context.trace.report(NO_THIS.on(expression));
                 break;
             case SUCCESS:
-                result = resolutionResult.getReceiverParameterDescriptor().getType();
+                ReceiverParameterDescriptor descriptor = resolutionResult.getReceiverParameterDescriptor();
+                context.trace.record(THIS_REFERENCE_TARGET, expression.getInstanceReference(), descriptor);
+                result = descriptor.getType();
                 context.trace.recordType(expression.getInstanceReference(), result);
                 break;
         }

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/LabelResolver.kt

@@ -25,8 +25,11 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.checkReservedYield
 import org.jetbrains.kotlin.psi.psiUtil.parents
-import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.BindingContext.*
+import org.jetbrains.kotlin.resolve.BindingContextUtils
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.DescriptorResolver
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
 import org.jetbrains.kotlin.resolve.scopes.utils.getDeclarationsByLabel
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -64,6 +67,12 @@ object LabelResolver {
             is KtFunctionLiteral -> return getLabelNamesIfAny(element.parent, false)
             is KtLambdaExpression -> result.addIfNotNull(getLabelForFunctionalExpression(element))
         }
+
+        if (element is KtClass) {
+            element.contextReceivers
+                .mapNotNullTo(result) { it.name()?.let { s -> Name.identifier(s) } }
+        }
+
         val functionOrProperty = when (element) {
             is KtNamedFunction -> {
                 result.addIfNotNull(element.nameAsName ?: getLabelForFunctionalExpression(element))
@@ -210,14 +219,15 @@ object LabelResolver {
                     trace.record(LABEL_TARGET, targetLabelExpression, it)
                 }
                 val declarationDescriptor = trace.bindingContext[DECLARATION_TO_DESCRIPTOR, element]
-                if (declarationDescriptor is FunctionDescriptor) {
+                if (declarationDescriptor is FunctionDescriptor || declarationDescriptor is ClassDescriptor) {
                     val labelNameToReceiverMap = trace.bindingContext[
                             DESCRIPTOR_TO_CONTEXT_RECEIVER_MAP,
                             if (declarationDescriptor is PropertyAccessorDescriptor) declarationDescriptor.correspondingProperty else declarationDescriptor
                     ]
                     val thisReceivers = labelNameToReceiverMap?.get(labelName.identifier)
                     val thisReceiver = when {
-                        thisReceivers.isNullOrEmpty() -> declarationDescriptor.extensionReceiverParameter
+                        thisReceivers.isNullOrEmpty() ->
+                            (declarationDescriptor as? FunctionDescriptor)?.extensionReceiverParameter
                         thisReceivers.size == 1 -> thisReceivers.single()
                         else -> {
                             BindingContextUtils.reportAmbiguousLabel(trace, targetLabelExpression, declarationsByLabel)
@@ -294,4 +304,4 @@ object LabelResolver {
             }
         }
     }
-}
+}

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ArgumentsGenerationUtils.kt

@@ -121,20 +121,7 @@ fun StatementGenerator.generateReceiver(defaultStartOffset: Int, defaultEndOffse
                             context.symbolTable.referenceValueParameter(receiverClassDescriptor.thisAsReceiverParameter)
                         )
                 }
-                is ContextClassReceiver -> {
+                is ContextClassReceiver -> loadContextReceiver(receiver, defaultStartOffset, defaultEndOffset)
-                    val receiverClassDescriptor = receiver.classDescriptor
-                    val thisAsReceiverParameter = receiverClassDescriptor.thisAsReceiverParameter
-                    val thisReceiver = IrGetValueImpl(
-                        defaultStartOffset, defaultEndOffset,
-                        thisAsReceiverParameter.type.toIrType(),
-                        context.symbolTable.referenceValue(thisAsReceiverParameter)
-                    )
-                    IrGetFieldImpl(
-                        defaultStartOffset, defaultEndOffset,
-                        context.additionalDescriptorStorage.getSyntheticField(receiver).symbol,
-                        irReceiverType, thisReceiver
-                    )
-                }
                 is ThisClassReceiver ->
                     generateThisOrSuperReceiver(receiver, receiver.classDescriptor)
                 is SuperCallReceiverValue ->
@@ -161,6 +148,26 @@ fun StatementGenerator.generateReceiver(defaultStartOffset: Int, defaultEndOffse
     }
 }
 
+fun StatementGenerator.loadContextReceiver(
+    receiver: ContextClassReceiver,
+    defaultStartOffset: Int, defaultEndOffset: Int,
+): IrGetFieldImpl {
+    val receiverClassDescriptor = receiver.classDescriptor
+    val thisAsReceiverParameter = receiverClassDescriptor.thisAsReceiverParameter
+    val thisReceiver = IrGetValueImpl(
+        defaultStartOffset, defaultEndOffset,
+        thisAsReceiverParameter.type.toIrType(),
+        context.symbolTable.referenceValue(thisAsReceiverParameter)
+    )
+
+    return IrGetFieldImpl(
+        defaultStartOffset, defaultEndOffset,
+        context.additionalDescriptorStorage.getSyntheticField(receiver).symbol,
+        receiver.type.toIrType(), thisReceiver
+    )
+}
+
+
 fun StatementGenerator.generateSingletonReference(
     descriptor: ClassDescriptor,
     startOffset: Int,

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/StatementGenerator.kt

@@ -17,10 +17,7 @@
 package org.jetbrains.kotlin.psi2ir.generators
 
 import org.jetbrains.kotlin.backend.common.BackendException
-import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.VariableDescriptorWithAccessors
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.Scope
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
@@ -46,6 +43,7 @@ import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
 import org.jetbrains.kotlin.resolve.calls.tasks.isDynamic
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
+import org.jetbrains.kotlin.resolve.scopes.receivers.ContextClassReceiver
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -415,23 +413,26 @@ class StatementGenerator(
 
     override fun visitThisExpression(expression: KtThisExpression, data: Nothing?): IrExpression {
         val referenceTarget = getOrFail(BindingContext.REFERENCE_TARGET, expression.instanceReference) { "No reference target for this" }
+        val receiverParameter =
+            getOrFail<KtReferenceExpression, ReceiverParameterDescriptor>(
+                BindingContext.THIS_REFERENCE_TARGET, expression.instanceReference
+            ) { "No reference target for this" }
         val startOffset = expression.startOffsetSkippingComments
         val endOffset = expression.endOffset
         return when (referenceTarget) {
             is ClassDescriptor ->
-                generateThisReceiver(startOffset, endOffset, referenceTarget.thisAsReceiverParameter.type, referenceTarget)
+                when (receiverParameter.value) {
-
+                    is ContextClassReceiver -> loadContextReceiver(receiverParameter.value as ContextClassReceiver, startOffset, endOffset)
+                    else -> generateThisReceiver(
+                        startOffset, endOffset, referenceTarget.thisAsReceiverParameter.type, referenceTarget
+                    )
+                }
             is CallableDescriptor -> {
-                val resolvedCall = getResolvedCall(expression)
+                val receiverType = receiverParameter.type.toIrType()
-                val receivers = listOfNotNull(referenceTarget.extensionReceiverParameter) + referenceTarget.contextReceiverParameters
-                val receiver = receivers.find {
-                    it == resolvedCall?.candidateDescriptor
-                } ?: referenceTarget.extensionReceiverParameter ?: error("No receiver: $referenceTarget")
-                val receiverType = receiver.type.toIrType()
                 IrGetValueImpl(
                     startOffset, endOffset,
                     receiverType,
-                    context.symbolTable.referenceValueParameter(receiver)
+                    context.symbolTable.referenceValueParameter(receiverParameter)
                 )
             }
 
@@ -513,4 +514,4 @@ abstract class StatementGeneratorExtension(val statementGenerator: StatementGene
     fun KtExpression.genStmt() = statementGenerator.generateStatement(this)
     fun KotlinType.toIrType() = with(statementGenerator) { toIrType() }
     fun translateType(kotlinType: KotlinType) = kotlinType.toIrType()
-}
+}

compiler/testData/codegen/box/extensionFunctions/contextReceivers/useFromAnotherModule.kt

@@ -22,11 +22,16 @@ class Foo {
     fun four(dummy: Any?) = this@Int
 }
 
+context(Int)
+class Bar {
+    fun five() = this@Int
+}
+
 // MODULE: main(lib)
 // FILE: B.kt
 
 fun box(): String {
     return with(1) {
-        if (a.one(null) + a.two + a.Foo().three + a.Foo().four(null) == 4) "OK" else "fail"
+        if (a.one(null) + a.two + a.Foo().three + a.Foo().four(null) + a.Bar().five() == 5) "OK" else "fail"
     }
 }

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

@@ -15,7 +15,7 @@ class B : A() {
     inner class C {
         fun g() {
             super@B.f()
-            <!DEBUG_INFO_MISSING_UNRESOLVED!>super<!><!UNRESOLVED_REFERENCE!>@Context<!>.<!DEBUG_INFO_MISSING_UNRESOLVED!>h<!>()
+            <!SUPERCLASS_NOT_ACCESSIBLE_FROM_INTERFACE!>super@Context<!>.<!UNRESOLVED_REFERENCE!>h<!>()
         }
     }
-}
+}

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

@@ -6,7 +6,7 @@ class A {
 }
 
 context(A) class B {
-    val prop = x + this<!UNRESOLVED_REFERENCE!>@A<!>.<!DEBUG_INFO_MISSING_UNRESOLVED!>x<!>
+    val prop = x + this@A.x
 
-    fun f() = x + this<!UNRESOLVED_REFERENCE!>@A<!>.<!DEBUG_INFO_MISSING_UNRESOLVED!>x<!>
+    fun f() = x + this@A.x
-}
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/thisWithReceiverLabelsClasses.txt

@@ -10,9 +10,9 @@ public final class A {
 
 context(A) public final class B {
     public constructor B()
-    public final val prop: [Error type: Not found recorded type for x + this@A.x]
+    public final val prop: kotlin.Int = 2
     public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
-    public final fun f(): [Error type: Return type for function cannot be resolved]
+    public final fun f(): kotlin.Int
     public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }