[PSI, FE] Review fixes

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

@@ -48,7 +48,6 @@ import org.jetbrains.kotlin.util.slicedMap.*;
 
 import java.util.Collection;
 import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.Map;
 
 import static org.jetbrains.kotlin.util.slicedMap.RewritePolicy.DO_NOTHING;
@@ -264,7 +263,7 @@ public interface BindingContext {
                     .setFurtherLookupSlices(DECLARATIONS_TO_DESCRIPTORS)
                     .build();
 
-    WritableSlice<DeclarationDescriptor, LinkedHashMap<ReceiverParameterDescriptor, String>> DESCRIPTOR_TO_NAMED_RECEIVERS = Slices.createSimpleSlice();
+    WritableSlice<DeclarationDescriptor, Map<ReceiverParameterDescriptor, String>> DESCRIPTOR_TO_NAMED_RECEIVERS = Slices.createSimpleSlice();
     WritableSlice<KtReferenceExpression, PsiElement> LABEL_TARGET = Slices.createSimpleSlice();
     WritableSlice<KtReferenceExpression, Collection<? extends PsiElement>> AMBIGUOUS_LABEL_TARGET = Slices.createSimpleSlice();
     WritableSlice<ValueParameterDescriptor, PropertyDescriptor> VALUE_PARAMETER_AS_PROPERTY = Slices.createSimpleSlice();

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/smartcasts/IdentifierInfo.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.resolve.calls.util.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.util.isSafeCall
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue.Kind.OTHER
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue.Kind.STABLE_VALUE
+import org.jetbrains.kotlin.resolve.scopes.receivers.ContextReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver
@@ -180,7 +181,12 @@ internal fun getIdForStableIdentifier(
 
         is KtThisExpression -> {
             val declarationDescriptor = bindingContext.get(BindingContext.REFERENCE_TARGET, expression.instanceReference)
-            getIdForThisReceiver(declarationDescriptor, bindingContext, expression.getLabelName())
+            val labelName = expression.getLabelName()
+            if (labelName == null) {
+                getIdForThisReceiver(declarationDescriptor)
+            } else {
+                getIdForThisReceiver(declarationDescriptor, bindingContext, labelName)
+            }
         }
 
         is KtPostfixExpression -> {
@@ -229,7 +235,7 @@ private fun getIdForSimpleNameExpression(
             if (implicitReceiver == null) {
                 selectorInfo
             } else {
-                val receiverInfo = getIdForImplicitReceiver(implicitReceiver, simpleNameExpression, bindingContext)
+                val receiverInfo = getIdForImplicitReceiver(implicitReceiver, simpleNameExpression)
 
                 if (receiverInfo == null) {
                     selectorInfo
@@ -255,9 +261,11 @@ private fun getIdForSimpleNameExpression(
     }
 }
 
-private fun getIdForImplicitReceiver(receiverValue: ReceiverValue?, expression: KtExpression?, bindingContext: BindingContext) =
+private fun getIdForImplicitReceiver(receiverValue: ReceiverValue?, expression: KtExpression?) =
     when (receiverValue) {
-        is ImplicitReceiver -> getIdForThisReceiver(receiverValue.declarationDescriptor, bindingContext)
+        is ContextReceiver -> IdentifierInfo.Receiver(receiverValue)
+
+        is ImplicitReceiver -> getIdForThisReceiver(receiverValue.declarationDescriptor)
 
         is TransientReceiver ->
             throw AssertionError("Transient receiver is implicit for an explicit expression: $expression. Receiver: $receiverValue")
@@ -265,16 +273,24 @@ private fun getIdForImplicitReceiver(receiverValue: ReceiverValue?, expression:
         else -> null
     }
 
-private fun getIdForThisReceiver(
+private fun getIdForThisReceiver(descriptorOfThisReceiver: DeclarationDescriptor?) = when (descriptorOfThisReceiver) {
-    descriptorOfThisReceiver: DeclarationDescriptor?,
+    is CallableDescriptor -> {
-    bindingContext: BindingContext,
+        val receiverParameter = descriptorOfThisReceiver.extensionReceiverParameter
-    labelName: String? = null
+            ?: error("'This' refers to the callable member without a receiver parameter: $descriptorOfThisReceiver")
-) =
+        IdentifierInfo.Receiver(receiverParameter.value)
+    }
+
+    is ClassDescriptor -> IdentifierInfo.Receiver(descriptorOfThisReceiver.thisAsReceiverParameter.value)
+
+    else -> IdentifierInfo.NO
+}
+
+private fun getIdForThisReceiver(descriptorOfThisReceiver: DeclarationDescriptor?, bindingContext: BindingContext, labelName: String) =
     when (descriptorOfThisReceiver) {
         is CallableDescriptor -> {
             val receiverParameter = findReceiverByLabelOrGetDefault(
                 descriptorOfThisReceiver,
-                descriptorOfThisReceiver.extensionReceiverParameter ?: descriptorOfThisReceiver.contextReceiverParameters.firstOrNull(),
+                descriptorOfThisReceiver.extensionReceiverParameter,
                 bindingContext,
                 labelName
             )
@@ -298,7 +314,7 @@ private fun findReceiverByLabelOrGetDefault(
     descriptorOfThisReceiver: DeclarationDescriptor,
     default: ReceiverParameterDescriptor?,
     bindingContext: BindingContext,
-    labelName: String? = null
+    labelName: String
 ): ReceiverParameterDescriptor {
     val receiverToLabelMap = bindingContext.get(
         BindingContext.DESCRIPTOR_TO_NAMED_RECEIVERS,

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

@@ -40,7 +40,7 @@ import org.jetbrains.kotlin.resolve.lazy.declarations.ClassMemberDeclarationProv
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope;
 import org.jetbrains.kotlin.resolve.scopes.MemberScope;
 import org.jetbrains.kotlin.resolve.scopes.StaticScopeForKotlinEnum;
-import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionClassReceiver;
+import org.jetbrains.kotlin.resolve.scopes.receivers.ContextClassReceiver;
 import org.jetbrains.kotlin.resolve.source.KotlinSourceElementKt;
 import org.jetbrains.kotlin.storage.MemoizedFunctionToNotNull;
 import org.jetbrains.kotlin.storage.NotNullLazyValue;
@@ -313,7 +313,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
                                 c.getTypeResolver().resolveType(getScopeForClassHeaderResolution(), typeReference, c.getTrace(), true);
                         return new ReceiverParameterDescriptorImpl(
                                 this,
-                                new ExtensionClassReceiver(this, kotlinType, null),
+                                new ContextClassReceiver(this, kotlinType, null),
                                 Annotations.Companion.getEMPTY()
                         );
             }).collect(Collectors.toList());

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

@@ -93,7 +93,7 @@ fun StatementGenerator.generateReceiver(defaultStartOffset: Int, defaultEndOffse
                             context.symbolTable.referenceValueParameter(receiverClassDescriptor.thisAsReceiverParameter)
                         )
                 }
-                is ExtensionClassReceiver -> {
+                is ContextClassReceiver -> {
                     val receiverClassDescriptor = receiver.classDescriptor
                     val thisAsReceiverParameter = receiverClassDescriptor.thisAsReceiverParameter
                     val thisReceiver = IrGetValueImpl(
@@ -226,7 +226,7 @@ fun StatementGenerator.generateCallReceiver(
             dispatchReceiverValue = generateReceiverOrNull(ktDefaultElement, dispatchReceiver)
             extensionReceiverValue = generateReceiverOrNull(ktDefaultElement, extensionReceiver)
             contextReceiverValues = if (ktDefaultElement is KtConstructorDelegationCall) contextReceivers.mapNotNull {
-                generateReceiverOrNull(ktDefaultElement, it as ExtensionClassReceiver)
+                generateReceiverOrNull(ktDefaultElement, it as ContextClassReceiver)
             }
             else contextReceivers.mapNotNull { generateReceiverOrNull(ktDefaultElement, it) }
         }

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

@@ -426,7 +426,7 @@ class StatementGenerator(
                 val receivers = listOfNotNull(referenceTarget.extensionReceiverParameter) + referenceTarget.contextReceiverParameters
                 val receiver = receivers.find {
                     it == resolvedCall?.candidateDescriptor
-                } ?: referenceTarget.extensionReceiverParameter ?: TODO("No receiver: $referenceTarget")
+                } ?: referenceTarget.extensionReceiverParameter ?: error("No receiver: $referenceTarget")
                 val receiverType = receiver.type.toIrType()
                 IrGetValueImpl(
                     startOffset, endOffset,

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinParsing.java

@@ -618,7 +618,7 @@ public class KotlinParsing extends AbstractKotlinParsing {
 
     /*
      * contextReceiverList
-     *   : "context" "(" (label? typeReference{","})+ ")"
+     *   : "context" "(" (contextReceiver{","})+ ")"
      */
     public void parseContextReceiverList() {
         assert _at(CONTEXT_KEYWORD);

compiler/psi/src/org/jetbrains/kotlin/psi/KtNamedFunction.java

@@ -183,37 +183,6 @@ public class KtNamedFunction extends KtTypeParameterListOwnerStub<KotlinFunction
         return getReceiverTypeRefByTree();
     }
 
-    @NotNull
-    @Override
-    public List<KtContextReceiver> getContextReceivers() {
-        KotlinFunctionStub stub = getStub();
-        if (stub != null) {
-            List<KtContextReceiverList> childContextReceivers = getStubOrPsiChildrenAsList(KtStubElementTypes.CONTEXT_RECEIVER_LIST);
-            if (!childContextReceivers.isEmpty()) {
-                return childContextReceivers.get(0).contextReceivers();
-            }
-            else {
-                return Collections.emptyList();
-            }
-        }
-        return getContextReceiverTypeRefsByTree();
-    }
-
-    @NotNull
-    private List<KtContextReceiver> getContextReceiverTypeRefsByTree() {
-        PsiElement child = getFirstChild();
-        while (child != null) {
-            IElementType tt = child.getNode().getElementType();
-            if (tt == KtTokens.LPAR || tt == KtTokens.COLON) break;
-            if (child instanceof KtContextReceiverList) {
-                return ((KtContextReceiverList) child).contextReceivers();
-            }
-            child = child.getNextSibling();
-        }
-
-        return Collections.emptyList();
-    }
-
     @Nullable
     private KtTypeReference getReceiverTypeRefByTree() {
         PsiElement child = getFirstChild();
@@ -229,6 +198,18 @@ public class KtNamedFunction extends KtTypeParameterListOwnerStub<KotlinFunction
         return null;
     }
 
+    @NotNull
+    @Override
+    public List<KtContextReceiver> getContextReceivers() {
+        KtContextReceiverList list = getContextReceiverList();
+        return list != null ? list.contextReceivers() : Collections.emptyList();
+    }
+
+    @Nullable
+    private KtContextReceiverList getContextReceiverList() {
+        return getStubOrPsiChild(KtStubElementTypes.CONTEXT_RECEIVER_LIST);
+    }
+
     @Override
     @Nullable
     public KtTypeReference getTypeReference() {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt

@@ -804,6 +804,24 @@ internal object ErrorDescriptorResolutionPart : ResolutionPart() {
 }
 
 internal object CheckContextReceiversResolutionPart : ResolutionPart() {
+    override fun ResolutionCandidate.process(workIndex: Int) {
+        val parentLexicalScopes = scopeTower.lexicalScope.parentsWithSelf.filterIsInstance<LexicalScope>()
+        val implicitReceiversGroups = mutableListOf<List<ReceiverValueWithSmartCastInfo>>()
+        for (scope in parentLexicalScopes) {
+            scopeTower.getImplicitReceiver(scope)?.let { implicitReceiversGroups.add(listOf(it)) }
+            val contextReceiversGroup = scopeTower.getContextReceivers(scope)
+            if (contextReceiversGroup.isNotEmpty()) {
+                implicitReceiversGroups.add(contextReceiversGroup)
+            }
+        }
+        val contextReceiversArguments = mutableListOf<SimpleKotlinCallArgument>()
+        for (candidateContextReceiverParameter in candidateDescriptor.contextReceiverParameters) {
+            val contextReceiverArgument = findContextReceiver(implicitReceiversGroups, candidateContextReceiverParameter) ?: return
+            contextReceiversArguments.add(contextReceiverArgument)
+        }
+        resolvedCall.contextReceiversArguments = contextReceiversArguments
+    }
+
     private data class ApplicableArgumentWithConstraint(
         val argument: SimpleKotlinCallArgument,
         val argumentType: UnwrappedType,
@@ -844,21 +862,4 @@ internal object CheckContextReceiversResolutionPart : ResolutionPart() {
         diagnosticsFromResolutionParts.add(NoContextReceiver(candidateContextReceiverParameter))
         return null
     }
-
-    override fun ResolutionCandidate.process(workIndex: Int) {
-        val parentLexicalScopes = scopeTower.lexicalScope.parentsWithSelf.filterIsInstance<LexicalScope>()
-        val implicitReceiversGroups = mutableListOf<List<ReceiverValueWithSmartCastInfo>>()
-        for (scope in parentLexicalScopes) {
-            scopeTower.getImplicitReceiver(scope)?.let { implicitReceiversGroups.add(listOf(it)) }
-            val contextReceiversGroup = scopeTower.getContextReceivers(scope)
-            if (contextReceiversGroup.isNotEmpty()) {
-                implicitReceiversGroups.add(contextReceiversGroup)
-            }
-        }
-        val contextReceiversArguments = mutableListOf<SimpleKotlinCallArgument>()
-        for (candidateContextReceiverParameter in candidateDescriptor.contextReceiverParameters) {
-            contextReceiversArguments.add(findContextReceiver(implicitReceiversGroups, candidateContextReceiverParameter) ?: return)
-        }
-        resolvedCall.contextReceiversArguments = contextReceiversArguments
-    }
 }

core/descriptors/src/org/jetbrains/kotlin/resolve/scopes/receivers/ContextClassReceiver.kt

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.types.KotlinType
 
-class ExtensionClassReceiver(
+class ContextClassReceiver(
     val classDescriptor: ClassDescriptor,
     receiverType: KotlinType,
     original: ReceiverValue?
@@ -17,7 +17,7 @@ class ExtensionClassReceiver(
     override val declarationDescriptor: DeclarationDescriptor
         get() = classDescriptor
 
-    override fun replaceType(newType: KotlinType): ReceiverValue = ExtensionClassReceiver(classDescriptor, newType, original)
+    override fun replaceType(newType: KotlinType): ReceiverValue = ContextClassReceiver(classDescriptor, newType, original)
 
-    override fun toString(): String = "$type: Ext { $classDescriptor }"
+    override fun toString(): String = "$type: Ctx { $classDescriptor }"
 }