FIR: Support context receivers in Fir2Ir

compiler/fir/fir2ir/jvm-backend/src/org/jetbrains/kotlin/fir/backend/jvm/FirJvmMangleComputer.kt

@@ -147,6 +147,11 @@ open class FirJvmMangleComputer(
             builder.appendSignature(MangleConstant.STATIC_MEMBER_MARK)
         }
 
+        contextReceivers.forEach {
+            builder.appendSignature(MangleConstant.CONTEXT_RECEIVER_PREFIX)
+            mangleType(builder, it.typeRef.coneType)
+        }
+
         receiverTypeRef?.let {
             builder.appendSignature(MangleConstant.EXTENSION_RECEIVER_PREFIX)
             mangleType(builder, it.coneType)

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

@@ -721,3 +721,9 @@ fun FirVariableAssignment.getIrAssignmentOrigin(): IrStatementOrigin {
     }
     return IrStatementOrigin.EQ
 }
+
+fun FirCallableDeclaration.contextReceiversForFunctionOrContainingProperty(): List<FirContextReceiver> =
+    if (this is FirPropertyAccessor)
+        this.propertySymbol?.fir?.contextReceivers.orEmpty()
+    else
+        this.contextReceivers

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

@@ -58,6 +58,8 @@ class Fir2IrClassifierStorage(
 
     private val enumEntryCache = mutableMapOf<FirEnumEntry, IrEnumEntry>()
 
+    private val fieldsForContextReceivers = mutableMapOf<IrClass, List<IrField>>()
+
     private val localStorage = Fir2IrLocalStorage()
 
     private fun FirTypeRef.toIrType(typeContext: ConversionTypeContext = ConversionTypeContext.DEFAULT): IrType =
@@ -126,9 +128,39 @@ class Fir2IrClassifierStorage(
         if (klass is FirRegularClass) {
             preCacheTypeParameters(klass)
             setTypeParameters(klass)
+            val fieldsForContextReceiversOfCurrentClass = createContextReceiverFields(klass)
+            if (fieldsForContextReceiversOfCurrentClass.isNotEmpty()) {
+                declarations.addAll(fieldsForContextReceiversOfCurrentClass)
+                fieldsForContextReceivers[this] = fieldsForContextReceiversOfCurrentClass
+            }
         }
     }
 
+    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(
+                UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                IrDeclarationOrigin.FIELD_FOR_CLASS_CONTEXT_RECEIVER,
+                IrFieldSymbolImpl(),
+                Name.identifier("contextReceiverField$index"),
+                contextReceiver.typeRef.toIrType(),
+                DescriptorVisibilities.PRIVATE,
+                isFinal = true,
+                isExternal = false,
+                isStatic = 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/Fir2IrDeclarationStorage.kt

@@ -28,7 +28,6 @@ import org.jetbrains.kotlin.fir.lazy.Fir2IrLazyConstructor
 import org.jetbrains.kotlin.fir.lazy.Fir2IrLazyProperty
 import org.jetbrains.kotlin.fir.lazy.Fir2IrLazySimpleFunction
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isLocalClassOrAnonymousObject
-import org.jetbrains.kotlin.fir.types.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.isKFunctionInvoke
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -335,14 +334,21 @@ class Fir2IrDeclarationStorage(
             val type = function.valueParameters.first().returnTypeRef.toIrType(ConversionTypeContext.DEFAULT.inSetter())
             declareDefaultSetterParameter(type)
         } else if (function != null) {
-            valueParameters = function.valueParameters.mapIndexed { index, valueParameter ->
-                createIrParameter(
-                    valueParameter, index,
-                    useStubForDefaultValueStub = function !is FirConstructor || containingClass?.name != Name.identifier("Enum"),
-                    typeContext,
-                    skipDefaultParameter = isFakeOverride
-                ).apply {
-                    this.parent = parent
+            val contextReceivers = function.contextReceiversForFunctionOrContainingProperty()
+
+            contextReceiverParametersCount = contextReceivers.size
+            valueParameters = buildList {
+                addContextReceiverParametersTo(contextReceivers, parent, this)
+
+                function.valueParameters.mapIndexedTo(this) { index, valueParameter ->
+                    createIrParameter(
+                        valueParameter, index + contextReceiverParametersCount,
+                        useStubForDefaultValueStub = function !is FirConstructor || containingClass?.name != Name.identifier("Enum"),
+                        typeContext,
+                        skipDefaultParameter = isFakeOverride
+                    ).apply {
+                        this.parent = parent
+                    }
                 }
             }
         }
@@ -391,6 +397,18 @@ class Fir2IrDeclarationStorage(
         }
     }
 
+    fun addContextReceiverParametersTo(
+        contextReceivers: List<FirContextReceiver>,
+        parent: IrFunction,
+        result: MutableList<IrValueParameter>,
+    ) {
+        contextReceivers.mapIndexedTo(result) { index, contextReceiver ->
+            createIrParameterFromContextReceiver(contextReceiver, index).apply {
+                this.parent = parent
+            }
+        }
+    }
+
     private fun <T : IrFunction> T.bindAndDeclareParameters(
         function: FirFunction?,
         irParent: IrDeclarationParent?,
@@ -404,7 +422,9 @@ class Fir2IrDeclarationStorage(
     }
 
     fun <T : IrFunction> T.putParametersInScope(function: FirFunction): T {
-        for ((firParameter, irParameter) in function.valueParameters.zip(valueParameters)) {
+        val contextReceivers = function.contextReceiversForFunctionOrContainingProperty()
+
+        for ((firParameter, irParameter) in function.valueParameters.zip(valueParameters.drop(contextReceivers.size))) {
             localStorage.putParameter(firParameter, irParameter)
         }
         return this
@@ -1055,6 +1075,22 @@ class Fir2IrDeclarationStorage(
         return irParameter
     }
 
+    private fun createIrParameterFromContextReceiver(
+        contextReceiver: FirContextReceiver,
+        index: Int,
+    ): IrValueParameter = convertCatching(contextReceiver) {
+        val type = contextReceiver.typeRef.toIrType()
+        return contextReceiver.convertWithOffsets { startOffset, endOffset ->
+            irFactory.createValueParameter(
+                startOffset, endOffset, IrDeclarationOrigin.DEFINED, IrValueParameterSymbolImpl(),
+                Name.identifier("_context_receiver_$index"), index, type,
+                null,
+                isCrossinline = false, isNoinline = false,
+                isHidden = false, isAssignable = false
+            )
+        }
+    }
+
     private var lastTemporaryIndex: Int = 0
     private fun nextTemporaryIndex(): Int = lastTemporaryIndex++
 

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

@@ -6,7 +6,10 @@
 package org.jetbrains.kotlin.fir.backend
 
 import com.intellij.psi.tree.IElementType
-import org.jetbrains.kotlin.*
+import org.jetbrains.kotlin.KtFakeSourceElementKind
+import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.KtPsiSourceElement
+import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -23,7 +26,6 @@ import org.jetbrains.kotlin.fir.declarations.utils.isSynthetic
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
-import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirStubStatement
 import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
 import org.jetbrains.kotlin.fir.references.FirReference
@@ -451,13 +453,39 @@ class Fir2IrVisitor(
             val dispatchReceiver = conversionScope.dispatchReceiverParameter(irClass)
             if (dispatchReceiver != null) {
                 return thisReceiverExpression.convertWithOffsets { startOffset, endOffset ->
-                    IrGetValueImpl(startOffset, endOffset, dispatchReceiver.type, dispatchReceiver.symbol)
+                    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")
+
+                        IrGetFieldImpl(
+                            startOffset, endOffset, contextReceivers[calleeReference.contextReceiverNumber].symbol,
+                            thisReceiverExpression.typeRef.toIrType(),
+                            thisRef,
+                        )
+                    } else {
+                        thisRef
+                    }
                 }
             }
         } else if (boundSymbol is FirCallableSymbol) {
-            val receiverSymbol =
-                calleeReference.toSymbolForCall(FirNoReceiverExpression, session, classifierStorage, declarationStorage, conversionScope)
-            val receiver = (receiverSymbol?.owner as? IrSimpleFunction)?.extensionReceiverParameter
+            val irFunction = when (boundSymbol) {
+                is FirFunctionSymbol -> declarationStorage.getIrFunctionSymbol(boundSymbol).owner
+                is FirPropertySymbol -> {
+                    val property = declarationStorage.getIrPropertySymbol(boundSymbol).owner as? IrProperty
+                    property?.let { conversionScope.parentAccessorOfPropertyFromStack(it) }
+                }
+                else -> null
+            }
+
+            val receiver = irFunction?.let { function ->
+                if (calleeReference.contextReceiverNumber != -1)
+                    function.valueParameters[calleeReference.contextReceiverNumber]
+                else
+                    function.extensionReceiverParameter
+            }
+
             if (receiver != null) {
                 return thisReceiverExpression.convertWithOffsets { startOffset, endOffset ->
                     IrGetValueImpl(startOffset, endOffset, receiver.type, receiver.symbol)
@@ -1208,4 +1236,4 @@ val KtSourceElement.operationToken: IElementType?
         return if (this is KtPsiSourceElement) (psi as? KtBinaryExpression)?.operationToken
         else treeStructure.findChildByType(lighterASTNode, KtNodeTypes.OPERATION_REFERENCE)?.tokenType
             ?: error("No operation reference for binary expression: $lighterASTNode")
-    }
+    }

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

@@ -90,7 +90,7 @@ class AnnotationGenerator(private val components: Fir2IrComponents) : Fir2IrComp
             propertyAccessor.annotations += property.annotations
                 .filter { it.useSiteTarget == AnnotationUseSiteTarget.PROPERTY_SETTER }
                 .toIrAnnotations()
-            val parameter = propertyAccessor.valueParameters.single()
+            val parameter = propertyAccessor.valueParameters.last()
             parameter.annotations += property.annotations
                 .filter { it.useSiteTarget == AnnotationUseSiteTarget.SETTER_PARAMETER }
                 .toIrAnnotations()

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

@@ -7,19 +7,21 @@ package org.jetbrains.kotlin.fir.backend.generators
 
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.backend.common.ir.isMethodOfAny
-import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isCompanion
+import org.jetbrains.kotlin.fir.dispatchReceiverClassOrNull
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
+import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.references.FirDelegateFieldReference
 import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.impl.FirReferencePlaceholderForResolvedAnnotations
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticFunctionSymbol
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -318,7 +320,7 @@ class CallAndReferenceGenerator(
                             getter != null -> IrCallImpl(
                                 startOffset, endOffset, type, getter.symbol,
                                 typeArgumentsCount = getter.typeParameters.size,
-                                valueArgumentsCount = 0,
+                                valueArgumentsCount = getter.valueParameters.size,
                                 origin = IrStatementOrigin.GET_PROPERTY,
                                 superQualifierSymbol = dispatchReceiver.superQualifierSymbol()
                             )
@@ -346,7 +348,7 @@ class CallAndReferenceGenerator(
                     else -> generateErrorCallExpression(startOffset, endOffset, calleeReference, type)
                 }
             }.applyTypeArguments(qualifiedAccess).applyReceivers(qualifiedAccess, explicitReceiverExpression)
-                .applyCallArguments(qualifiedAccess as? FirCall, annotationMode)
+                .applyCallArguments(qualifiedAccess, annotationMode)
         } catch (e: Throwable) {
             throw IllegalStateException(
                 "Error while translating ${qualifiedAccess.render()} " +
@@ -375,10 +377,11 @@ class CallAndReferenceGenerator(
                             setter != null -> IrCallImpl(
                                 startOffset, endOffset, type, setter.symbol,
                                 typeArgumentsCount = setter.typeParameters.size,
-                                valueArgumentsCount = 1,
+                                valueArgumentsCount = setter.valueParameters.size,
                                 origin = origin,
                                 superQualifierSymbol = variableAssignment.dispatchReceiver.superQualifierSymbol()
                             ).apply {
+                                putContextReceiverArguments(variableAssignment)
                                 putValueArgument(0, assignedValue)
                             }
                             else -> generateErrorCallExpression(startOffset, endOffset, calleeReference)
@@ -392,11 +395,11 @@ class CallAndReferenceGenerator(
                             setter != null -> IrCallImpl(
                                 startOffset, endOffset, type, setter.symbol,
                                 typeArgumentsCount = setter.typeParameters.size,
-                                valueArgumentsCount = 1,
+                                valueArgumentsCount = setter.valueParameters.size,
                                 origin = origin,
                                 superQualifierSymbol = variableAssignment.dispatchReceiver.superQualifierSymbol()
                             ).apply {
-                                putValueArgument(0, assignedValue)
+                                putValueArgument(putContextReceiverArguments(variableAssignment), assignedValue)
                             }
                             backingField != null -> IrSetFieldImpl(
                                 startOffset, endOffset, backingField.symbol, type,
@@ -548,10 +551,16 @@ class CallAndReferenceGenerator(
         return ConeSubstitutorByMap(map, session)
     }
 
-    internal fun IrExpression.applyCallArguments(call: FirCall?, annotationMode: Boolean): IrExpression {
-        if (call == null) return this
+    internal fun IrExpression.applyCallArguments(
+        statement: FirStatement?,
+        annotationMode: Boolean
+    ): IrExpression {
+        val qualifiedAccess = statement as? FirQualifiedAccess
+        val call = statement as? FirCall
         return when (this) {
             is IrMemberAccessExpression<*> -> {
+                val contextReceiverCount = putContextReceiverArguments(qualifiedAccess)
+                if (call == null) return this
                 val argumentsCount = call.arguments.size
                 if (argumentsCount <= valueArgumentsCount) {
                     apply {
@@ -573,7 +582,10 @@ class CallAndReferenceGenerator(
                         val substitutor = (call as? FirFunctionCall)?.buildSubstitutorByCalledFunction(function) ?: ConeSubstitutor.Empty
                         if (argumentMapping != null && (annotationMode || argumentMapping.isNotEmpty())) {
                             if (valueParameters != null) {
-                                return applyArgumentsWithReorderingIfNeeded(argumentMapping, valueParameters, substitutor, annotationMode)
+                                return applyArgumentsWithReorderingIfNeeded(
+                                    argumentMapping, valueParameters, substitutor, annotationMode,
+                                    contextReceiverCount,
+                                )
                             }
                         }
                         // Case without argument mapping (deserialized annotation)
@@ -584,7 +596,10 @@ class CallAndReferenceGenerator(
                                 else -> null
                             } ?: valueParameters?.get(index)
                             val argumentExpression = convertArgument(argument, valueParameter, substitutor)
-                            putValueArgument(valueParameters?.indexOf(valueParameter)?.takeIf { it >= 0 } ?: index, argumentExpression)
+                            putValueArgument(
+                                (valueParameters?.indexOf(valueParameter)?.takeIf { it >= 0 } ?: index) + contextReceiverCount,
+                                argumentExpression
+                            )
                         }
                     }
                 } else {
@@ -600,7 +615,7 @@ class CallAndReferenceGenerator(
                 }
             }
             is IrErrorCallExpressionImpl -> apply {
-                for (argument in call.arguments) {
+                for (argument in call?.arguments.orEmpty()) {
                     addArgument(visitor.convertToIrExpression(argument))
                 }
             }
@@ -608,11 +623,26 @@ class CallAndReferenceGenerator(
         }
     }
 
+    private fun IrMemberAccessExpression<*>.putContextReceiverArguments(statement: FirStatement?): Int {
+        val contextReceiverCount = (statement as? FirQualifiedAccess)?.contextReceiverArguments?.size ?: 0
+        if (statement is FirQualifiedAccess && contextReceiverCount > 0) {
+            for (index in 0 until contextReceiverCount) {
+                putValueArgument(
+                    index,
+                    visitor.convertToIrExpression(statement.contextReceiverArguments[index]),
+                )
+            }
+        }
+
+        return contextReceiverCount
+    }
+
     private fun IrMemberAccessExpression<*>.applyArgumentsWithReorderingIfNeeded(
         argumentMapping: Map<FirExpression, FirValueParameter>,
         valueParameters: List<FirValueParameter>,
         substitutor: ConeSubstitutor,
-        annotationMode: Boolean
+        annotationMode: Boolean,
+        contextReceiverCount: Int,
     ): IrExpression {
         val converted = argumentMapping.entries.map { (argument, parameter) ->
             parameter to convertArgument(argument, parameter, substitutor, annotationMode)
@@ -633,13 +663,16 @@ class CallAndReferenceGenerator(
                 dispatchReceiver = dispatchReceiver?.freeze("\$this")
                 extensionReceiver = extensionReceiver?.freeze("\$receiver")
                 for ((parameter, irArgument) in converted) {
-                    putValueArgument(valueParameters.indexOf(parameter), irArgument.freeze(parameter.name.asString()))
+                    putValueArgument(
+                        valueParameters.indexOf(parameter) + contextReceiverCount,
+                        irArgument.freeze(parameter.name.asString())
+                    )
                 }
                 statements.add(this@applyArgumentsWithReorderingIfNeeded)
             }
         } else {
             for ((parameter, irArgument) in converted) {
-                putValueArgument(valueParameters.indexOf(parameter), irArgument)
+                putValueArgument(valueParameters.indexOf(parameter) + contextReceiverCount, irArgument)
             }
             if (annotationMode) {
                 for ((index, parameter) in valueParameters.withIndex()) {

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

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolvedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrFieldAccessExpression
@@ -104,8 +105,9 @@ internal class ClassMemberGenerator(
                         irFunction.putParametersInScope(firFunction)
                     }
                 }
+                val irParameters = valueParameters.drop(firFunction.contextReceivers.size)
                 val annotationMode = containingClass?.classKind == ClassKind.ANNOTATION_CLASS && irFunction is IrConstructor
-                for ((valueParameter, firValueParameter) in valueParameters.zip(firFunction.valueParameters)) {
+                for ((valueParameter, firValueParameter) in irParameters.zip(firFunction.valueParameters)) {
                     valueParameter.setDefaultValue(firValueParameter, annotationMode)
                     annotationGenerator.generate(valueParameter, firValueParameter, irFunction is IrConstructor)
                 }
@@ -118,12 +120,42 @@ internal class ClassMemberGenerator(
                     val irDelegatingConstructorCall = delegatedConstructor.toIrDelegatingConstructorCall()
                     body.statements += irDelegatingConstructorCall
                 }
+                val irClass = parent as IrClass
                 if (delegatedConstructor?.isThis == false) {
                     val instanceInitializerCall = IrInstanceInitializerCallImpl(
-                        startOffset, endOffset, (parent as IrClass).symbol, irFunction.constructedClassType
+                        startOffset, endOffset, irClass.symbol, irFunction.constructedClassType
                     )
                     body.statements += instanceInitializerCall
                 }
+
+                if (containingClass is FirRegularClass && containingClass.contextReceivers.isNotEmpty()) {
+                    val contextReceiverFields =
+                        components.classifierStorage.getFieldsWithContextReceiversForClass(irClass)
+                            ?: error("Not found context receiver fields")
+
+                    val thisParameter =
+                        conversionScope.dispatchReceiverParameter(irClass) ?: error("No found this parameter for $irClass")
+
+                    val receiver = IrGetValueImpl(
+                        UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                        thisParameter.type,
+                        thisParameter.symbol,
+                    )
+
+                    for (index in containingClass.contextReceivers.indices) {
+                        val irValueParameter = valueParameters[index]
+                        body.statements.add(
+                            IrSetFieldImpl(
+                                UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                                contextReceiverFields[index].symbol,
+                                receiver,
+                                IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, irValueParameter.type, irValueParameter.symbol),
+                                components.irBuiltIns.unitType,
+                            )
+                        )
+                    }
+                }
+
                 val regularBody = firFunction.body?.let { visitor.convertToIrBlockBody(it) }
                 if (regularBody != null) {
                     body.statements += regularBody.statements

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

@@ -197,6 +197,10 @@ class Fir2IrLazyClass(
             }
         }
 
+        with(classifierStorage) {
+            result.addAll(this@Fir2IrLazyClass.createContextReceiverFields(fir))
+        }
+
         for (name in scope.getCallableNames()) {
             result += getFakeOverridesByName(name)
         }

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

@@ -6,10 +6,15 @@
 package org.jetbrains.kotlin.fir.lazy
 
 import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
+import org.jetbrains.kotlin.descriptors.DescriptorVisibility
 import org.jetbrains.kotlin.fir.backend.Fir2IrComponents
 import org.jetbrains.kotlin.fir.backend.declareThisReceiverParameter
 import org.jetbrains.kotlin.fir.backend.toIrType
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirConstructor
+import org.jetbrains.kotlin.fir.declarations.utils.isExpect
+import org.jetbrains.kotlin.fir.declarations.utils.isExternal
+import org.jetbrains.kotlin.fir.declarations.utils.isInline
+import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.symbols.Fir2IrConstructorSymbol
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.declarations.*
@@ -20,13 +25,8 @@ import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
-import org.jetbrains.kotlin.descriptors.DescriptorVisibility
-import org.jetbrains.kotlin.fir.declarations.utils.isExpect
-import org.jetbrains.kotlin.fir.declarations.utils.isExternal
-import org.jetbrains.kotlin.fir.declarations.utils.isInline
-import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 
 class Fir2IrLazyConstructor(
     components: Fir2IrComponents,
@@ -100,16 +100,25 @@ class Fir2IrLazyConstructor(
             error("Mutating Fir2Ir lazy elements is not possible")
         }
 
-    override var contextReceiverParametersCount: Int = 0
+    override var contextReceiverParametersCount: Int = fir.contextReceivers.size
 
     override var valueParameters: List<IrValueParameter> by lazyVar(lock) {
         declarationStorage.enterScope(this)
-        fir.valueParameters.mapIndexed { index, valueParameter ->
-            declarationStorage.createIrParameter(
-                valueParameter, index,
-                useStubForDefaultValueStub = (parent as? IrClass)?.name != Name.identifier("Enum")
-            ).apply {
-                this.parent = this@Fir2IrLazyConstructor
+
+        buildList {
+            declarationStorage.addContextReceiverParametersTo(
+                fir.contextReceivers,
+                this@Fir2IrLazyConstructor,
+                this@buildList,
+            )
+
+            fir.valueParameters.mapIndexedTo(this) { index, valueParameter ->
+                declarationStorage.createIrParameter(
+                    valueParameter, index,
+                    useStubForDefaultValueStub = (parent as? IrClass)?.name != Name.identifier("Enum")
+                ).apply {
+                    this.parent = this@Fir2IrLazyConstructor
+                }
             }
         }.apply {
             declarationStorage.leaveScope(this@Fir2IrLazyConstructor)

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

@@ -5,10 +5,7 @@
 
 package org.jetbrains.kotlin.fir.lazy
 
-import org.jetbrains.kotlin.fir.backend.ConversionTypeContext
-import org.jetbrains.kotlin.fir.backend.Fir2IrComponents
-import org.jetbrains.kotlin.fir.backend.generateOverriddenAccessorSymbols
-import org.jetbrains.kotlin.fir.backend.toIrType
+import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
@@ -76,21 +73,32 @@ class Fir2IrLazyPropertyAccessor(
         }
     }
 
-    override var contextReceiverParametersCount: Int = 0
+    override var contextReceiverParametersCount: Int = fir.contextReceiversForFunctionOrContainingProperty().size
 
     override var valueParameters: List<IrValueParameter> by lazyVar(lock) {
-        if (!isSetter) emptyList()
+        if (!isSetter && contextReceiverParametersCount == 0) emptyList()
         else {
             declarationStorage.enterScope(this)
-            listOf(
-                declarationStorage.createDefaultSetterParameter(
-                    startOffset, endOffset,
-                    (firAccessor?.valueParameters?.firstOrNull()?.returnTypeRef ?: firParentProperty.returnTypeRef).toIrType(
-                        typeConverter, conversionTypeContext
-                    ),
-                    parent = this@Fir2IrLazyPropertyAccessor
+
+            buildList {
+                declarationStorage.addContextReceiverParametersTo(
+                    fir.contextReceiversForFunctionOrContainingProperty(),
+                    this@Fir2IrLazyPropertyAccessor,
+                    this@buildList,
                 )
-            ).apply {
+
+                if (isSetter) {
+                    add(
+                        declarationStorage.createDefaultSetterParameter(
+                            startOffset, endOffset,
+                            (firAccessor?.valueParameters?.firstOrNull()?.returnTypeRef ?: firParentProperty.returnTypeRef).toIrType(
+                                typeConverter, conversionTypeContext
+                            ),
+                            parent = this@Fir2IrLazyPropertyAccessor
+                        )
+                    )
+                }
+            }.apply {
                 declarationStorage.leaveScope(this@Fir2IrLazyPropertyAccessor)
             }
         }

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.lazy
 
 import org.jetbrains.kotlin.fir.backend.Fir2IrComponents
+import org.jetbrains.kotlin.fir.backend.contextReceiversForFunctionOrContainingProperty
 import org.jetbrains.kotlin.fir.backend.generateOverriddenFunctionSymbols
 import org.jetbrains.kotlin.fir.backend.toIrType
 import org.jetbrains.kotlin.fir.declarations.FirFunction
@@ -64,15 +65,24 @@ class Fir2IrLazySimpleFunction(
         }
     }
 
-    override var contextReceiverParametersCount: Int = 0
+    override var contextReceiverParametersCount: Int = fir.contextReceiversForFunctionOrContainingProperty().size
 
     override var valueParameters: List<IrValueParameter> by lazyVar(lock) {
         declarationStorage.enterScope(this)
-        fir.valueParameters.mapIndexed { index, valueParameter ->
-            declarationStorage.createIrParameter(
-                valueParameter, index, skipDefaultParameter = isFakeOverride
-            ).apply {
-                this.parent = this@Fir2IrLazySimpleFunction
+
+        buildList {
+            declarationStorage.addContextReceiverParametersTo(
+                fir.contextReceiversForFunctionOrContainingProperty(),
+                this@Fir2IrLazySimpleFunction,
+                this@buildList,
+            )
+
+            fir.valueParameters.mapIndexedTo(this) { index, valueParameter ->
+                declarationStorage.createIrParameter(
+                    valueParameter, index, skipDefaultParameter = isFakeOverride
+                ).apply {
+                    this.parent = this@Fir2IrLazySimpleFunction
+                }
             }
         }.apply {
             declarationStorage.leaveScope(this@Fir2IrLazySimpleFunction)

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/calls/FirReceivers.kt

@@ -79,6 +79,7 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
     protected val useSiteSession: FirSession,
     protected val scopeSession: ScopeSession,
     private val mutable: Boolean,
+    val contextReceiverNumber: Int = -1,
 ) : ReceiverValue {
     final override var type: ConeKotlinType = type
         private set
@@ -92,7 +93,7 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
 
     override fun scope(useSiteSession: FirSession, scopeSession: ScopeSession): FirTypeScope? = implicitScope
 
-    private val originalReceiverExpression: FirThisReceiverExpression = receiverExpression(boundSymbol, type)
+    private val originalReceiverExpression: FirThisReceiverExpression = receiverExpression(boundSymbol, type, contextReceiverNumber)
     final override var receiverExpression: FirExpression = originalReceiverExpression
         private set
 
@@ -119,13 +120,18 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
     abstract fun createSnapshot(): ImplicitReceiverValue<S>
 }
 
-private fun receiverExpression(symbol: FirBasedSymbol<*>, type: ConeKotlinType): FirThisReceiverExpression =
+private fun receiverExpression(
+    symbol: FirBasedSymbol<*>,
+    type: ConeKotlinType,
+    contextReceiverNumber: Int,
+): FirThisReceiverExpression =
     buildThisReceiverExpression {
         // NB: we can't use `symbol.fir.source` as the source of `this` receiver. For instance, if this is an implicit receiver for a class,
         // the entire class itself will be set as a source. If combined with an implicit type operation, a certain assertion, like null
         // check assertion, will retrieve source as an assertion message, which is literally the entire class (!).
         calleeReference = buildImplicitThisReference {
             boundSymbol = symbol
+            this.contextReceiverNumber = contextReceiverNumber
         }
         typeRef = buildResolvedTypeRef {
             this.type = type
@@ -193,8 +199,9 @@ sealed class ContextReceiverValue<S : FirBasedSymbol<*>>(
     useSiteSession: FirSession,
     scopeSession: ScopeSession,
     mutable: Boolean = true,
+    contextReceiverNumber: Int,
 ) : ImplicitReceiverValue<S>(
-    boundSymbol, type, useSiteSession, scopeSession, mutable
+    boundSymbol, type, useSiteSession, scopeSession, mutable, contextReceiverNumber,
 ) {
     abstract override fun createSnapshot(): ContextReceiverValue<S>
 }
@@ -206,11 +213,12 @@ class ContextReceiverValueForCallable(
     useSiteSession: FirSession,
     scopeSession: ScopeSession,
     mutable: Boolean = true,
+    contextReceiverNumber: Int,
 ) : ContextReceiverValue<FirCallableSymbol<*>>(
-    boundSymbol, type, labelName, useSiteSession, scopeSession, mutable
+    boundSymbol, type, labelName, useSiteSession, scopeSession, mutable, contextReceiverNumber
 ) {
     override fun createSnapshot(): ContextReceiverValue<FirCallableSymbol<*>> =
-        ContextReceiverValueForCallable(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false)
+        ContextReceiverValueForCallable(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
 
     override val isContextReceiver: Boolean
         get() = true
@@ -223,11 +231,12 @@ class ContextReceiverValueForClass(
     useSiteSession: FirSession,
     scopeSession: ScopeSession,
     mutable: Boolean = true,
+    contextReceiverNumber: Int,
 ) : ContextReceiverValue<FirClassSymbol<*>>(
-    boundSymbol, type, labelName, useSiteSession, scopeSession, mutable
+    boundSymbol, type, labelName, useSiteSession, scopeSession, mutable, contextReceiverNumber
 ) {
     override fun createSnapshot(): ContextReceiverValue<FirClassSymbol<*>> =
-        ContextReceiverValueForClass(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false)
+        ContextReceiverValueForClass(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
 
     override val isContextReceiver: Boolean
         get() = true

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

@@ -121,6 +121,7 @@ class FirCallResolver(
                 dispatchReceiver = candidate.dispatchReceiverExpression()
                 extensionReceiver = candidate.chosenExtensionReceiverExpression()
                 argumentList = candidate.callInfo.argumentList
+                contextReceiverArguments.addAll(candidate.contextReceiverArguments())
             }
         } else {
             resultExpression
@@ -382,6 +383,7 @@ class FirCallResolver(
             val candidate = reducedCandidates.single()
             resultExpression = resultExpression.transformDispatchReceiver(StoreReceiver, candidate.dispatchReceiverExpression())
             resultExpression = resultExpression.transformExtensionReceiver(StoreReceiver, candidate.chosenExtensionReceiverExpression())
+            resultExpression.replaceContextReceiverArguments(candidate.contextReceiverArguments())
         }
         if (resultExpression is FirExpression) transformer.storeTypeFromCallee(resultExpression)
         return resultExpression

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

@@ -104,6 +104,9 @@ class Candidate(
     fun chosenExtensionReceiverExpression(): FirExpression =
         chosenExtensionReceiverValue?.receiverExpression?.takeIf { it !is FirExpressionStub } ?: FirNoReceiverExpression
 
+    fun contextReceiverArguments(): List<FirExpression> =
+        contextReceiverArguments ?: emptyList()
+
     var hasVisibleBackingField = false
 
     override fun equals(other: Any?): Boolean {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -147,6 +147,9 @@ class FirCallCompletionResultsWriterTransformer(
             )
             .transformDispatchReceiver(StoreReceiver, dispatchReceiver)
             .transformExtensionReceiver(StoreReceiver, extensionReceiver) as T
+
+        result.replaceContextReceiverArguments(subCandidate.contextReceiverArguments())
+
         if (result is FirPropertyAccessExpressionImpl && calleeReference.candidate.currentApplicability == CandidateApplicability.PROPERTY_AS_OPERATOR) {
             result.nonFatalDiagnostics.add(ConePropertyAsOperator(calleeReference.candidate.symbol as FirPropertySymbol))
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -97,8 +97,11 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             is FirThisReference -> {
                 val labelName = callee.labelName
                 val implicitReceiver = implicitReceiverStack[labelName]
-                implicitReceiver?.boundSymbol?.let {
-                    callee.replaceBoundSymbol(it)
+                implicitReceiver?.let {
+                    callee.replaceBoundSymbol(it.boundSymbol)
+                    if (it is ContextReceiverValue) {
+                        callee.replaceContextReceiverNumber(it.contextReceiverNumber)
+                    }
                 }
                 val implicitType = implicitReceiver?.originalType
                 qualifiedAccessExpression.resultType = when {
@@ -598,6 +601,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                     (leftArgument as? FirQualifiedAccess)?.let {
                         dispatchReceiver = it.dispatchReceiver
                         extensionReceiver = it.extensionReceiver
+                        contextReceiverArguments.addAll(it.contextReceiverArguments)
                     }
                     annotations += assignmentOperatorStatement.annotations
                 }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/declarations/ImplicitReceiverUtils.kt

@@ -90,9 +90,10 @@ fun SessionHolder.collectTowerDataElementsForClass(owner: FirClass, defaultType:
     }
 
     val thisReceiver = ImplicitDispatchReceiverValue(owner.symbol, defaultType, session, scopeSession)
-    val contextReceivers = (owner as? FirRegularClass)?.contextReceivers?.map {
+    val contextReceivers = (owner as? FirRegularClass)?.contextReceivers?.mapIndexed { index, receiver ->
         ContextReceiverValueForClass(
-            owner.symbol, it.typeRef.coneType, it.labelName, session, scopeSession,
+            owner.symbol, receiver.typeRef.coneType, receiver.labelName, session, scopeSession,
+            contextReceiverNumber = index,
         )
     }.orEmpty()
 
@@ -248,6 +249,9 @@ typealias FirLocalScopes = PersistentList<FirLocalScope>
 fun FirCallableDeclaration.createContextReceiverValues(
     sessionHolder: SessionHolder,
 ): List<ContextReceiverValueForCallable> =
-    contextReceivers.map {
-        ContextReceiverValueForCallable(symbol, it.typeRef.coneType, it.labelName, sessionHolder.session, sessionHolder.scopeSession)
+    contextReceivers.mapIndexed { index, receiver ->
+        ContextReceiverValueForCallable(
+            symbol, receiver.typeRef.coneType, receiver.labelName, sessionHolder.session, sessionHolder.scopeSession,
+            contextReceiverNumber = index,
+        )
     }