FIR2IR: Fix substitution work for context receivers

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

@@ -16752,12 +16752,24 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/simpleCall.kt");
             }
 
+            @Test
+            @TestMetadata("substitutedContextReceivers.kt")
+            public void testSubstitutedContextReceivers() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/substitutedContextReceivers.kt");
+            }
+
             @Test
             @TestMetadata("superClassAndSubClassWithContextReceiver.kt")
             public void testSuperClassAndSubClassWithContextReceiver() throws Exception {
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt");
             }
 
+            @Test
+            @TestMetadata("superClassAndSubClassWithContextReceiverSubstituted.kt")
+            public void testSuperClassAndSubClassWithContextReceiverSubstituted() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiverSubstituted.kt");
+            }
+
             @Test
             @TestMetadata("suspendContextualWithExtension.kt")
             public void testSuspendContextualWithExtension() throws Exception {

compiler/fir/java/src/org/jetbrains/kotlin/fir/scopes/jvm/JvmMappedScope.kt

@@ -130,6 +130,7 @@ class JvmMappedScope(
                 newReturnType = substitutor.substituteOrSelf(oldConstructor.returnTypeRef.coneType),
                 newParameterTypes = oldConstructor.valueParameters.map { substitutor.substituteOrSelf(it.returnTypeRef.coneType) },
                 newTypeParameters = null,
+                newContextReceiverTypes = emptyList(),
                 isExpect = false,
                 fakeOverrideSubstitution = null
             ).apply {

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirClassSubstitutionScope.kt

@@ -135,11 +135,16 @@ class FirClassSubstitutionScope(
             it.returnTypeRef.coneType.substitute(newSubstitutor)
         }
 
+        val newContextReceiverTypes = member.contextReceivers.map {
+            it.typeRef.coneType.substitute(newSubstitutor)
+        }
+
         if (newReceiverType == null &&
             newReturnType == null &&
             newParameterTypes.all { it == null } &&
             newTypeParameters === member.typeParameters &&
-            fakeOverrideSubstitution == null
+            fakeOverrideSubstitution == null &&
+            newContextReceiverTypes.all { it == null }
         ) {
             if (original.dispatchReceiverType?.substitute(substitutor) != null) {
                 return FirFakeOverrideGenerator.createSubstitutionOverrideFunction(
@@ -164,6 +169,7 @@ class FirClassSubstitutionScope(
             member,
             newDispatchReceiverType ?: dispatchReceiverTypeForSubstitutedMembers,
             newReceiverType,
+            newContextReceiverTypes,
             newReturnType,
             newParameterTypes,
             newTypeParameters as List<FirTypeParameter>,
@@ -188,7 +194,15 @@ class FirClassSubstitutionScope(
             it.returnTypeRef.coneType.substitute(newSubstitutor)
         }
 
-        if (newReturnType == null && newParameterTypes.all { it == null } && newTypeParameters === constructor.typeParameters) {
+        val newContextReceiverTypes = constructor.contextReceivers.map {
+            it.typeRef.coneType.substitute(newSubstitutor)
+        }
+
+        if (newReturnType == null &&
+            newParameterTypes.all { it == null } &&
+            newTypeParameters === constructor.typeParameters &&
+            newContextReceiverTypes.all { it == null }
+        ) {
             return original
         }
 
@@ -200,6 +214,7 @@ class FirClassSubstitutionScope(
             newDispatchReceiverType,
             newReturnType,
             newParameterTypes,
+            newContextReceiverTypes,
             newTypeParameters,
             makeExpect,
             fakeOverrideSubstitution
@@ -217,10 +232,15 @@ class FirClassSubstitutionScope(
             member, symbolForOverride
         )
 
+        val newContextReceiverTypes = member.contextReceivers.map {
+            it.typeRef.coneType.substitute(substitutor)
+        }
+
         if (newReceiverType == null &&
             newReturnType == null &&
             newTypeParameters === member.typeParameters &&
-            fakeOverrideSubstitution == null
+            fakeOverrideSubstitution == null &&
+            newContextReceiverTypes.all { it == null }
         ) {
             if (original.dispatchReceiverType?.substitute(substitutor) != null) {
                 return FirFakeOverrideGenerator.createSubstitutionOverrideProperty(
@@ -241,6 +261,7 @@ class FirClassSubstitutionScope(
             member,
             newDispatchReceiverType ?: dispatchReceiverTypeForSubstitutedMembers,
             newReceiverType,
+            newContextReceiverTypes,
             newReturnType,
             newTypeParameters as List<FirTypeParameter>,
             makeExpect,
@@ -316,6 +337,10 @@ class FirClassSubstitutionScope(
             it.returnTypeRef.coneType.substitute()
         }.orEmpty()
 
+        val newContextReceiverTypes = member.contextReceivers.map {
+            it.typeRef.coneType.substitute()
+        }
+
         if (original.dispatchReceiverType?.substitute(substitutor) == null &&
             newReturnType == null &&
             newGetterParameterTypes.all { it == null } &&
@@ -329,6 +354,7 @@ class FirClassSubstitutionScope(
             member,
             original,
             substitutor.substituteOrSelf(dispatchReceiverTypeForSubstitutedMembers) as ConeSimpleKotlinType?,
+            newContextReceiverTypes,
             newReturnType,
             newGetterParameterTypes,
             newSetterParameterTypes,

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirFakeOverrideGenerator.kt

@@ -36,6 +36,7 @@ object FirFakeOverrideGenerator {
         baseFunction: FirSimpleFunction,
         newDispatchReceiverType: ConeSimpleKotlinType?,
         newReceiverType: ConeKotlinType? = null,
+        newContextReceiverTypes: List<ConeKotlinType?>? = null,
         newReturnType: ConeKotlinType? = null,
         newParameterTypes: List<ConeKotlinType?>? = null,
         newTypeParameters: List<FirTypeParameter>? = null,
@@ -43,8 +44,8 @@ object FirFakeOverrideGenerator {
         fakeOverrideSubstitution: FakeOverrideSubstitution? = null
     ): FirNamedFunctionSymbol {
         createSubstitutionOverrideFunction(
-            symbolForSubstitutionOverride, session, baseFunction, newDispatchReceiverType, newReceiverType, newReturnType,
+            symbolForSubstitutionOverride, session, baseFunction, newDispatchReceiverType, newReceiverType, newContextReceiverTypes,
-            newParameterTypes, newTypeParameters, isExpect, fakeOverrideSubstitution
+            newReturnType, newParameterTypes, newTypeParameters, isExpect, fakeOverrideSubstitution
         )
         return symbolForSubstitutionOverride
     }
@@ -63,6 +64,7 @@ object FirFakeOverrideGenerator {
         baseFunction: FirSimpleFunction,
         newDispatchReceiverType: ConeSimpleKotlinType?,
         newReceiverType: ConeKotlinType?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newReturnType: ConeKotlinType?,
         newParameterTypes: List<ConeKotlinType?>?,
         newTypeParameters: List<FirTypeParameter>?,
@@ -81,6 +83,7 @@ object FirFakeOverrideGenerator {
             newParameterTypes,
             newTypeParameters,
             newReceiverType,
+            newContextReceiverTypes,
             newReturnType,
             fakeOverrideSubstitution = fakeOverrideSubstitution
         ).apply {
@@ -98,6 +101,7 @@ object FirFakeOverrideGenerator {
         newParameterTypes: List<ConeKotlinType?>? = null,
         newTypeParameters: List<FirTypeParameter>? = null,
         newReceiverType: ConeKotlinType? = null,
+        newContextReceiverTypes: List<ConeKotlinType?>? = null,
         newReturnType: ConeKotlinType? = null,
         newModality: Modality? = null,
         newVisibility: Visibility? = null,
@@ -116,7 +120,7 @@ object FirFakeOverrideGenerator {
             attributes = baseFunction.attributes.copy()
             typeParameters += configureAnnotationsTypeParametersAndSignature(
                 session, baseFunction, newParameterTypes, newTypeParameters,
-                newReceiverType, newReturnType, fakeOverrideSubstitution, newSymbol
+                newReceiverType, newContextReceiverTypes, newReturnType, fakeOverrideSubstitution, newSymbol
             ).filterIsInstance<FirTypeParameter>()
             deprecation = baseFunction.deprecation
         }
@@ -130,6 +134,7 @@ object FirFakeOverrideGenerator {
         newDispatchReceiverType: ConeSimpleKotlinType?,
         newReturnType: ConeKotlinType?,
         newParameterTypes: List<ConeKotlinType?>?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newTypeParameters: List<FirTypeParameterRef>?,
         isExpect: Boolean,
         fakeOverrideSubstitution: FakeOverrideSubstitution?
@@ -149,6 +154,7 @@ object FirFakeOverrideGenerator {
                 newParameterTypes,
                 newTypeParameters,
                 newReceiverType = null,
+                newContextReceiverTypes,
                 newReturnType,
                 fakeOverrideSubstitution,
                 fakeOverrideSymbol
@@ -171,6 +177,7 @@ object FirFakeOverrideGenerator {
         newParameterTypes: List<ConeKotlinType?>?,
         newTypeParameters: List<FirTypeParameterRef>?,
         newReceiverType: ConeKotlinType?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newReturnType: ConeKotlinType?,
         fakeOverrideSubstitution: FakeOverrideSubstitution?,
         symbolForOverride: FirBasedSymbol<*>,
@@ -181,6 +188,7 @@ object FirFakeOverrideGenerator {
                     baseFunction,
                     newParameterTypes,
                     newReceiverType,
+                    newContextReceiverTypes,
                     newReturnType,
                     fakeOverrideSubstitution
                 )
@@ -194,8 +202,8 @@ object FirFakeOverrideGenerator {
                     substitutor.substituteOrNull(it.returnTypeRef.coneType)
                 }
                 val symbol = baseFunction.symbol
-                val (copiedReceiverType, possibleReturnType) = substituteReceiverAndReturnType(
+                val (copiedReceiverType, copiedContextReceiverTypes, possibleReturnType) = substituteReceiverAndReturnType(
-                    baseFunction as FirCallableDeclaration, newReceiverType, newReturnType, substitutor
+                    baseFunction as FirCallableDeclaration, newReceiverType, newContextReceiverTypes, newReturnType, substitutor
                 )
                 val (copiedReturnType, newFakeOverrideSubstitution) = when (possibleReturnType) {
                     is Maybe.Value -> possibleReturnType.value to null
@@ -205,6 +213,7 @@ object FirFakeOverrideGenerator {
                     baseFunction,
                     copiedParameterTypes,
                     copiedReceiverType,
+                    copiedContextReceiverTypes,
                     copiedReturnType,
                     newFakeOverrideSubstitution
                 )
@@ -215,6 +224,7 @@ object FirFakeOverrideGenerator {
                     baseFunction,
                     newParameterTypes,
                     newReceiverType,
+                    newContextReceiverTypes,
                     newReturnType,
                     fakeOverrideSubstitution
                 )
@@ -227,8 +237,9 @@ object FirFakeOverrideGenerator {
         baseFunction: FirFunction,
         newParameterTypes: List<ConeKotlinType?>?,
         newReceiverType: ConeKotlinType?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newReturnType: ConeKotlinType?,
-        fakeOverrideSubstitution: FakeOverrideSubstitution?
+        fakeOverrideSubstitution: FakeOverrideSubstitution?,
     ) {
         annotations += baseFunction.annotations
 
@@ -256,6 +267,14 @@ object FirFakeOverrideGenerator {
                 symbol = FirValueParameterSymbol(valueParameter.name)
             }
         }
+
+        contextReceivers += baseFunction.contextReceivers.zip(
+            newContextReceiverTypes ?: List(baseFunction.contextReceivers.size) { null }
+        ) { contextReceiver, newType ->
+            buildContextReceiverCopy(contextReceiver) {
+                typeRef = contextReceiver.typeRef.withReplacedConeType(newType)
+            }
+        }
     }
 
     fun createSubstitutionOverrideProperty(
@@ -264,6 +283,7 @@ object FirFakeOverrideGenerator {
         baseProperty: FirProperty,
         newDispatchReceiverType: ConeSimpleKotlinType?,
         newReceiverType: ConeKotlinType? = null,
+        newContextReceiverTypes: List<ConeKotlinType?>? = null,
         newReturnType: ConeKotlinType? = null,
         newTypeParameters: List<FirTypeParameter>? = null,
         isExpect: Boolean = baseProperty.isExpect,
@@ -271,7 +291,7 @@ object FirFakeOverrideGenerator {
     ): FirPropertySymbol {
         createCopyForFirProperty(
             symbolForSubstitutionOverride, baseProperty, session, FirDeclarationOrigin.SubstitutionOverride, isExpect,
-            newDispatchReceiverType, newTypeParameters, newReceiverType, newReturnType,
+            newDispatchReceiverType, newTypeParameters, newReceiverType, newContextReceiverTypes, newReturnType,
             fakeOverrideSubstitution = fakeOverrideSubstitution
         ).apply {
             originalForSubstitutionOverrideAttr = baseProperty
@@ -296,6 +316,7 @@ object FirFakeOverrideGenerator {
         newDispatchReceiverType: ConeSimpleKotlinType?,
         newTypeParameters: List<FirTypeParameter>? = null,
         newReceiverType: ConeKotlinType? = null,
+        newContextReceiverTypes: List<ConeKotlinType?>? = null,
         newReturnType: ConeKotlinType? = null,
         newModality: Modality? = null,
         newVisibility: Visibility? = null,
@@ -319,6 +340,7 @@ object FirFakeOverrideGenerator {
                 baseProperty,
                 newTypeParameters,
                 newReceiverType,
+                newContextReceiverTypes,
                 newReturnType,
                 fakeOverrideSubstitution
             )
@@ -331,30 +353,37 @@ object FirFakeOverrideGenerator {
         baseProperty: FirProperty,
         newTypeParameters: List<FirTypeParameter>?,
         newReceiverType: ConeKotlinType?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newReturnType: ConeKotlinType?,
         fakeOverrideSubstitution: FakeOverrideSubstitution?
     ): List<FirTypeParameter> {
         return when {
             baseProperty.typeParameters.isEmpty() -> {
-                configureAnnotationsAndSignature(baseProperty, newReceiverType, newReturnType, fakeOverrideSubstitution)
+                configureAnnotationsAndSignature(
+                    baseProperty, newReceiverType, newContextReceiverTypes, newReturnType, fakeOverrideSubstitution
+                )
                 emptyList()
             }
             newTypeParameters == null -> {
                 val (copiedTypeParameters, substitutor) = createNewTypeParametersAndSubstitutor(
                     useSiteSession, baseProperty, symbol, ConeSubstitutor.Empty
                 )
-                val (copiedReceiverType, possibleReturnType) = substituteReceiverAndReturnType(
+                val (copiedReceiverType, copiedContextReceiverTypes, possibleReturnType) = substituteReceiverAndReturnType(
-                    baseProperty, newReceiverType, newReturnType, substitutor
+                    baseProperty, newReceiverType, newContextReceiverTypes, newReturnType, substitutor
                 )
                 val (copiedReturnType, newFakeOverrideSubstitution) = when (possibleReturnType) {
                     is Maybe.Value -> possibleReturnType.value to null
                     else -> null to FakeOverrideSubstitution(substitutor, baseProperty.symbol)
                 }
-                configureAnnotationsAndSignature(baseProperty, copiedReceiverType, copiedReturnType, newFakeOverrideSubstitution)
+                configureAnnotationsAndSignature(
+                    baseProperty, copiedReceiverType, copiedContextReceiverTypes, copiedReturnType, newFakeOverrideSubstitution
+                )
                 copiedTypeParameters.filterIsInstance<FirTypeParameter>()
             }
             else -> {
-                configureAnnotationsAndSignature(baseProperty, newReceiverType, newReturnType, fakeOverrideSubstitution)
+                configureAnnotationsAndSignature(
+                    baseProperty, newReceiverType, newContextReceiverTypes, newReturnType, fakeOverrideSubstitution
+                )
                 newTypeParameters
             }
         }
@@ -363,27 +392,39 @@ object FirFakeOverrideGenerator {
     private fun substituteReceiverAndReturnType(
         baseCallable: FirCallableDeclaration,
         newReceiverType: ConeKotlinType?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newReturnType: ConeKotlinType?,
         substitutor: ConeSubstitutor
-    ): Pair<ConeKotlinType?, Maybe<ConeKotlinType?>> {
+    ): Triple<ConeKotlinType?, List<ConeKotlinType?>, Maybe<ConeKotlinType?>> {
         val copiedReceiverType = newReceiverType?.let {
             substitutor.substituteOrNull(it)
         } ?: baseCallable.receiverTypeRef?.let {
             substitutor.substituteOrNull(it.coneType)
         }
 
+        val copiedContextReceiverTypes = newContextReceiverTypes?.map {
+            it?.type?.let(substitutor::substituteOrNull)
+        } ?: baseCallable.contextReceivers.map {
+            substitutor.substituteOrNull(it.typeRef.coneType)
+        }
+
         val copiedReturnType = newReturnType?.let {
             substitutor.substituteOrNull(it)
         } ?: baseCallable.returnTypeRef.let {
-            val coneType = baseCallable.returnTypeRef.coneTypeSafe<ConeKotlinType>() ?: return copiedReceiverType to Maybe.Nothing
+            val coneType = baseCallable.returnTypeRef.coneTypeSafe<ConeKotlinType>() ?: return Triple(
+                copiedReceiverType,
+                copiedContextReceiverTypes,
+                Maybe.Nothing,
+            )
             substitutor.substituteOrNull(coneType)
         }
-        return copiedReceiverType to Maybe.Value(copiedReturnType)
+        return Triple(copiedReceiverType, copiedContextReceiverTypes, Maybe.Value(copiedReturnType))
     }
 
     private fun FirPropertyBuilder.configureAnnotationsAndSignature(
         baseProperty: FirProperty,
         newReceiverType: ConeKotlinType?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newReturnType: ConeKotlinType?,
         fakeOverrideSubstitution: FakeOverrideSubstitution?
     ) {
@@ -400,6 +441,13 @@ object FirFakeOverrideGenerator {
             returnTypeRef = baseProperty.returnTypeRef.withReplacedReturnType(newReturnType)
         }
         receiverTypeRef = baseProperty.receiverTypeRef?.withReplacedConeType(newReceiverType)
+        contextReceivers += baseProperty.contextReceivers.zip(
+            newContextReceiverTypes ?: List(baseProperty.contextReceivers.size) { null }
+        ) { contextReceiver, newType ->
+            buildContextReceiverCopy(contextReceiver) {
+                typeRef = contextReceiver.typeRef.withReplacedConeType(newType)
+            }
+        }
     }
 
     fun createSubstitutionOverrideField(
@@ -438,6 +486,7 @@ object FirFakeOverrideGenerator {
         baseProperty: FirSyntheticProperty,
         baseSymbol: FirSyntheticPropertySymbol,
         newDispatchReceiverType: ConeSimpleKotlinType?,
+        newContextReceiverTypes: List<ConeKotlinType?>?,
         newReturnType: ConeKotlinType?,
         newGetterParameterTypes: List<ConeKotlinType?>?,
         newSetterParameterTypes: List<ConeKotlinType?>?,
@@ -450,6 +499,7 @@ object FirFakeOverrideGenerator {
             baseProperty.getter.delegate,
             newDispatchReceiverType,
             newReceiverType = null,
+            newContextReceiverTypes,
             newReturnType,
             newGetterParameterTypes,
             newTypeParameters = null,
@@ -463,6 +513,7 @@ object FirFakeOverrideGenerator {
             baseSetter.delegate,
             newDispatchReceiverType,
             newReceiverType = null,
+            newContextReceiverTypes,
             StandardClassIds.Unit.constructClassLikeType(emptyArray(), isNullable = false),
             newSetterParameterTypes,
             newTypeParameters = null,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ConstructorProcessing.kt

@@ -15,12 +15,12 @@ import org.jetbrains.kotlin.fir.declarations.utils.isInner
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
-import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirFakeOverrideGenerator
 import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
 import org.jetbrains.kotlin.fir.scopes.scopeForClass
+import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visibilityChecker
@@ -197,6 +197,10 @@ private fun FirTypeAliasSymbol.findSAMConstructorForTypeAlias(
         valueParameter.returnTypeRef.coneType.let(substitutor::substituteOrNull)
     }
 
+    val newContextReceiverTypes = samConstructorForClass.contextReceivers.map { contextReceiver ->
+        contextReceiver.typeRef.coneType.let(substitutor::substituteOrNull)
+    }
+
     if (newReturnType == null && newParameterTypes.all { it == null }) return samConstructorForClass
 
     val symbolForOverride = FirFakeOverrideGenerator.createSymbolForSubstitutionOverride(namedSymbol, expansionRegularClass.classId)
@@ -205,6 +209,7 @@ private fun FirTypeAliasSymbol.findSAMConstructorForTypeAlias(
         session, symbolForOverride, samConstructorForClass,
         newDispatchReceiverType = null,
         newReceiverType = null,
+        newContextReceiverTypes,
         newReturnType, newParameterTypes, typeParameters,
     ).fir
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirContextReceiverBuilder.kt

@@ -44,3 +44,16 @@ inline fun buildContextReceiver(init: FirContextReceiverBuilder.() -> Unit): Fir
     }
     return FirContextReceiverBuilder().apply(init).build()
 }
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildContextReceiverCopy(original: FirContextReceiver, init: FirContextReceiverBuilder.() -> Unit): FirContextReceiver {
+    contract {
+        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
+    }
+    val copyBuilder = FirContextReceiverBuilder()
+    copyBuilder.source = original.source
+    copyBuilder.typeRef = original.typeRef
+    copyBuilder.customLabelName = original.customLabelName
+    copyBuilder.labelNameFromTypeRef = original.labelNameFromTypeRef
+    return copyBuilder.apply(init).build()
+}

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt

@@ -280,6 +280,10 @@ object BuilderConfigurator : AbstractBuilderConfigurator<FirTreeBuilder>(FirTree
             parents += loopJumpBuilder
         }
 
+        builder(contextReceiver) {
+            withCopy()
+        }
+
         builder(valueParameter, type = "FirValueParameterImpl") {
             openBuilder()
             withCopy()

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

@@ -1,6 +1,5 @@
 // !LANGUAGE: +ContextReceivers
 // TARGET_BACKEND: JVM_IR
-// IGNORE_BACKEND_FIR: JVM_IR
 // FIR status: context receivers aren't yet supported
 // WITH_STDLIB
 
@@ -12,4 +11,4 @@ fun box(): String {
     with (listOf(1, 2, 3)) {
         return A().result
     }
-}
+}

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

@@ -1,6 +1,5 @@
 // !LANGUAGE: +ContextReceivers
 // TARGET_BACKEND: JVM_IR
-// IGNORE_BACKEND_FIR: JVM_IR
 // FIR status: context receivers aren't yet supported
 
 context(T) class B<T : CharSequence> {
@@ -9,4 +8,4 @@ context(T) class B<T : CharSequence> {
 
 fun box() = with("OK") {
     B().result
-}
+}

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

@@ -0,0 +1,29 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+
+class Box<E>(val x: E)
+
+class A<X, Y> {
+    context(Box<X>, Y)
+    fun foo(): String = x.toString() + this@Y.toString()
+
+    context(Box<X>, Y)
+    val p1: String get() = x.toString() + this@Y.toString()
+}
+
+context(Box<X>, Y)
+fun <X, Y> bar(): String = x.toString() + this@Y.toString()
+
+fun box(): String {
+    return with(Box("OK")) {
+        with(56) {
+            val a = A<String, Int>()
+            if (a.foo() != "OK56") return "fail 1"
+            if (a.p1 != "OK56") return "fail 2"
+
+            val b = bar<String, Int>()
+            if (b != "OK56") return "fail 3"
+            return "OK"
+        }
+    }
+}

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

@@ -0,0 +1,18 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+
+class Components(val x: String)
+
+context(Components)
+abstract class A<F : CharSequence>(val y: F) {
+    fun foo(): String = x + y
+}
+
+context(Components)
+class B(y: String) : A<String>(y)
+
+fun box(): String {
+    return with(Components("O")) {
+        B("K").foo()
+    }
+}

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

@@ -16752,12 +16752,24 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/simpleCall.kt");
             }
 
+            @Test
+            @TestMetadata("substitutedContextReceivers.kt")
+            public void testSubstitutedContextReceivers() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/substitutedContextReceivers.kt");
+            }
+
             @Test
             @TestMetadata("superClassAndSubClassWithContextReceiver.kt")
             public void testSuperClassAndSubClassWithContextReceiver() throws Exception {
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiver.kt");
             }
 
+            @Test
+            @TestMetadata("superClassAndSubClassWithContextReceiverSubstituted.kt")
+            public void testSuperClassAndSubClassWithContextReceiverSubstituted() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/superClassAndSubClassWithContextReceiverSubstituted.kt");
+            }
+
             @Test
             @TestMetadata("suspendContextualWithExtension.kt")
             public void testSuspendContextualWithExtension() throws Exception {