FIR: Support context receivers in Fir2Ir

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
                 }