[PSI2IR] Support context receivers on classes

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -15931,6 +15931,40 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
         }
     }
 
+    @Nested
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    public class ExtensionClasses {
+        @Test
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
+        }
+
+        @Test
+        @TestMetadata("constructors.kt")
+        public void testConstructors() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/constructors.kt");
+        }
+
+        @Test
+        @TestMetadata("genericCollection.kt")
+        public void testGenericCollection() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/genericCollection.kt");
+        }
+
+        @Test
+        @TestMetadata("generics.kt")
+        public void testGenerics() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
+        }
+
+        @Test
+        @TestMetadata("simple.kt")
+        public void testSimple() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/simple.kt");
+        }
+    }
+
     @Nested
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")

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

@@ -1269,7 +1269,6 @@ public class DescriptorResolver {
         });
     }
 
-    @NotNull
     public PropertyDescriptor resolvePrimaryConstructorParameterToAProperty(
             @NotNull ClassDescriptor classDescriptor,
             @NotNull ValueParameterDescriptor valueParameter,

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

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtConstructorDelegationCall
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.ValueArgument
@@ -92,6 +93,20 @@ fun StatementGenerator.generateReceiver(defaultStartOffset: Int, defaultEndOffse
                             context.symbolTable.referenceValueParameter(receiverClassDescriptor.thisAsReceiverParameter)
                         )
                 }
+                is ExtensionClassReceiver -> {
+                    val receiverClassDescriptor = receiver.classDescriptor
+                    val thisAsReceiverParameter = receiverClassDescriptor.thisAsReceiverParameter
+                    val thisReceiver = IrGetValueImpl(
+                        defaultStartOffset, defaultEndOffset,
+                        thisAsReceiverParameter.type.toIrType(),
+                        context.symbolTable.referenceValue(thisAsReceiverParameter)
+                    )
+                    IrGetFieldImpl(
+                        defaultStartOffset, defaultEndOffset,
+                        context.symbolTable.referenceField(context.additionalDescriptorStorage.getField(receiver)),
+                        irReceiverType, thisReceiver
+                    )
+                }
                 is ThisClassReceiver ->
                     generateThisOrSuperReceiver(receiver, receiver.classDescriptor)
                 is SuperCallReceiverValue ->
@@ -210,7 +225,10 @@ fun StatementGenerator.generateCallReceiver(
         else -> {
             dispatchReceiverValue = generateReceiverOrNull(ktDefaultElement, dispatchReceiver)
             extensionReceiverValue = generateReceiverOrNull(ktDefaultElement, extensionReceiver)
-            contextReceiverValues = contextReceivers.mapNotNull { generateReceiverOrNull(ktDefaultElement, it) }
+            contextReceiverValues = if (ktDefaultElement is KtConstructorDelegationCall) contextReceivers.mapNotNull {
+                generateReceiverOrNull(ktDefaultElement, it as ExtensionClassReceiver)
+            }
+            else contextReceivers.mapNotNull { generateReceiverOrNull(ktDefaultElement, it) }
         }
     }
 

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

