[FE] Support context receivers on classes

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -10582,6 +10582,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("genericOuterClass.kt")
+                public void testGenericOuterClass() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/genericOuterClass.kt");
+                }
+
                 @Test
                 @TestMetadata("insideDeclaration.kt")
                 public void testInsideDeclaration() throws Exception {
@@ -10606,6 +10612,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/noExplicitReceiver.kt");
                 }
 
+                @Test
+                @TestMetadata("outerClass.kt")
+                public void testOuterClass() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/outerClass.kt");
+                }
+
                 @Test
                 @TestMetadata("plusMatrix.kt")
                 public void testPlusMatrix() throws Exception {

compiler/frontend/src/org/jetbrains/kotlin/psi/synthetics/SyntheticClassOrObjectDescriptor.kt

@@ -167,6 +167,7 @@ class SyntheticClassOrObjectDescriptor(
         override fun getPrimaryConstructorModifierList(): KtModifierList? = null
         override fun getPrimaryConstructorParameters(): List<KtParameter> = emptyList()
         override fun getSecondaryConstructors(): List<KtSecondaryConstructor> = emptyList()
+        override fun getContextReceiverTypeReferences(): List<KtTypeReference> = emptyList()
 
         override fun getPsiOrParent() = _parent.psiOrParent
         override fun getParent() = _parent.psiOrParent

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/ClassResolutionScopesSupport.kt

@@ -25,8 +25,6 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.getAllSuperclassesWithoutAny
 import org.jetbrains.kotlin.resolve.scopes.*
 import org.jetbrains.kotlin.resolve.scopes.utils.ErrorLexicalScope
 import org.jetbrains.kotlin.storage.StorageManager
-import org.jetbrains.kotlin.utils.addIfNotNull
-import java.util.*
 
 class ClassResolutionScopesSupport(
     private val classDescriptor: ClassDescriptor,
@@ -69,7 +67,7 @@ class ClassResolutionScopesSupport(
             scopeWithGenerics,
             classDescriptor,
             true,
-            listOf(classDescriptor.thisAsReceiverParameter),
+            classDescriptor.contextReceivers + classDescriptor.thisAsReceiverParameter,
             LexicalScopeKind.CLASS_MEMBER_SCOPE
         )
     }

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.descriptors.impl.ClassDescriptorBase;
 import org.jetbrains.kotlin.descriptors.impl.FunctionDescriptorImpl;
+import org.jetbrains.kotlin.descriptors.impl.ReceiverParameterDescriptorImpl;
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory0;
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation;
 import org.jetbrains.kotlin.lexer.KtTokens;
@@ -39,6 +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.source.KotlinSourceElementKt;
 import org.jetbrains.kotlin.storage.MemoizedFunctionToNotNull;
 import org.jetbrains.kotlin.storage.NotNullLazyValue;
@@ -51,6 +53,7 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import static kotlin.collections.CollectionsKt.firstOrNull;
 import static org.jetbrains.kotlin.descriptors.DescriptorVisibilities.PRIVATE;
@@ -104,6 +107,8 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
 
     private final NotNullLazyValue<Collection<ClassDescriptor>> sealedSubclasses;
 
+    private final NotNullLazyValue<List<ReceiverParameterDescriptor>> contextReceivers;
+
     public LazyClassDescriptor(
             @NotNull LazyClassContext c,
             @NotNull DeclarationDescriptor containingDeclaration,
@@ -294,9 +299,23 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
             storageManager.createLazyValue(() -> {
                 if (getModality() == Modality.SEALED) {
                     return c.getSealedClassInheritorsProvider().computeSealedSubclasses(this, freedomForSealedInterfacesSupported);
-                } else {
+                }
+                else {
                     return Collections.emptyList();
                 }
+            });
+        this.contextReceivers = storageManager.createLazyValue(() -> {
+            if (classOrObject == null) {
+                return CollectionsKt.emptyList();
+            }
+            return classOrObject.getContextReceiverTypeReferences().stream().map(typeReference -> {
+                KotlinType kotlinType = c.getTypeResolver().resolveType(getScopeForClassHeaderResolution(), typeReference, c.getTrace(), true);
+                return new ReceiverParameterDescriptorImpl(
+                        this,
+                        new ExtensionClassReceiver(this, kotlinType, null),
+                        Annotations.Companion.getEMPTY()
+                );
+            }).collect(Collectors.toList());
         });
     }
 
@@ -432,6 +451,12 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
         return ((LazyClassMemberScope) getUnsubstitutedMemberScope()).getConstructors();
     }
 
+    @NotNull
+    @Override
+    public List<ReceiverParameterDescriptor> getContextReceivers() {
+        return contextReceivers.invoke();
+    }
+
     @Override
     public ClassConstructorDescriptor getUnsubstitutedPrimaryConstructor() {
         return ((LazyClassMemberScope) getUnsubstitutedMemberScope()).getPrimaryConstructor();
@@ -642,6 +667,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
         ForceResolveUtil.forceResolveAllContents(getDescriptorsForExtraCompanionObjects());
         ForceResolveUtil.forceResolveAllContents(getUnsubstitutedMemberScope());
         ForceResolveUtil.forceResolveAllContents(getTypeConstructor());
+        ForceResolveUtil.forceResolveAllContents(this.getContextReceivers());
     }
 
     // Note: headers of member classes' members are not resolved
@@ -671,6 +697,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
         }
         getUnsubstitutedPrimaryConstructor();
         getVisibility();
