[fir2ir] Just refactoring before classes receiver fix for evaluator

KT-61144

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

@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlin.utils.addToStdlib.getOrPut
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import org.jetbrains.kotlin.utils.addToStdlib.runUnless
 
@@ -135,40 +136,34 @@ class Fir2IrClassifierStorage(
         preCacheTypeParameters(klass, symbol)
         setTypeParameters(klass)
         if (klass is FirRegularClass) {
-            val fieldsForContextReceiversOfCurrentClass = createContextReceiverFields(klass)
+            val fieldsForContextReceiversOfCurrentClass = getFieldsWithContextReceiversForClass(this, klass)
-            if (fieldsForContextReceiversOfCurrentClass.isNotEmpty()) {
+            declarations.addAll(fieldsForContextReceiversOfCurrentClass)
-                declarations.addAll(fieldsForContextReceiversOfCurrentClass)
+        }
-                fieldsForContextReceivers[this] = fieldsForContextReceiversOfCurrentClass
+    }
+
+    fun getFieldsWithContextReceiversForClass(irClass: IrClass, klass: FirClass): List<IrField> {
+        if (klass !is FirRegularClass || klass.contextReceivers.isEmpty()) return emptyList()
+
+        return fieldsForContextReceivers.getOrPut(irClass) {
+            klass.contextReceivers.withIndex().map { (index, contextReceiver) ->
+                components.irFactory.createField(
+                    startOffset = UNDEFINED_OFFSET,
+                    endOffset = UNDEFINED_OFFSET,
+                    origin = IrDeclarationOrigin.FIELD_FOR_CLASS_CONTEXT_RECEIVER,
+                    name = Name.identifier("contextReceiverField$index"),
+                    visibility = DescriptorVisibilities.PRIVATE,
+                    symbol = IrFieldSymbolImpl(),
+                    type = contextReceiver.typeRef.toIrType(),
+                    isFinal = true,
+                    isStatic = false,
+                    isExternal = false,
+                ).also {
+                    it.parent = irClass
+                }
             }
         }
     }
 
-    fun IrClass.createContextReceiverFields(klass: FirRegularClass): List<IrField> {
-        if (klass.contextReceivers.isEmpty()) return emptyList()
-
-        val contextReceiverFields = mutableListOf<IrField>()
-        for ((index, contextReceiver) in klass.contextReceivers.withIndex()) {
-            val irField = components.irFactory.createField(
-                startOffset = UNDEFINED_OFFSET,
-                endOffset = UNDEFINED_OFFSET,
-                origin = IrDeclarationOrigin.FIELD_FOR_CLASS_CONTEXT_RECEIVER,
-                name = Name.identifier("contextReceiverField$index"),
-                visibility = DescriptorVisibilities.PRIVATE,
-                symbol = IrFieldSymbolImpl(),
-                type = contextReceiver.typeRef.toIrType(),
-                isFinal = true,
-                isStatic = false,
-                isExternal = false,
-            )
-            irField.parent = this@createContextReceiverFields
-            contextReceiverFields.add(irField)
-        }
-
-        return contextReceiverFields
-    }
-
-    fun getFieldsWithContextReceiversForClass(irClass: IrClass): List<IrField>? = fieldsForContextReceivers[irClass]
-
     private fun IrClass.declareSupertypes(klass: FirClass) {
         superTypes = klass.superTypeRefs.map { superTypeRef -> superTypeRef.toIrType() }
     }

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

@@ -691,8 +691,11 @@ class Fir2IrVisitor(
                                 IrGetValueImpl(startOffset, endOffset, constructorParameter.type, constructorParameter.symbol)
                             } else {
                                 val contextReceivers =
-                                    components.classifierStorage.getFieldsWithContextReceiversForClass(irClass)
+                                    components.classifierStorage.getFieldsWithContextReceiversForClass(irClass, firClass)
-                                        ?: error("Not defined context receivers for $irClass")
+                                require(contextReceivers.size > calleeReference.contextReceiverNumber) {
+                                    "Not defined context receiver #${calleeReference.contextReceiverNumber} for $irClass. " +
+                                            "Context receivers found: $contextReceivers"
+                                }
 
                                 IrGetFieldImpl(
                                     startOffset, endOffset, contextReceivers[calleeReference.contextReceiverNumber].symbol,

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

@@ -130,13 +130,16 @@ internal class ClassMemberGenerator(
 
                 if (containingClass is FirRegularClass && containingClass.contextReceivers.isNotEmpty()) {
                     val contextReceiverFields =
-                        components.classifierStorage.getFieldsWithContextReceiversForClass(irClass)
+                        components.classifierStorage.getFieldsWithContextReceiversForClass(irClass, containingClass)
-                            ?: error("Not found context receiver fields")
 
                     val thisParameter =
                         conversionScope.dispatchReceiverParameter(irClass) ?: error("No found this parameter for $irClass")
 
                     for (index in containingClass.contextReceivers.indices) {
+                        require(contextReceiverFields.size > index) {
+                            "Not defined context receiver #${index} for $irClass. " +
+                                    "Context receivers found: $contextReceiverFields"
+                        }
                         val irValueParameter = valueParameters[index]
                         body.statements.add(
                             IrSetFieldImpl(

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt

@@ -209,7 +209,7 @@ class Fir2IrLazyClass(
         addDeclarationsFromScope(fir.staticScope(session, scopeSession))
 
         with(classifierStorage) {
-            result.addAll(this@Fir2IrLazyClass.createContextReceiverFields(fir))
+            result.addAll(getFieldsWithContextReceiversForClass(this@Fir2IrLazyClass, fir))
         }
 
         for (name in scope.getCallableNames()) {