@@ -18,7 +18,9 @@ package org.jetbrains.kotlin.psi2ir.generators
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.Scope
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
@@ -175,12 +177,15 @@ class BodyGenerator(
     fun getLoop(expression: KtExpression): IrLoop? =
         loopTable[expression]
 
-    fun generatePrimaryConstructorBody(ktClassOrObject: KtPureClassOrObject): IrBody {
+    fun generatePrimaryConstructorBody(ktClassOrObject: KtPureClassOrObject, irConstructor: IrConstructor): IrBody {
         val irBlockBody = context.irFactory.createBlockBody(ktClassOrObject.pureStartOffset, ktClassOrObject.pureEndOffset)
 
         generateSuperConstructorCall(irBlockBody, ktClassOrObject)
 
         val classDescriptor = (scopeOwner as ClassConstructorDescriptor).containingDeclaration
+        if (classDescriptor.contextReceivers.isNotEmpty()) {
+            generateSetAdditionalFieldForPrimaryConstructorBody(classDescriptor, irConstructor, irBlockBody)
+        }
         irBlockBody.statements.add(
             IrInstanceInitializerCallImpl(
                 ktClassOrObject.pureStartOffset, ktClassOrObject.pureEndOffset,
@@ -329,4 +334,28 @@ class BodyGenerator(
             pregenerateCall(constructorCall)
         )
 
+    private fun generateSetAdditionalFieldForPrimaryConstructorBody(
+        classDescriptor: ClassDescriptor,
+        irConstructor: IrConstructor,
+        irBlockBody: IrBlockBody
+    ) {
+        val thisAsReceiverParameter = classDescriptor.thisAsReceiverParameter
+        val receiver = IrGetValueImpl(
+            UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+            thisAsReceiverParameter.type.toIrType(),
+            context.symbolTable.referenceValue(thisAsReceiverParameter)
+        )
+        for ((index, receiverDescriptor) in classDescriptor.contextReceivers.withIndex()) {
+            val irValueParameter = irConstructor.valueParameters[index]
+            irBlockBody.statements.add(
+                IrSetFieldImpl(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                    context.symbolTable.referenceField(context.additionalDescriptorStorage.getField(receiverDescriptor.value)),
+                    receiver,
+                    IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, irValueParameter.type, irValueParameter.symbol),
+                    context.irBuiltIns.unitType
+                )
+            )
+        }
+    }
 }

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

@@ -20,6 +20,11 @@ import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
+import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.impl.FieldDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.PropertyDescriptorImpl
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrImplementingDelegateDescriptorImpl
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
@@ -35,6 +40,8 @@ import org.jetbrains.kotlin.ir.util.SYNTHETIC_OFFSET
 import org.jetbrains.kotlin.ir.util.createIrClassFromDescriptor
 import org.jetbrains.kotlin.ir.util.declareSimpleFunctionWithOverrides
 import org.jetbrains.kotlin.ir.util.properties
+import org.jetbrains.kotlin.ir.util.referenceFunction
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtDelegatedSuperTypeEntry
 import org.jetbrains.kotlin.psi.KtEnumEntry
@@ -47,6 +54,7 @@ import org.jetbrains.kotlin.psi.synthetics.SyntheticClassOrObjectDescriptor
 import org.jetbrains.kotlin.psi.synthetics.findClassDescriptor
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DelegationResolver
+import org.jetbrains.kotlin.resolve.DescriptorFactory
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.propertyIfAccessor
 import org.jetbrains.kotlin.resolve.descriptorUtil.setSingleOverridden
@@ -89,6 +97,8 @@ class ClassGenerator(
                 classDescriptor.thisAsReceiverParameter.type.toIrType()
             )
 
+            generateFieldsForAdditionalReceivers(irClass, classDescriptor)
+
             val irPrimaryConstructor = generatePrimaryConstructor(irClass, ktClassOrObject)
             if (irPrimaryConstructor != null) {
                 generateDeclarationsForPrimaryConstructorParameters(irClass, irPrimaryConstructor, ktClassOrObject)
@@ -423,6 +433,41 @@ class ClassGenerator(
         EnumClassMembersGenerator(declarationGenerator).generateSpecialMembers(irClass)
     }
 
+    private fun generateFieldsForAdditionalReceivers(irClass: IrClass, classDescriptor: ClassDescriptor) {
+        for ((fieldIndex, receiverDescriptor) in classDescriptor.contextReceivers.withIndex()) {
+            val descriptor = PropertyDescriptorImpl.create(
+                classDescriptor,
+                Annotations.EMPTY,
+                Modality.FINAL,
+                DescriptorVisibilities.DEFAULT_VISIBILITY,
+                false,
+                Name.identifier("additionalReceiverField$fieldIndex"),
+                CallableMemberDescriptor.Kind.DECLARATION,
+                SourceElement.NO_SOURCE,
+                false, false,
+                classDescriptor.isExpect,
+                false, false, false
+            )
+            descriptor.setType(
+                receiverDescriptor.type, emptyList(), DescriptorUtils.getDispatchReceiverParameterIfNeeded(classDescriptor), null,
+                emptyList()
+            )
+            descriptor.initialize(
+                null, null,
+                FieldDescriptorImpl(Annotations.EMPTY, descriptor),
+                null
+            )
+            context.additionalDescriptorStorage.put(receiverDescriptor.value, descriptor)
+            val irField = context.symbolTable.declareField(
+                UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                IrDeclarationOrigin.FIELD_FOR_CLASS_CONTEXT_RECEIVER,
+                descriptor, descriptor.type.toIrType(),
+                descriptor.visibility
+            )
+            irClass.addMember(irField)
+        }
+    }
+
     private fun generatePrimaryConstructor(irClass: IrClass, ktClassOrObject: KtPureClassOrObject): IrConstructor? {
         val classDescriptor = irClass.descriptor
         val primaryConstructorDescriptor = classDescriptor.unsubstitutedPrimaryConstructor ?: return null
@@ -459,7 +504,7 @@ class ClassGenerator(
         // generate real body declarations
         ktClassOrObject.body?.let { ktClassBody ->
             ktClassBody.declarations.mapNotNullTo(irClass.declarations) { ktDeclaration ->
-                declarationGenerator.generateClassMemberDeclaration(ktDeclaration, irClass)
+                declarationGenerator.generateClassMemberDeclaration(ktDeclaration, irClass, ktClassOrObject)
             }
         }
 

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

@@ -85,12 +85,16 @@ class DeclarationGenerator(override val context: GeneratorContext) : Generator {
         return generateClassOrObjectDeclaration(syntheticDeclaration)
     }
 
-    fun generateClassMemberDeclaration(ktDeclaration: KtDeclaration, irClass: IrClass): IrDeclaration? =
+    fun generateClassMemberDeclaration(
+        ktDeclaration: KtDeclaration,
+        irClass: IrClass,
+        ktClassOrObject: KtPureClassOrObject
+    ): IrDeclaration? =
         when (ktDeclaration) {
             is KtAnonymousInitializer ->
                 AnonymousInitializerGenerator(this).generateAnonymousInitializerDeclaration(ktDeclaration, irClass)
             is KtSecondaryConstructor ->
-                FunctionGenerator(this).generateSecondaryConstructor(ktDeclaration)
+                FunctionGenerator(this).generateSecondaryConstructor(ktDeclaration, ktClassOrObject)
             is KtEnumEntry ->
                 generateEnumEntryDeclaration(ktDeclaration)
             else ->

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

@@ -0,0 +1,29 @@
+/*
+ * 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.psi2ir.generators
+
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
+
+class DescriptorStorageForAdditionalReceivers {
+    private val fieldStorage: MutableMap<ReceiverValue, PropertyDescriptor> = mutableMapOf()
+    private val variableStorage: MutableMap<KtExpression, VariableDescriptor> = mutableMapOf()
+
+    fun put(receiverValue: ReceiverValue, descriptor: PropertyDescriptor) {
+        fieldStorage[receiverValue] = descriptor
+    }
+
+    fun put(expression: KtExpression, descriptor: VariableDescriptor) {
+        variableStorage[expression] = descriptor
+    }
+
+    fun getField(receiverValue: ReceiverValue) =
+        fieldStorage[receiverValue] ?: error("No field descriptor for receiver value $receiverValue")
+
+    fun getVariable(expression: KtExpression) = variableStorage[expression] ?: error("No variable descriptor for receiver $expression")
+}

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

@@ -234,19 +234,29 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
         primaryConstructorDescriptor: ClassConstructorDescriptor,
         ktClassOrObject: KtPureClassOrObject
     ): IrConstructor =
-        declareConstructor(ktClassOrObject, ktClassOrObject.primaryConstructor ?: ktClassOrObject, primaryConstructorDescriptor) {
+        declareConstructor(
+            ktClassOrObject,
+            ktClassOrObject.primaryConstructor ?: ktClassOrObject,
+            ktClassOrObject.contextReceiverTypeReferences,
+            primaryConstructorDescriptor
+        ) { irConstructor ->
             if (
                 primaryConstructorDescriptor.isExpect ||
                 primaryConstructorDescriptor.constructedClass.isEffectivelyExternal()
             )
                 null
             else
-                generatePrimaryConstructorBody(ktClassOrObject)
+                generatePrimaryConstructorBody(ktClassOrObject, irConstructor)
         }
 
-    fun generateSecondaryConstructor(ktConstructor: KtSecondaryConstructor): IrConstructor {
+    fun generateSecondaryConstructor(ktConstructor: KtSecondaryConstructor, ktClassOrObject: KtPureClassOrObject): IrConstructor {
         val constructorDescriptor = getOrFail(BindingContext.CONSTRUCTOR, ktConstructor) as ClassConstructorDescriptor
-        return declareConstructor(ktConstructor, ktConstructor, constructorDescriptor) {
+        return declareConstructor(
+            ktConstructor,
+            ktConstructor,
+            ktClassOrObject.contextReceiverTypeReferences,
+            constructorDescriptor
+        ) {
             when {
                 constructorDescriptor.constructedClass.isEffectivelyExternal() ->
                     null
@@ -263,8 +273,9 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
     private inline fun declareConstructor(
         ktConstructorElement: KtPureElement,
         ktParametersElement: KtPureElement,
+        ktContextReceiversElements: List<KtPureElement>,
         constructorDescriptor: ClassConstructorDescriptor,
-        generateBody: BodyGenerator.() -> IrBody?
+        generateBody: BodyGenerator.(IrConstructor) -> IrBody?
     ): IrConstructor {
         val startOffset = ktConstructorElement.getStartOffsetOfConstructorDeclarationKeywordOrNull() ?: ktConstructorElement.pureStartOffset
         val endOffset = ktConstructorElement.pureEndOffset
@@ -279,8 +290,8 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
                 metadata = DescriptorMetadataSource.Function(it.descriptor)
             }
         }.buildWithScope { irConstructor ->
-            generateValueParameterDeclarations(irConstructor, ktParametersElement, null, emptyList())
+            generateValueParameterDeclarations(irConstructor, ktParametersElement, null, ktContextReceiversElements)
-            irConstructor.body = createBodyGenerator(irConstructor.symbol).generateBody()
+            irConstructor.body = createBodyGenerator(irConstructor.symbol).generateBody(irConstructor)
             irConstructor.returnType = constructorDescriptor.returnType.toIrType()
         }
     }

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

@@ -71,6 +71,8 @@ class GeneratorContext private constructor(
     // TODO: inject a correct StorageManager instance, or store NotFoundClasses inside ModuleDescriptor
     val reflectionTypes = ReflectionTypes(moduleDescriptor, NotFoundClasses(LockBasedStorageManager.NO_LOCKS, moduleDescriptor))
 
+    val additionalDescriptorStorage: DescriptorStorageForAdditionalReceivers = DescriptorStorageForAdditionalReceivers()
+
     val samTypeApproximator = SamTypeApproximator(moduleDescriptor.builtIns, languageVersionSettings)
 
     fun createFileScopeContext(ktFile: KtFile): GeneratorContext {

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrDeclarationOrigin.kt

@@ -68,6 +68,7 @@ interface IrDeclarationOrigin {
     object SYNTHETIC_HELPER_FOR_ENUM_VALUES : IrDeclarationOriginImpl("SYNTHETIC_HELPER_FOR_ENUM_VALUES", isSynthetic = true)
     object FIELD_FOR_ENUM_VALUES : IrDeclarationOriginImpl("FIELD_FOR_ENUM_VALUES", isSynthetic = true)
     object FIELD_FOR_OBJECT_INSTANCE : IrDeclarationOriginImpl("FIELD_FOR_OBJECT_INSTANCE")
+    object FIELD_FOR_CLASS_CONTEXT_RECEIVER : IrDeclarationOriginImpl("FIELD_FOR_CLASS_CONTEXT_RECEIVER", isSynthetic = true)
 
     object ADAPTER_FOR_CALLABLE_REFERENCE : IrDeclarationOriginImpl("ADAPTER_FOR_CALLABLE_REFERENCE", isSynthetic = true)
     object ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE : IrDeclarationOriginImpl("ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE")

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrConstructorCallImpl.kt

@@ -40,7 +40,7 @@ class IrConstructorCallImpl(
             val constructorDescriptor = constructorSymbol.descriptor
             val classTypeParametersCount = constructorDescriptor.constructedClass.original.declaredTypeParameters.size
             val totalTypeParametersCount = constructorDescriptor.typeParameters.size
-            val valueParametersCount = constructorDescriptor.valueParameters.size
+            val valueParametersCount = constructorDescriptor.valueParameters.size + constructorDescriptor.contextReceiverParameters.size
             return IrConstructorCallImpl(
                 startOffset, endOffset,
                 type,

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrDelegatingConstructorCallImpl.kt

@@ -48,7 +48,7 @@ class IrDelegatingConstructorCallImpl(
             type: IrType,
             symbol: IrConstructorSymbol,
             typeArgumentsCount: Int = symbol.descriptor.typeParametersCount,
-            valueArgumentsCount: Int = symbol.descriptor.valueParameters.size
+            valueArgumentsCount: Int = symbol.descriptor.valueParameters.size + symbol.descriptor.contextReceiverParameters.size
         ) = IrDelegatingConstructorCallImpl(startOffset, endOffset, type, symbol, typeArgumentsCount, valueArgumentsCount)
 
         fun fromSymbolOwner(

compiler/testData/codegen/box/extensionClasses/constructors.kt

@@ -0,0 +1,32 @@
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+// WITH_RUNTIME
+
+class A(val ok: String)
+
+context(A)
+class B(oValue: Boolean = true, kValue: Boolean = true) {
+    var o: Boolean
+    var k: Boolean
+
+    init {
+        o = oValue
+        k = kValue
+    }
+
+    constructor(oValue: String, kValue: String) : this(oValue == "O", kValue == "K")
+
+    fun result() = if (o && k) ok else "fail"
+}
+
+fun box(): String {
+    val a = A("OK")
+    with (a) {
+        val results = listOf(
+            B(true, true).result(),
+            B("O", "K").result(),
+            B().result()
+        )
+        return if (results.all { it == "OK" }) "OK" else "fail"
+    }
+}

compiler/testData/codegen/box/extensionClasses/genericCollection.kt

@@ -0,0 +1,13 @@
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+// WITH_RUNTIME
+
+context(Collection<P>) class A<P> {
+    val result = if (!isEmpty()) "OK" else "fail"
+}
+
+fun box(): String {
+    with (listOf(1, 2, 3)) {
+        return A().result
+    }
+}

compiler/testData/codegen/box/extensionClasses/generics.kt

@@ -0,0 +1,10 @@
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+
+context(T) class B<T : CharSequence> {
+    val result = if (length == 2) "OK" else "fail"
+}
+
+fun box() = with("OK") {
+    B().result
+}

compiler/testData/codegen/box/extensionClasses/simple.kt

@@ -0,0 +1,15 @@
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_FIR: JVM_IR
+
+class A {
+    val ok = "OK"
+}
+
+context(A)
+class B {
+    fun result() = ok
+}
+
+fun box() = with(A()) {
+    B().result()
+}

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

@@ -15853,6 +15853,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
         }
     }
 
+    @Nested
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    public class ExtensionClasses {
+        @Test
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM, true);
+        }
+    }
+
     @Nested
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")

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

@@ -15931,6 +15931,40 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
         }
     }
 
+    @Nested
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    public class ExtensionClasses {
+        @Test
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
+        }
+
+        @Test
+        @TestMetadata("constructors.kt")
+        public void testConstructors() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/constructors.kt");
+        }
+
+        @Test
+        @TestMetadata("genericCollection.kt")
+        public void testGenericCollection() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/genericCollection.kt");
+        }
+
+        @Test
+        @TestMetadata("generics.kt")
+        public void testGenerics() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/generics.kt");
+        }
+
+        @Test
+        @TestMetadata("simple.kt")
+        public void testSimple() throws Exception {
+            runTest("compiler/testData/codegen/box/extensionClasses/simple.kt");
+        }
+    }
+
     @Nested
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -13041,6 +13041,19 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
         }
     }
 
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ExtensionClasses extends AbstractLightAnalysisModeTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, TargetBackend.JVM, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM, true);
+        }
+    }
+
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

compiler/visualizer/tests-gen/org/jetbrains/kotlin/visualizer/VisualizerBlackBoxTestGenerated.java

@@ -14614,6 +14614,20 @@ public class VisualizerBlackBoxTestGenerated extends AbstractVisualizerBlackBoxT
         }
     }
 
+    @Nested
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ExtensionClasses extends AbstractIrJsCodegenBoxES6Test {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest0(this::doTest, TargetBackend.JS_IR_ES6, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS_IR_ES6, true);
+        }
+    }
+
     @Nested
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/JsCodegenBoxTestGenerated.java

@@ -12267,6 +12267,19 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
         }
     }
 
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ExtensionClasses extends AbstractJsCodegenBoxTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest0(this::doTest, TargetBackend.JS, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS, true);
+        }
+    }
+
     @Nested
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsCodegenBoxTestGenerated.java

@@ -12309,6 +12309,20 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
         }
     }
 
+    @Nested
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ExtensionClasses extends AbstractIrJsCodegenBoxTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest0(this::doTest, TargetBackend.JS_IR, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS_IR, true);
+        }
+    }
+
     @Nested
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -10208,6 +10208,19 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
         }
     }
 
+    @TestMetadata("compiler/testData/codegen/box/extensionClasses")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ExtensionClasses extends AbstractIrCodegenBoxWasmTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest0(this::doTest, TargetBackend.WASM, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInExtensionClasses() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/extensionClasses"), Pattern.compile("^([^_](.+))\\.kt$"), null, TargetBackend.WASM, true);
+        }
+    }
+
     @TestMetadata("compiler/testData/codegen/box/extensionFunctions")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)