+        getContextReceivers();
     }
 
     @NotNull

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/descriptors/IrBasedDescriptors.kt

@@ -608,6 +608,10 @@ open class IrBasedClassDescriptor(owner: IrClass) : ClassDescriptor, IrBasedDecl
 
     override fun getThisAsReceiverParameter() = owner.thisReceiver?.toIrBasedDescriptor() as ReceiverParameterDescriptor
 
+    override fun getContextReceivers(): List<ReceiverParameterDescriptor> {
+        TODO("Not yet implemented")
+    }
+
     override fun getUnsubstitutedPrimaryConstructor() =
         owner.declarations.filterIsInstance<IrConstructor>().singleOrNull { it.isPrimary }?.toIrBasedDescriptor()
 
@@ -734,6 +738,10 @@ open class IrBasedEnumEntryDescriptor(owner: IrEnumEntry) : ClassDescriptor, IrB
 
     override fun getThisAsReceiverParameter() = (owner.parent as IrClass).toIrBasedDescriptor().thisAsReceiverParameter
 
+    override fun getContextReceivers(): List<ReceiverParameterDescriptor> {
+        TODO("Not yet implemented")
+    }
+
     override fun getUnsubstitutedPrimaryConstructor(): ClassConstructorDescriptor? {
         TODO("not implemented")
     }

compiler/psi/src/org/jetbrains/kotlin/psi/KtClassOrObject.kt

@@ -25,6 +25,7 @@ import com.intellij.psi.impl.CheckUtil
 import com.intellij.psi.stubs.IStubElementType
 import com.intellij.psi.tree.TokenSet
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -186,6 +187,23 @@ abstract class KtClassOrObject :
         parts.reverse()
         return parts.joinToString(separator = ".")
     }
+
+    override fun getContextReceiverTypeReferences(): List<KtTypeReference> {
+        val stub = stub
+        if (stub != null) {
+            return getStubOrPsiChildrenAsList(KtStubElementTypes.TYPE_REFERENCE)
+        }
+        var node = node.firstChildNode
+        while (node != null) {
+            val tt = node.elementType
+            if (tt === KtNodeTypes.CONTEXT_RECEIVER) {
+                val contextReceiver = node.psi as KtContextReceiver
+                return contextReceiver.typeReferences()
+            }
+            node = node.treeNext
+        }
+        return emptyList()
+    }
 }
 
 

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

@@ -58,6 +58,10 @@ public interface KtPureClassOrObject extends KtPureElement, KtDeclarationContain
     @ReadOnly
     List<KtSecondaryConstructor> getSecondaryConstructors();
 
+    @NotNull
+    @ReadOnly
+    List<KtTypeReference> getContextReceiverTypeReferences();
+
     @Nullable
     KtClassBody getBody();
 }

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

@@ -0,0 +1,13 @@
+context(T) class A<T>
+
+context(Collection<P>) class B<P>
+
+fun Int.foo() {
+    A<Int>()
+    A<String>()
+}
+
+fun Collection<Int>.bar() {
+    B<Int>()
+    B<String>()
+}

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

@@ -0,0 +1,13 @@
+context(T) class A<T>
+
+context(Collection<P>) class B<P>
+
+fun Int.foo() {
+    A<Int>()
+    <!UNRESOLVED_REFERENCE_WRONG_RECEIVER!>A<!><String>()
+}
+
+fun Collection<Int>.bar() {
+    B<Int>()
+    <!UNRESOLVED_REFERENCE_WRONG_RECEIVER!>B<!><String>()
+}

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

