FIR: Fix behavior for context-receiver contained super class constructor

^KT-51889 Fixed

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConversionScope.kt

@@ -17,6 +17,7 @@ class Fir2IrConversionScope {
     private val parentStack = mutableListOf<IrDeclarationParent>()
 
     private val containingFirClassStack = mutableListOf<FirClass>()
+    private val currentlyGeneratedDelegatedConstructors = mutableMapOf<IrClass, IrConstructor>()
 
     fun <T : IrDeclarationParent?> withParent(parent: T, f: T.() -> Unit): T {
         if (parent == null) return parent
@@ -26,6 +27,16 @@ class Fir2IrConversionScope {
         return parent
     }
 
+    fun <T> forDelegatingConstructorCall(constructor: IrConstructor, irClass: IrClass, f: () -> T): T {
+        currentlyGeneratedDelegatedConstructors[irClass] = constructor
+        val result = f()
+        currentlyGeneratedDelegatedConstructors.remove(irClass)
+        return result
+    }
+
+    fun getConstructorForCurrentlyGeneratedDelegatedConstructor(itClass: IrClass): IrConstructor? =
+        currentlyGeneratedDelegatedConstructors[itClass]
+
     fun containingFileIfAny(): IrFile? = parentStack.getOrNull(0) as? IrFile
 
     fun withContainingFirClass(containingFirClass: FirClass, f: () -> Unit) {

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -455,15 +455,24 @@ class Fir2IrVisitor(
                 return thisReceiverExpression.convertWithOffsets { startOffset, endOffset ->
                     val thisRef = IrGetValueImpl(startOffset, endOffset, dispatchReceiver.type, dispatchReceiver.symbol)
                     if (calleeReference.contextReceiverNumber != -1) {
-                        val contextReceivers =
-                            components.classifierStorage.getFieldsWithContextReceiversForClass(irClass)
-                                ?: error("Not defined context receivers for $irClass")
+                        val constructorForCurrentlyGeneratedDelegatedConstructor =
+                            conversionScope.getConstructorForCurrentlyGeneratedDelegatedConstructor(irClass)
 
-                        IrGetFieldImpl(
-                            startOffset, endOffset, contextReceivers[calleeReference.contextReceiverNumber].symbol,
-                            thisReceiverExpression.typeRef.toIrType(),
-                            thisRef,
-                        )
+                        if (constructorForCurrentlyGeneratedDelegatedConstructor != null) {
+                            val constructorParameter =
+                                constructorForCurrentlyGeneratedDelegatedConstructor.valueParameters[calleeReference.contextReceiverNumber]
+                            IrGetValueImpl(startOffset, endOffset, constructorParameter.type, constructorParameter.symbol)
+                        } else {
+                            val contextReceivers =
+                                components.classifierStorage.getFieldsWithContextReceiversForClass(irClass)
+                                    ?: error("Not defined context receivers for $irClass")
+
+                            IrGetFieldImpl(
+                                startOffset, endOffset, contextReceivers[calleeReference.contextReceiverNumber].symbol,
+                                thisReceiverExpression.typeRef.toIrType(),
+                                thisRef,
+                            )
+                        }
                     } else {
                         thisRef
                     }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/CallAndReferenceGenerator.kt

@@ -555,11 +555,10 @@ class CallAndReferenceGenerator(
         statement: FirStatement?,
         annotationMode: Boolean
     ): IrExpression {
-        val qualifiedAccess = statement as? FirQualifiedAccess
         val call = statement as? FirCall
         return when (this) {
             is IrMemberAccessExpression<*> -> {
-                val contextReceiverCount = putContextReceiverArguments(qualifiedAccess)
+                val contextReceiverCount = putContextReceiverArguments(statement)
                 if (call == null) return this
                 val argumentsCount = call.arguments.size
                 if (argumentsCount <= valueArgumentsCount) {
@@ -624,8 +623,10 @@ class CallAndReferenceGenerator(
     }
 
     private fun IrMemberAccessExpression<*>.putContextReceiverArguments(statement: FirStatement?): Int {
-        val contextReceiverCount = (statement as? FirQualifiedAccess)?.contextReceiverArguments?.size ?: 0
-        if (statement is FirQualifiedAccess && contextReceiverCount > 0) {
+        if (statement !is FirContextReceiverArgumentListOwner) return 0
+
+        val contextReceiverCount = statement.contextReceiverArguments.size
+        if (contextReceiverCount > 0) {
             for (index in 0 until contextReceiverCount) {
                 putValueArgument(
                     index,

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/ClassMemberGenerator.kt

@@ -116,11 +116,13 @@ internal class ClassMemberGenerator(
             if (firFunction is FirConstructor && irFunction is IrConstructor && !firFunction.isExpect) {
                 val body = factory.createBlockBody(startOffset, endOffset)
                 val delegatedConstructor = firFunction.delegatedConstructor
+                val irClass = parent as IrClass
                 if (delegatedConstructor != null) {
-                    val irDelegatingConstructorCall = delegatedConstructor.toIrDelegatingConstructorCall()
+                    val irDelegatingConstructorCall = conversionScope.forDelegatingConstructorCall(irFunction, irClass) {
+                        delegatedConstructor.toIrDelegatingConstructorCall()
+                    }
                     body.statements += irDelegatingConstructorCall
                 }
-                val irClass = parent as IrClass
                 if (delegatedConstructor?.isThis == false) {
                     val instanceInitializerCall = IrInstanceInitializerCallImpl(
                         startOffset, endOffset, irClass.symbol, irFunction.constructedClassType

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -648,6 +648,7 @@ class FirCallResolver(
                 if (symbol is FirConstructorSymbol && symbol.fir.isInner) {
                     transformDispatchReceiver(StoreReceiver, singleCandidate.dispatchReceiverExpression())
                 }
+                replaceContextReceiverArguments(singleCandidate.contextReceiverArguments())
             }
         }
     }

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

@@ -67,6 +67,7 @@ sealed class CallKind(vararg resolutionSequence: ResolutionStage) {
         CollectTypeVariableUsagesInfo,
         CheckDispatchReceiver,
         CheckExtensionReceiver,
+        CheckContextReceivers,
         CheckDslScopeViolation,
         CheckArguments,
         EagerResolveOfCallableReferences,

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

@@ -1,6 +1,5 @@
 // !LANGUAGE: +ContextReceivers
 // TARGET_BACKEND: JVM_IR
-// IGNORE_BACKEND_FIR: JVM_IR
 
 class Components(val x: String)