@@ -0,0 +1,18 @@
+package
+
+public fun kotlin.collections.Collection<kotlin.Int>.bar(): kotlin.Unit
+public fun kotlin.Int.foo(): kotlin.Unit
+
+public final class A</*0*/ T> {
+    public constructor A</*0*/ T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class B</*0*/ P> {
+    public constructor B</*0*/ P>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

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

@@ -0,0 +1,17 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+class Outer {
+    val x: Int = 1
+}
+
+context(Outer)
+class Inner(arg: Any) {
+    fun bar() = <!UNRESOLVED_REFERENCE!>x<!>
+}
+
+fun f(outer: Outer) {
+    Inner(1)
+    with(outer) {
+        Inner(3)
+    }
+}

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

@@ -0,0 +1,17 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+class Outer {
+    val x: Int = 1
+}
+
+context(Outer)
+class Inner(arg: Any) {
+    fun bar() = x
+}
+
+fun f(outer: Outer) {
+    <!NO_CONTEXT_RECEIVER!>Inner(1)<!>
+    with(outer) {
+        Inner(3)
+    }
+}

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

@@ -0,0 +1,19 @@
+package
+
+public fun f(/*0*/ outer: Outer): kotlin.Unit
+
+public final class Inner {
+    public constructor Inner(/*0*/ arg: kotlin.Any)
+    public final fun bar(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Outer {
+    public constructor Outer()
+    public final val x: kotlin.Int = 1
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -10588,6 +10588,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("genericOuterClass.kt")
+                public void testGenericOuterClass() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/genericOuterClass.kt");
+                }
+
                 @Test
                 @TestMetadata("insideDeclaration.kt")
                 public void testInsideDeclaration() throws Exception {
@@ -10612,6 +10618,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/noExplicitReceiver.kt");
                 }
 
+                @Test
+                @TestMetadata("outerClass.kt")
+                public void testOuterClass() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/outerClass.kt");
+                }
+
                 @Test
                 @TestMetadata("plusMatrix.kt")
                 public void testPlusMatrix() throws Exception {

core/descriptors/src/org/jetbrains/kotlin/descriptors/ClassDescriptor.java

@@ -77,6 +77,10 @@ public interface ClassDescriptor extends ClassifierDescriptorWithTypeParameters,
     @NotNull
     ReceiverParameterDescriptor getThisAsReceiverParameter();
 
+    @NotNull
+    @ReadOnly
+    List<ReceiverParameterDescriptor> getContextReceivers();
+
     @Nullable
     ClassConstructorDescriptor getUnsubstitutedPrimaryConstructor();
 

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/AbstractClassDescriptor.java

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.storage.StorageManager;
 import org.jetbrains.kotlin.types.*;
 import org.jetbrains.kotlin.types.checker.KotlinTypeRefiner;
 
+import java.util.Collections;
 import java.util.List;
 
 public abstract class AbstractClassDescriptor extends ModuleAwareClassDescriptor {
@@ -114,6 +115,12 @@ public abstract class AbstractClassDescriptor extends ModuleAwareClassDescriptor
         return thisAsReceiverParameter.invoke();
     }
 
+    @NotNull
+    @Override
+    public List<ReceiverParameterDescriptor> getContextReceivers() {
+        return Collections.emptyList();
+    }
+
     @NotNull
     @Override
     public MemberScope getMemberScope(@NotNull List<? extends TypeProjection> typeArguments, @NotNull KotlinTypeRefiner kotlinTypeRefiner) {

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/ClassConstructorDescriptorImpl.java

@@ -70,7 +70,7 @@ public class ClassConstructorDescriptorImpl extends FunctionDescriptorImpl imple
             @NotNull List<TypeParameterDescriptor> typeParameterDescriptors
     ) {
         super.initialize(
-                null, calculateDispatchReceiverParameter(), Collections.<ReceiverParameterDescriptor>emptyList(),
+                null, calculateDispatchReceiverParameter(), calculateAdditionalReceiverParameters(),
                 typeParameterDescriptors,
                 unsubstitutedValueParameters, null,
                 Modality.FINAL, visibility);
@@ -97,6 +97,15 @@ public class ClassConstructorDescriptorImpl extends FunctionDescriptorImpl imple
         return null;
     }
 
+    @NotNull
+    List<ReceiverParameterDescriptor> calculateAdditionalReceiverParameters() {
+        ClassDescriptor classDescriptor = getContainingDeclaration();
+        if (!classDescriptor.getContextReceivers().isEmpty()) {
+            return classDescriptor.getContextReceivers();
+        }
+        return Collections.emptyList();
+    }
+
     @NotNull
     @Override
     public ClassDescriptor getContainingDeclaration() {

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/LazySubstitutingClassDescriptor.java

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.types.checker.KotlinTypeRefiner;
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 
 public class LazySubstitutingClassDescriptor extends ModuleAwareClassDescriptor {
@@ -156,6 +157,12 @@ public class LazySubstitutingClassDescriptor extends ModuleAwareClassDescriptor
         throw new UnsupportedOperationException(); // TODO
     }
 
+    @NotNull
+    @Override
+    public List<ReceiverParameterDescriptor> getContextReceivers() {
+        return Collections.emptyList();
+    }
+
     @NotNull
     @Override
     public Collection<ClassConstructorDescriptor> getConstructors() {

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.resolve.scopes.receivers
+
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.types.KotlinType
+
+class ExtensionClassReceiver(
+    val classDescriptor: ClassDescriptor,
+    receiverType: KotlinType,
+    original: ReceiverValue?
+): AbstractReceiverValue(receiverType, original), ImplicitReceiver {
+    override val declarationDescriptor: DeclarationDescriptor
+        get() = classDescriptor
+
+    override fun replaceType(newType: KotlinType): ReceiverValue = ExtensionClassReceiver(classDescriptor, newType, original)
+
+    override fun toString(): String = "$type: Ext { $classDescriptor }"
+}