FIR: Add preliminary support for context receivers in resolution

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt

@@ -335,14 +335,14 @@ internal class KtFirCallResolver(
                 // Implicit invoke (e.g., `x()`) will have a different callee symbol (e.g., `x`) than the candidate (e.g., `invoke`).
                 createKtPartiallyAppliedSymbolForImplicitInvoke(
                     candidate.dispatchReceiverValue?.receiverExpression ?: FirNoReceiverExpression,
-                    candidate.extensionReceiverValue?.receiverExpression ?: FirNoReceiverExpression,
+                    candidate.chosenExtensionReceiverValue?.receiverExpression ?: FirNoReceiverExpression,
                     candidate.explicitReceiverKind
                 )
             } else {
                 KtPartiallyAppliedSymbol(
                     unsubstitutedKtSignature.substitute(substitutor),
                     candidate.dispatchReceiverValue?.receiverExpression?.toKtReceiverValue(),
-                    candidate.extensionReceiverValue?.receiverExpression?.toKtReceiverValue(),
+                    candidate.chosenExtensionReceiverValue?.receiverExpression?.toKtReceiverValue(),
                 )
             }
         } else if (fir is FirQualifiedAccess) {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDataClassConverters.kt

@@ -1608,6 +1608,32 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.NO_CONTEXT_RECEIVER) { firDiagnostic ->
+        NoContextReceiverImpl(
+            firSymbolBuilder.typeBuilder.buildKtType(firDiagnostic.a),
+            firDiagnostic as KtPsiDiagnostic,
+            token,
+        )
+    }
+    add(FirErrors.MULTIPLE_ARGUMENTS_APPLICABLE_FOR_CONTEXT_RECEIVER) { firDiagnostic ->
+        MultipleArgumentsApplicableForContextReceiverImpl(
+            firSymbolBuilder.typeBuilder.buildKtType(firDiagnostic.a),
+            firDiagnostic as KtPsiDiagnostic,
+            token,
+        )
+    }
+    add(FirErrors.AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER) { firDiagnostic ->
+        AmbiguousCallWithImplicitContextReceiverImpl(
+            firDiagnostic as KtPsiDiagnostic,
+            token,
+        )
+    }
+    add(FirErrors.UNSUPPORTED_CONTEXTUAL_DECLARATION_CALL) { firDiagnostic ->
+        UnsupportedContextualDeclarationCallImpl(
+            firDiagnostic as KtPsiDiagnostic,
+            token,
+        )
+    }
     add(FirErrors.RECURSION_IN_IMPLICIT_TYPES) { firDiagnostic ->
         RecursionInImplicitTypesImpl(
             firDiagnostic as KtPsiDiagnostic,

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -1145,6 +1145,24 @@ sealed class KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         abstract val candidates: List<KtSymbol>
     }
 
+    abstract class NoContextReceiver : KtFirDiagnostic<KtElement>() {
+        override val diagnosticClass get() = NoContextReceiver::class
+        abstract val contextReceiverRepresentation: KtType
+    }
+
+    abstract class MultipleArgumentsApplicableForContextReceiver : KtFirDiagnostic<KtElement>() {
+        override val diagnosticClass get() = MultipleArgumentsApplicableForContextReceiver::class
+        abstract val contextReceiverRepresentation: KtType
+    }
+
+    abstract class AmbiguousCallWithImplicitContextReceiver : KtFirDiagnostic<KtElement>() {
+        override val diagnosticClass get() = AmbiguousCallWithImplicitContextReceiver::class
+    }
+
+    abstract class UnsupportedContextualDeclarationCall : KtFirDiagnostic<KtElement>() {
+        override val diagnosticClass get() = UnsupportedContextualDeclarationCall::class
+    }
+
     abstract class RecursionInImplicitTypes : KtFirDiagnostic<PsiElement>() {
         override val diagnosticClass get() = RecursionInImplicitTypes::class
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -1373,6 +1373,28 @@ internal class NextAmbiguityImpl(
     override val token: ValidityToken,
 ) : KtFirDiagnostic.NextAmbiguity(), KtAbstractFirDiagnostic<PsiElement>
 
+internal class NoContextReceiverImpl(
+    override val contextReceiverRepresentation: KtType,
+    override val firDiagnostic: KtPsiDiagnostic,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.NoContextReceiver(), KtAbstractFirDiagnostic<KtElement>
+
+internal class MultipleArgumentsApplicableForContextReceiverImpl(
+    override val contextReceiverRepresentation: KtType,
+    override val firDiagnostic: KtPsiDiagnostic,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.MultipleArgumentsApplicableForContextReceiver(), KtAbstractFirDiagnostic<KtElement>
+
+internal class AmbiguousCallWithImplicitContextReceiverImpl(
+    override val firDiagnostic: KtPsiDiagnostic,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.AmbiguousCallWithImplicitContextReceiver(), KtAbstractFirDiagnostic<KtElement>
+
+internal class UnsupportedContextualDeclarationCallImpl(
+    override val firDiagnostic: KtPsiDiagnostic,
+    override val token: ValidityToken,
+) : KtFirDiagnostic.UnsupportedContextualDeclarationCall(), KtAbstractFirDiagnostic<KtElement>
+
 internal class RecursionInImplicitTypesImpl(
     override val firDiagnostic: KtPsiDiagnostic,
     override val token: ValidityToken,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolver/SingleCandidateResolver.kt

@@ -53,11 +53,12 @@ class SingleCandidateResolver(
             resolutionParameters.callableSymbol,
             explicitReceiverKind = explicitReceiverKind,
             dispatchReceiverValue = dispatchReceiverValue,
-            extensionReceiverValue =
+            givenExtensionReceiverOptions = listOfNotNull(
-            if (explicitReceiverKind.isExtensionReceiver)
+                if (explicitReceiverKind.isExtensionReceiver)
-                callInfo.explicitReceiver?.let { ExpressionReceiverValue(it) }
+                    callInfo.explicitReceiver?.let { ExpressionReceiverValue(it) }
-            else
+                else
-                implicitExtensionReceiverValue,
+                    implicitExtensionReceiverValue
+            ),
             scope = null,
         )
 

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/FirDiagnosticsList.kt

@@ -531,6 +531,17 @@ object DIAGNOSTICS_LIST : DiagnosticList("FirErrors") {
         }
     }
 
+    val CONTEXT_RECEIVERS_RESOLUTION by object : DiagnosticGroup("Context receivers resolution") {
+        val NO_CONTEXT_RECEIVER by error<KtElement>(PositioningStrategy.REFERENCE_BY_QUALIFIED) {
+            parameter<ConeKotlinType>("contextReceiverRepresentation")
+        }
+        val MULTIPLE_ARGUMENTS_APPLICABLE_FOR_CONTEXT_RECEIVER by error<KtElement>(PositioningStrategy.REFERENCE_BY_QUALIFIED) {
+            parameter<ConeKotlinType>("contextReceiverRepresentation")
+        }
+        val AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER by error<KtElement>(PositioningStrategy.REFERENCE_BY_QUALIFIED)
+        val UNSUPPORTED_CONTEXTUAL_DECLARATION_CALL by error<KtElement>()
+    }
+
     val TYPES_AND_TYPE_PARAMETERS by object : DiagnosticGroup("Types & type parameters") {
         val RECURSION_IN_IMPLICIT_TYPES by error<PsiElement>()
         val INFERENCE_ERROR by error<PsiElement>()

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt

@@ -355,6 +355,12 @@ object FirErrors {
     val HAS_NEXT_FUNCTION_AMBIGUITY by error1<PsiElement, Collection<FirBasedSymbol<*>>>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
     val NEXT_AMBIGUITY by error1<PsiElement, Collection<FirBasedSymbol<*>>>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
 
+    // Context receivers resolution
+    val NO_CONTEXT_RECEIVER by error1<KtElement, ConeKotlinType>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
+    val MULTIPLE_ARGUMENTS_APPLICABLE_FOR_CONTEXT_RECEIVER by error1<KtElement, ConeKotlinType>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
+    val AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER by error0<KtElement>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
+    val UNSUPPORTED_CONTEXTUAL_DECLARATION_CALL by error0<KtElement>()
+
     // Types & type parameters
     val RECURSION_IN_IMPLICIT_TYPES by error0<PsiElement>()
     val INFERENCE_ERROR by error0<PsiElement>()

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrorsDefaultMessages.kt

@@ -57,6 +57,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ACTUAL_TYPE_ALIAS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ACTUAL_WITHOUT_EXPECT
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_ACTUALS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_ANONYMOUS_TYPE_INFERRED
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_EXPECTS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_SUPER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ANNOTATION_ARGUMENT_KCLASS_LITERAL_OF_TYPE_PARAMETER_ERROR
@@ -294,6 +295,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MISPLACED_TYPE_PA
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MISSING_STDLIB_CLASS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MISSING_VAL_ON_ANNOTATION_PARAMETER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MODIFIER_FORM_FOR_NON_BUILT_IN_SUSPEND
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MULTIPLE_ARGUMENTS_APPLICABLE_FOR_CONTEXT_RECEIVER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MULTIPLE_VARARG_PARAMETERS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MUST_BE_INITIALIZED
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.MUST_BE_INITIALIZED_OR_BE_ABSTRACT
@@ -336,6 +338,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NOT_YET_SUPPORTED
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NO_ACTUAL_CLASS_MEMBER_FOR_EXPECTED_CLASS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NO_ACTUAL_FOR_EXPECT
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NO_COMPANION_OBJECT
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NO_CONTEXT_RECEIVER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NO_ELSE_IN_WHEN
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NO_EXPLICIT_RETURN_TYPE_IN_API_MODE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NO_EXPLICIT_RETURN_TYPE_IN_API_MODE_WARNING
@@ -517,6 +520,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UNSAFE_IMPLICIT_I
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UNSAFE_INFIX_CALL
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UNSAFE_OPERATOR_CALL
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UNSUPPORTED
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UNSUPPORTED_CONTEXTUAL_DECLARATION_CALL
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UNSUPPORTED_FEATURE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UNUSED_VARIABLE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.UPPER_BOUND_IS_EXTENSION_FUNCTION_TYPE
@@ -974,6 +978,21 @@ object FirErrorsDefaultMessages : BaseDiagnosticRendererFactory() {
         map.put(NEXT_AMBIGUITY, "Method ''next()'' is ambiguous for this expression: {0}", SYMBOLS)
         map.put(NEXT_NONE_APPLICABLE, "None of the ''next()'' functions is applicable for ''iterator()'' of type ''{0}''", SYMBOLS)
 
+        map.put(NO_CONTEXT_RECEIVER, "No required context receiver found: {0}", RENDER_TYPE)
+        map.put(
+            MULTIPLE_ARGUMENTS_APPLICABLE_FOR_CONTEXT_RECEIVER,
+            "Multiple arguments applicable for context receiver: {0}",
+            RENDER_TYPE
+        )
+        map.put(
+            AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER,
+            "With implicit context receiver, call is ambiguous. Specify the receiver explicitly"
+        )
+        map.put(
+            UNSUPPORTED_CONTEXTUAL_DECLARATION_CALL,
+            "To use contextual declarations, specify the `-Xcontext-receivers` compiler option"
+        )
+
         // Ambiguity
         map.put(OVERLOAD_RESOLUTION_AMBIGUITY, "Overload resolution ambiguity between candidates: {0}", SYMBOLS)
         map.put(ASSIGN_OPERATOR_AMBIGUITY, "Ambiguity between assign operator candidates: {0}", SYMBOLS)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/coneDiagnosticToFirDiagnostic.kt

@@ -186,6 +186,7 @@ private fun mapInapplicableCandidateError(
     source: KtSourceElement,
     qualifiedAccessSource: KtSourceElement?,
 ): List<KtDiagnostic> {
+    val typeContext = session.typeContext
     val genericDiagnostic = FirErrors.INAPPLICABLE_CANDIDATE.createOn(source, diagnostic.candidate.symbol)
     val diagnostics = diagnostic.candidate.diagnostics.filter { it.applicability == diagnostic.applicability }.mapNotNull { rootCause ->
         when (rootCause) {
@@ -197,7 +198,6 @@ private fun mapInapplicableCandidateError(
                 rootCause.forbiddenNamedArgumentsTarget
             )
             is ArgumentTypeMismatch -> {
-                val typeContext = session.typeContext
                 FirErrors.ARGUMENT_TYPE_MISMATCH.createOn(
                     rootCause.argument.source ?: source,
                     rootCause.expectedType.removeTypeVariableTypes(typeContext),
@@ -205,6 +205,18 @@ private fun mapInapplicableCandidateError(
                     rootCause.isMismatchDueToNullability
                 )
             }
+            is MultipleContextReceiversApplicableForExtensionReceivers ->
+                FirErrors.AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER.createOn(qualifiedAccessSource ?: source)
+            is NoApplicableValueForContextReceiver ->
+                FirErrors.NO_CONTEXT_RECEIVER.createOn(
+                    qualifiedAccessSource ?: source,
+                    rootCause.expectedContextReceiverType.removeTypeVariableTypes(typeContext)
+                )
+            is AmbiguousValuesForContextReceiverParameter ->
+                FirErrors.MULTIPLE_ARGUMENTS_APPLICABLE_FOR_CONTEXT_RECEIVER.createOn(
+                    qualifiedAccessSource ?: source,
+                    rootCause.expectedContextReceiverType.removeTypeVariableTypes(typeContext)
+                )
             is NullForNotNullType -> FirErrors.NULL_FOR_NONNULL_TYPE.createOn(
                 rootCause.argument.source ?: source
             )

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

@@ -22,8 +22,12 @@ import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeErrorType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.SmartcastStability
 
 interface Receiver
@@ -79,6 +83,8 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
     final override var type: ConeKotlinType = type
         private set
 
+    abstract val isContextReceiver: Boolean
+
     val originalType: ConeKotlinType = type
 
     var implicitScope: FirTypeScope? = type.scope(useSiteSession, scopeSession, FakeOverrideTypeCalculator.DoNothing)
@@ -144,6 +150,9 @@ class ImplicitDispatchReceiverValue(
     override fun createSnapshot(): ImplicitReceiverValue<FirClassSymbol<*>> {
         return ImplicitDispatchReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
     }
+
+    override val isContextReceiver: Boolean
+        get() = false
 }
 
 class ImplicitExtensionReceiverValue(
@@ -156,6 +165,9 @@ class ImplicitExtensionReceiverValue(
     override fun createSnapshot(): ImplicitReceiverValue<FirCallableSymbol<*>> {
         return ImplicitExtensionReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
     }
+
+    override val isContextReceiver: Boolean
+        get() = false
 }
 
 
@@ -169,4 +181,54 @@ class InaccessibleImplicitReceiverValue(
     override fun createSnapshot(): ImplicitReceiverValue<FirClassSymbol<*>> {
         return InaccessibleImplicitReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
     }
+
+    override val isContextReceiver: Boolean
+        get() = false
+}
+
+sealed class ContextReceiverValue<S : FirBasedSymbol<*>>(
+    boundSymbol: S,
+    type: ConeKotlinType,
+    val labelName: Name?,
+    useSiteSession: FirSession,
+    scopeSession: ScopeSession,
+    mutable: Boolean = true,
+) : ImplicitReceiverValue<S>(
+    boundSymbol, type, useSiteSession, scopeSession, mutable
+) {
+    abstract override fun createSnapshot(): ContextReceiverValue<S>
+}
+
+class ContextReceiverValueForCallable(
+    boundSymbol: FirCallableSymbol<*>,
+    type: ConeKotlinType,
+    labelName: Name?,
+    useSiteSession: FirSession,
+    scopeSession: ScopeSession,
+    mutable: Boolean = true,
+) : ContextReceiverValue<FirCallableSymbol<*>>(
+    boundSymbol, type, labelName, useSiteSession, scopeSession, mutable
+) {
+    override fun createSnapshot(): ContextReceiverValue<FirCallableSymbol<*>> =
+        ContextReceiverValueForCallable(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false)
+
+    override val isContextReceiver: Boolean
+        get() = true
+}
+
+class ContextReceiverValueForClass(
+    boundSymbol: FirClassSymbol<*>,
+    type: ConeKotlinType,
+    labelName: Name?,
+    useSiteSession: FirSession,
+    scopeSession: ScopeSession,
+    mutable: Boolean = true,
+) : ContextReceiverValue<FirClassSymbol<*>>(
+    boundSymbol, type, labelName, useSiteSession, scopeSession, mutable
+) {
+    override fun createSnapshot(): ContextReceiverValue<FirClassSymbol<*>> =
+        ContextReceiverValueForClass(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false)
+
+    override val isContextReceiver: Boolean
+        get() = true
 }

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/InferenceUtils.kt

@@ -218,7 +218,9 @@ private fun ConeTypeProjection.typeOrDefault(default: ConeKotlinType): ConeKotli
 
 fun ConeKotlinType.receiverType(session: FirSession): ConeKotlinType? {
     if (!isBuiltinFunctionalType(session) || !isExtensionFunctionType(session)) return null
-    return fullyExpandedType(session).typeArguments.first().typeOrDefault(session.builtinTypes.nothingType.type)
+    return fullyExpandedType(session).let { expanded ->
+        expanded.typeArguments[expanded.contextReceiversNumberForFunctionType].typeOrDefault(session.builtinTypes.nothingType.type)
+    }
 }
 
 fun ConeKotlinType.returnType(session: FirSession): ConeKotlinType {

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

@@ -119,7 +119,7 @@ class FirCallResolver(
             functionCall.copyAsImplicitInvokeCall {
                 explicitReceiver = candidate.callInfo.explicitReceiver
                 dispatchReceiver = candidate.dispatchReceiverExpression()
-                extensionReceiver = candidate.extensionReceiverExpression()
+                extensionReceiver = candidate.chosenExtensionReceiverExpression()
                 argumentList = candidate.callInfo.argumentList
             }
         } else {
@@ -381,7 +381,7 @@ class FirCallResolver(
         if (reducedCandidates.size == 1) {
             val candidate = reducedCandidates.single()
             resultExpression = resultExpression.transformDispatchReceiver(StoreReceiver, candidate.dispatchReceiverExpression())
-            resultExpression = resultExpression.transformExtensionReceiver(StoreReceiver, candidate.extensionReceiverExpression())
+            resultExpression = resultExpression.transformExtensionReceiver(StoreReceiver, candidate.chosenExtensionReceiverExpression())
         }
         if (resultExpression is FirExpression) transformer.storeTypeFromCallee(resultExpression)
         return resultExpression

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

@@ -47,9 +47,10 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.resolve.ForbiddenNamedArgumentsTarget
-import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
+import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.types.SmartcastStability
+import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
@@ -73,7 +74,10 @@ fun FirFunction.constructFunctionalType(isSuspend: Boolean = false): ConeLookupT
     }
     val rawReturnType = (this as FirCallableDeclaration).returnTypeRef.coneType
 
-    return createFunctionalType(parameters, receiverTypeRef?.coneType, rawReturnType, isSuspend = isSuspend)
+    return createFunctionalType(
+        parameters, receiverTypeRef?.coneType, rawReturnType, isSuspend = isSuspend,
+        contextReceivers = contextReceivers.map { it.typeRef.coneType }
+    )
 }
 
 fun FirFunction.constructFunctionalTypeRef(isSuspend: Boolean = false): FirResolvedTypeRef {
@@ -88,9 +92,16 @@ fun createFunctionalType(
     receiverType: ConeKotlinType?,
     rawReturnType: ConeKotlinType,
     isSuspend: Boolean,
+    contextReceivers: List<ConeKotlinType> = emptyList(),
     isKFunctionType: Boolean = false
 ): ConeLookupTagBasedType {
-    val receiverAndParameterTypes = listOfNotNull(receiverType) + parameters + listOf(rawReturnType)
+    val receiverAndParameterTypes =
+        buildList {
+            addAll(contextReceivers)
+            addIfNotNull(receiverType)
+            addAll(parameters)
+            add(rawReturnType)
+        }
 
     val kind = if (isSuspend) {
         if (isKFunctionType) FunctionClassKind.KSuspendFunction else FunctionClassKind.SuspendFunction
@@ -99,7 +110,18 @@ fun createFunctionalType(
     }
 
     val functionalTypeId = ClassId(kind.packageFqName, kind.numberedClassName(receiverAndParameterTypes.size - 1))
-    val attributes = if (receiverType != null) ConeAttributes.WithExtensionFunctionType else ConeAttributes.Empty
+    val attributes = when {
+        contextReceivers.isNotEmpty() -> ConeAttributes.create(
+            buildList {
+                add(CompilerConeAttributes.ContextFunctionTypeParams(contextReceivers.size))
+                if (receiverType != null) {
+                    add(CompilerConeAttributes.ExtensionFunctionType)
+                }
+            }
+        )
+        receiverType != null -> ConeAttributes.WithExtensionFunctionType
+        else -> ConeAttributes.Empty
+    }
     return ConeClassLikeTypeImpl(
         ConeClassLikeLookupTagImpl(functionalTypeId),
         receiverAndParameterTypes.toTypedArray(),

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

@@ -16,6 +16,7 @@ sealed class CallKind(vararg resolutionSequence: ResolutionStage) {
         CollectTypeVariableUsagesInfo,
         CheckDispatchReceiver,
         CheckExtensionReceiver,
+        CheckContextReceivers,
         CheckDslScopeViolation,
         CheckLowPriorityInOverloadResolution,
         PostponedVariablesInitializerResolutionStage,
@@ -43,6 +44,7 @@ sealed class CallKind(vararg resolutionSequence: ResolutionStage) {
         CollectTypeVariableUsagesInfo,
         CheckDispatchReceiver,
         CheckExtensionReceiver,
+        CheckContextReceivers,
         CheckDslScopeViolation,
         CheckArguments,
         CheckCallModifiers,
@@ -113,6 +115,7 @@ class ResolutionSequenceBuilder(
     var mapTypeArguments: Boolean = false,
     var resolveCallableReferenceArguments: Boolean = false,
     var checkCallableReferenceExpectedType: Boolean = false,
+    val checkContextReceivers: Boolean = false,
 ) {
     fun build(): CallKind {
         val stages = mutableListOf<ResolutionStage>().apply {
@@ -125,6 +128,7 @@ class ResolutionSequenceBuilder(
             if (checkDispatchReceiver) add(CheckDispatchReceiver)
             if (checkExtensionReceiver) add(CheckExtensionReceiver)
             if (checkArguments) add(CheckArguments)
+            if (checkContextReceivers) add(CheckContextReceivers)
             if (resolveCallableReferenceArguments) add(EagerResolveOfCallableReferences)
             if (checkLowPriorityInOverloadResolution) add(CheckLowPriorityInOverloadResolution)
             if (initializePostponedVariables) add(PostponedVariablesInitializerResolutionStage)

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

@@ -28,9 +28,11 @@ import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 class Candidate(
     override val symbol: FirBasedSymbol<*>,
     override val dispatchReceiverValue: ReceiverValue?,
-    override val extensionReceiverValue: ReceiverValue?,
+    // In most cases, it contains zero or single element
+    // More than one, only in case of context receiver group
+    val givenExtensionReceiverOptions: List<ReceiverValue>,
     override val explicitReceiverKind: ExplicitReceiverKind,
-    val constraintSystemFactory: InferenceComponents.ConstraintSystemFactory,
+    private val constraintSystemFactory: InferenceComponents.ConstraintSystemFactory,
     private val baseSystem: ConstraintStorage,
     override val callInfo: CallInfo,
     val originScope: FirScope?,
@@ -73,6 +75,10 @@ class Candidate(
     var currentApplicability = CandidateApplicability.RESOLVED
         private set
 
+    override var chosenExtensionReceiverValue: ReceiverValue? = givenExtensionReceiverOptions.singleOrNull()
+
+    var contextReceiverArguments: List<FirExpression>? = null
+
     override val applicability: CandidateApplicability
         get() = currentApplicability
 
@@ -95,8 +101,8 @@ class Candidate(
     fun dispatchReceiverExpression(): FirExpression =
         dispatchReceiverValue?.receiverExpression?.takeIf { it !is FirExpressionStub } ?: FirNoReceiverExpression
 
-    fun extensionReceiverExpression(): FirExpression =
+    fun chosenExtensionReceiverExpression(): FirExpression =
-        extensionReceiverValue?.receiverExpression?.takeIf { it !is FirExpressionStub } ?: FirNoReceiverExpression
+        chosenExtensionReceiverValue?.receiverExpression?.takeIf { it !is FirExpressionStub } ?: FirNoReceiverExpression
 
     var hasVisibleBackingField = false
 

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

@@ -49,19 +49,20 @@ class CandidateFactory private constructor(
         explicitReceiverKind: ExplicitReceiverKind,
         scope: FirScope?,
         dispatchReceiverValue: ReceiverValue? = null,
-        extensionReceiverValue: ReceiverValue? = null,
+        givenExtensionReceiverOptions: List<ReceiverValue> = emptyList(),
         objectsByName: Boolean = false
     ): Candidate {
         @Suppress("NAME_SHADOWING")
         val symbol = symbol.unwrapIntegerOperatorSymbolIfNeeded(callInfo)
 
         val result = Candidate(
-            symbol, dispatchReceiverValue, extensionReceiverValue,
+            symbol, dispatchReceiverValue, givenExtensionReceiverOptions,
             explicitReceiverKind, context.inferenceComponents.constraintSystemFactory, baseSystem,
             callInfo,
             scope,
             isFromCompanionObjectTypeScope = when (explicitReceiverKind) {
-                ExplicitReceiverKind.EXTENSION_RECEIVER -> extensionReceiverValue.isCandidateFromCompanionObjectTypeScope()
+                ExplicitReceiverKind.EXTENSION_RECEIVER ->
+                    givenExtensionReceiverOptions.singleOrNull().isCandidateFromCompanionObjectTypeScope()
                 ExplicitReceiverKind.DISPATCH_RECEIVER -> dispatchReceiverValue.isCandidateFromCompanionObjectTypeScope()
                 // The following cases are not applicable for companion objects.
                 ExplicitReceiverKind.NO_EXPLICIT_RECEIVER, ExplicitReceiverKind.BOTH_RECEIVERS -> false
@@ -121,7 +122,7 @@ class CandidateFactory private constructor(
         return Candidate(
             symbol,
             dispatchReceiverValue = null,
-            extensionReceiverValue = null,
+            givenExtensionReceiverOptions = emptyList(),
             explicitReceiverKind = ExplicitReceiverKind.NO_EXPLICIT_RECEIVER,
             context.inferenceComponents.constraintSystemFactory,
             baseSystem,

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

@@ -34,6 +34,8 @@ import org.jetbrains.kotlin.fir.visibilityChecker
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.calls.inference.isSubtypeConstraintCompatible
+import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind.*
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.resolve.deprecation.DeprecationLevelValue
@@ -76,18 +78,43 @@ internal object CheckExplicitReceiverConsistency : ResolutionStage() {
 object CheckExtensionReceiver : ResolutionStage() {
     override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
         val expectedReceiverType = candidate.getReceiverType(context) ?: return
-
-        val argumentExtensionReceiverValue = candidate.extensionReceiverValue ?: return
         val expectedType = candidate.substitutor.substituteOrSelf(expectedReceiverType.type)
-        val argumentType = captureFromTypeParameterUpperBoundIfNeeded(
+
-            argumentType = argumentExtensionReceiverValue.type,
+        // Probably, we should add an assertion here since we check consistency on the level of scope tower levels
-            expectedType = expectedType,
+        if (candidate.givenExtensionReceiverOptions.isEmpty()) return
-            session = context.session
+
-        )
+        val preparedReceivers = candidate.givenExtensionReceiverOptions.map {
+            candidate.prepareReceivers(it, expectedType, context)
+        }
+
+        if (preparedReceivers.size == 1) {
+            resolveExtensionReceiver(preparedReceivers, candidate, expectedType, sink, context)
+            return
+        }
+
+        val successfulReceivers = preparedReceivers.filter {
+            candidate.system.isSubtypeConstraintCompatible(it.type, expectedType, SimpleConstraintSystemConstraintPosition)
+        }
+
+        when (successfulReceivers.size) {
+            0 -> sink.yieldDiagnostic(InapplicableWrongReceiver())
+            1 -> resolveExtensionReceiver(successfulReceivers, candidate, expectedType, sink, context)
+            else -> sink.yieldDiagnostic(MultipleContextReceiversApplicableForExtensionReceivers())
+        }
+    }
+
+    private suspend fun resolveExtensionReceiver(
+        receivers: List<ReceiverDescription>,
+        candidate: Candidate,
+        expectedType: ConeKotlinType,
+        sink: CheckerSink,
+        context: ResolutionContext
+    ) {
+        val receiver = receivers.single()
         candidate.resolvePlainArgumentType(
             candidate.csBuilder,
-            argumentExtensionReceiverValue.receiverExpression,
+            receiver.expression,
-            argumentType = argumentType,
+            argumentType = receiver.type,
             expectedType = expectedType,
             sink = sink,
             context = context,
@@ -109,6 +136,25 @@ object CheckExtensionReceiver : ResolutionStage() {
     }
 }
 
+private fun Candidate.prepareReceivers(
+    argumentExtensionReceiverValue: ReceiverValue,
+    expectedType: ConeKotlinType,
+    context: ResolutionContext,
+): ReceiverDescription {
+    val argumentType = captureFromTypeParameterUpperBoundIfNeeded(
+        argumentType = argumentExtensionReceiverValue.type,
+        expectedType = expectedType,
+        session = context.session
+    ).let { prepareCapturedType(it, context) }
+
+    return ReceiverDescription(argumentExtensionReceiverValue.receiverExpression, argumentType)
+}
+
+private class ReceiverDescription(
+    val expression: FirExpression,
+    val type: ConeKotlinType,
+)
+
 object CheckDispatchReceiver : ResolutionStage() {
     @OptIn(SymbolInternals::class)
     override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
@@ -156,6 +202,58 @@ object CheckDispatchReceiver : ResolutionStage() {
     }
 }
 
+object CheckContextReceivers : ResolutionStage() {
+    override suspend fun check(candidate: Candidate, callInfo: CallInfo, sink: CheckerSink, context: ResolutionContext) {
+        val contextReceiverExpectedTypes = (candidate.symbol as? FirCallableSymbol<*>)?.fir?.contextReceivers?.map {
+            candidate.substitutor.substituteOrSelf(it.typeRef.coneType)
+        }?.takeUnless { it.isEmpty() } ?: return
+
+        val receiverGroups: List<List<ImplicitReceiverValue<*>>> =
+            context.bodyResolveContext.towerDataContext.towerDataElements.mapNotNull { towerDataElement ->
+                towerDataElement.implicitReceiver?.let(::listOf) ?: towerDataElement.contextReceiverGroup
+            }
+
+        val resultingContextReceiverArguments = mutableListOf<FirExpression>()
+        for (expectedType in contextReceiverExpectedTypes) {
+            val matchingReceivers = candidate.findClosestMatchingReceivers(expectedType, receiverGroups, context)
+            when (matchingReceivers.size) {
+                0 -> {
+                    sink.reportDiagnostic(NoApplicableValueForContextReceiver(expectedType))
+                    return
+                }
+                1 -> {
+                    val matchingReceiver = matchingReceivers.single()
+                    resultingContextReceiverArguments.add(matchingReceiver.expression)
+                    candidate.system.addSubtypeConstraint(matchingReceiver.type, expectedType, SimpleConstraintSystemConstraintPosition)
+                }
+                else -> {
+                    sink.reportDiagnostic(AmbiguousValuesForContextReceiverParameter(expectedType))
+                    return
+                }
+            }
+        }
+
+        candidate.contextReceiverArguments = resultingContextReceiverArguments
+    }
+}
+
+private fun Candidate.findClosestMatchingReceivers(
+    expectedType: ConeKotlinType,
+    receiverGroups: List<List<ImplicitReceiverValue<*>>>,
+    context: ResolutionContext,
+): List<ReceiverDescription> {
+    for (receiverGroup in receiverGroups) {
+        val currentResult =
+            receiverGroup
+                .map { prepareReceivers(it, expectedType, context) }
+                .filter { system.isSubtypeConstraintCompatible(it.type, expectedType, SimpleConstraintSystemConstraintPosition) }
+
+        if (currentResult.isNotEmpty()) return currentResult
+    }
+
+    return emptyList()
+}
+
 /**
  * See https://kotlinlang.org/api/latest/jvm/stdlib/kotlin/-dsl-marker/ for more details and
  * /compiler/testData/diagnostics/tests/resolve/dslMarker for the test files.
@@ -176,7 +274,7 @@ object CheckDslScopeViolation : ResolutionStage() {
             }
         }
         checkReceiverValue(candidate.dispatchReceiverValue)
-        checkReceiverValue(candidate.extensionReceiverValue)
+        checkReceiverValue(candidate.chosenExtensionReceiverValue)
 
         // For value of builtin functional type with implicit extension receiver, the receiver is passed as the first argument rather than
         // an extension receiver of the `invoke` call. Hence, we need to specially handle this case.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/FirInvokeResolveTowerExtension.kt

@@ -162,7 +162,7 @@ internal class FirInvokeResolveTowerExtension(
                 continue
             }
 
-            val extensionReceiverExpression = invokeReceiverCandidate.extensionReceiverExpression()
+            val extensionReceiverExpression = invokeReceiverCandidate.chosenExtensionReceiverExpression()
             val useImplicitReceiverAsBuiltinInvokeArgument =
                 !invokeBuiltinExtensionMode && isExtensionFunctionType &&
                         invokeReceiverCandidate.explicitReceiverKind == ExplicitReceiverKind.NO_EXPLICIT_RECEIVER
@@ -391,6 +391,13 @@ private class InvokeFunctionResolveTask(
                     info, group.Member,
                     ExplicitReceiverKind.EXTENSION_RECEIVER
                 )
+            },
+            onContextReceiverGroup = { contextReceiverGroup, towerGroup ->
+                processLevelForRegularInvoke(
+                    contextReceiverGroup.toMemberScopeTowerLevel(extensionReceiver = invokeReceiverValue),
+                    info, towerGroup.Member,
+                    ExplicitReceiverKind.EXTENSION_RECEIVER
+                )
             }
         )
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/FirTowerResolveTask.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls.tower
 
+import org.jetbrains.kotlin.fir.declarations.ContextReceiverGroup
 import org.jetbrains.kotlin.fir.declarations.FirTowerDataContext
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
@@ -47,6 +48,12 @@ internal class TowerDataElementsForName(
         }
     }
 
+    val contextReceiverGroups by lazy(LazyThreadSafetyMode.NONE) {
+        nonLocalTowerDataElements.mapIndexedNotNull { index, towerDataElement ->
+            towerDataElement.contextReceiverGroup?.let { receiver -> IndexedValue(index, receiver) }
+        }
+    }
+
     val emptyScopes = mutableSetOf<FirScope>()
     val implicitReceiverValuesWithEmptyScopes = mutableSetOf<ImplicitReceiverValue<*>>()
 }
@@ -81,22 +88,34 @@ internal abstract class FirBaseTowerResolveTask(
         extensionReceiver: ReceiverValue? = null,
         withHideMembersOnly: Boolean = false,
         includeInnerConstructors: Boolean = extensionReceiver != null,
+        contextReceiverGroup: ContextReceiverGroup? = null,
     ): ScopeTowerLevel = ScopeTowerLevel(
         components, this,
-        extensionReceiver, withHideMembersOnly, includeInnerConstructors
+        givenExtensionReceiverOptions = contextReceiverGroup ?: listOfNotNull(extensionReceiver),
+        withHideMembersOnly, includeInnerConstructors
     )
 
     protected fun ReceiverValue.toMemberScopeTowerLevel(
-        extensionReceiver: ReceiverValue? = null
+        extensionReceiver: ReceiverValue? = null,
+        contextReceiverGroup: ContextReceiverGroup? = null,
     ) = MemberScopeTowerLevel(
         components, this,
-        extensionReceiver,
+        givenExtensionReceiverOptions = contextReceiverGroup ?: listOfNotNull(extensionReceiver),
+    )
+
+    protected fun ContextReceiverGroup.toMemberScopeTowerLevel(
+        extensionReceiver: ReceiverValue? = null,
+        otherContextReceiverGroup: ContextReceiverGroup? = null,
+    ) = ContextReceiverGroupMemberScopeTowerLevel(
+        components, this,
+        givenExtensionReceiverOptions = otherContextReceiverGroup ?: listOfNotNull(extensionReceiver),
     )
 
     protected inline fun enumerateTowerLevels(
         parentGroup: TowerGroup = TowerGroup.EmptyRoot,
         onScope: (FirScope, TowerGroup) -> Unit,
         onImplicitReceiver: (ImplicitReceiverValue<*>, TowerGroup) -> Unit,
+        onContextReceiverGroup: (ContextReceiverGroup, TowerGroup) -> Unit,
     ) {
         for ((index, localScope) in towerDataElementsForName.reversedFilteredLocalScopes) {
             onScope(localScope, parentGroup.Local(index))
@@ -118,6 +137,10 @@ internal abstract class FirBaseTowerResolveTask(
                 onImplicitReceiver(receiver, parentGroup.Implicit(depth))
             }
         }
+
+        for ((depth, contextReceiverGroup) in towerDataElementsForName.contextReceiverGroups) {
+            onContextReceiverGroup(contextReceiverGroup, parentGroup.ContextReceiverGroup(depth))
+        }
     }
 
     /**
@@ -243,7 +266,17 @@ internal open class FirTowerResolveTask(
                 )
             },
             onImplicitReceiver = { implicitReceiverValue, group ->
-                processCombinationOfReceivers(implicitReceiverValue, explicitReceiverValue, info, group)
+                // Member extensions
+                processLevel(
+                    implicitReceiverValue.toMemberScopeTowerLevel(extensionReceiver = explicitReceiverValue),
+                    info, group.Member, ExplicitReceiverKind.EXTENSION_RECEIVER
+                )
+            },
+            onContextReceiverGroup = { contextReceiverGroup, towerGroup ->
+                processLevel(
+                    contextReceiverGroup.toMemberScopeTowerLevel(extensionReceiver = explicitReceiverValue),
+                    info, towerGroup, ExplicitReceiverKind.EXTENSION_RECEIVER,
+                )
             }
         )
     }
@@ -277,7 +310,13 @@ internal open class FirTowerResolveTask(
                     implicitReceiverValuesWithEmptyScopes,
                     emptyScopes
                 )
-            }
+            },
+            onContextReceiverGroup = { contextReceiverGroup, towerGroup ->
+                processCandidatesWithGivenContextReceiverGroup(
+                    contextReceiverGroup,
+                    info, towerGroup,
+                )
+            },
         )
     }
 
@@ -308,6 +347,7 @@ internal open class FirTowerResolveTask(
                     ExplicitReceiverKind.EXTENSION_RECEIVER, parentGroup
                 )
             } else {
+                // context?
                 for ((depth, implicitReceiverValue) in towerDataElementsForName.implicitReceivers) {
                     processHideMembersLevel(
                         implicitReceiverValue, topLevelScope, info, index, depth,
@@ -331,19 +371,6 @@ internal open class FirTowerResolveTask(
 
     }
 
-    private suspend fun processCombinationOfReceivers(
-        implicitReceiverValue: ImplicitReceiverValue<*>,
-        explicitReceiverValue: ExpressionReceiverValue,
-        info: CallInfo,
-        parentGroup: TowerGroup
-    ) {
-        // Member extensions
-        processLevel(
-            implicitReceiverValue.toMemberScopeTowerLevel(extensionReceiver = explicitReceiverValue),
-            info, parentGroup.Member, ExplicitReceiverKind.EXTENSION_RECEIVER
-        )
-    }
-
     private suspend fun processCandidatesWithGivenImplicitReceiverAsValue(
         receiver: ImplicitReceiverValue<*>,
         info: CallInfo,
@@ -377,11 +404,49 @@ internal open class FirTowerResolveTask(
                     implicitReceiverValue.toMemberScopeTowerLevel(extensionReceiver = receiver),
                     info, group
                 )
-            }
+            },
+            onContextReceiverGroup = { contextReceiverGroup, towerGroup ->
+                processLevel(
+                    contextReceiverGroup.toMemberScopeTowerLevel(extensionReceiver = receiver),
+                    info, towerGroup
+                )
+            },
         )
 
     }
 
+    private suspend fun processCandidatesWithGivenContextReceiverGroup(
+        contextReceiverGroup: ContextReceiverGroup,
+        info: CallInfo,
+        parentGroup: TowerGroup,
+    ) {
+        processLevel(
+            contextReceiverGroup.toMemberScopeTowerLevel(), info, parentGroup.Member,
+        )
+
+        enumerateTowerLevels(
+            parentGroup,
+            onScope = { scope, towerGroup ->
+                processLevel(
+                    scope.toScopeTowerLevel(contextReceiverGroup = contextReceiverGroup),
+                    info, towerGroup,
+                )
+            },
+            onImplicitReceiver = { implicitReceiverValue, towerGroup ->
+                processLevel(
+                    implicitReceiverValue.toMemberScopeTowerLevel(contextReceiverGroup = contextReceiverGroup),
+                    info, towerGroup
+                )
+            },
+            onContextReceiverGroup = { otherContextReceiverGroup, towerGroup ->
+                processLevel(
+                    contextReceiverGroup.toMemberScopeTowerLevel(otherContextReceiverGroup = otherContextReceiverGroup),
+                    info, towerGroup,
+                )
+            }
+        )
+    }
+
     private suspend fun processHideMembersLevel(
         receiverValue: ReceiverValue,
         topLevelScope: FirScope,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/FirTowerResolver.kt

@@ -116,7 +116,7 @@ class FirTowerResolver(
                     ExplicitReceiverKind.NO_EXPLICIT_RECEIVER,
                     scope,
                     dispatchReceiver,
-                    extensionReceiverValue = null
+                    givenExtensionReceiverOptions = emptyList()
                 ),
                 context
             )

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerGroup.kt

@@ -40,9 +40,11 @@ sealed class TowerGroupKind(val index: Byte) : Comparable<TowerGroupKind> {
 
     class ImplicitOrNonLocal(depth: Int, val kindForDebugSake: String) : WithDepth(6, depth)
 
-    object InvokeExtension : TowerGroupKind(7)
+    class ContextReceiverGroup(depth: Int) : WithDepth(7, depth)
 
-    object QualifierValue : TowerGroupKind(8)
+    object InvokeExtension : TowerGroupKind(8)
+
+    object QualifierValue : TowerGroupKind(9)
 
     class UnqualifiedEnum(depth: Int) : WithDepth(9, depth)
 
@@ -166,6 +168,8 @@ private constructor(
         fun Implicit(depth: Int) = kindOf(TowerGroupKind.Implicit(depth))
         fun NonLocal(depth: Int) = kindOf(TowerGroupKind.NonLocal(depth))
 
+        fun ContextReceiverGroup(depth: Int) = kindOf(TowerGroupKind.ContextReceiverGroup(depth))
+
         fun TopPrioritized(depth: Int) = kindOf(TowerGroupKind.TopPrioritized(depth))
 
         val Last = kindOf(TowerGroupKind.Last)
@@ -180,6 +184,8 @@ private constructor(
     fun Implicit(depth: Int) = kindOf(TowerGroupKind.Implicit(depth))
     fun NonLocal(depth: Int) = kindOf(TowerGroupKind.NonLocal(depth))
 
+    fun ContextReceiverGroup(depth: Int) = kindOf(TowerGroupKind.ContextReceiverGroup(depth))
+
     val InvokeExtension get() = kindOf(TowerGroupKind.InvokeExtension)
 
     fun TopPrioritized(depth: Int) = kindOf(TowerGroupKind.TopPrioritized(depth))

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerLevelHandler.kt

@@ -45,7 +45,7 @@ internal class TowerLevelHandler {
             CallKind.VariableAccess -> {
                 processResult += towerLevel.processPropertiesByName(info, processor)
 
-                if (!collector.isSuccess() && towerLevel is ScopeTowerLevel && towerLevel.extensionReceiver == null) {
+                if (!collector.isSuccess() && towerLevel is ScopeTowerLevel && !towerLevel.areThereExtensionReceiverOptions()) {
                     processResult += towerLevel.processObjectsByName(info, processor)
                 }
             }
@@ -74,7 +74,7 @@ private class TowerScopeLevelProcessor(
     override fun consumeCandidate(
         symbol: FirBasedSymbol<*>,
         dispatchReceiverValue: ReceiverValue?,
-        extensionReceiverValue: ReceiverValue?,
+        givenExtensionReceiverOptions: List<ReceiverValue>,
         scope: FirScope,
         objectsByName: Boolean
     ) {
@@ -85,7 +85,7 @@ private class TowerScopeLevelProcessor(
                 explicitReceiverKind,
                 scope,
                 dispatchReceiverValue,
-                extensionReceiverValue,
+                givenExtensionReceiverOptions,
                 objectsByName
             ), candidateFactory.context
         )

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerLevels.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.calls.tower
 
 import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.declarations.ContextReceiverGroup
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
@@ -56,7 +57,7 @@ abstract class TowerScopeLevel {
         fun consumeCandidate(
             symbol: T,
             dispatchReceiverValue: ReceiverValue?,
-            extensionReceiverValue: ReceiverValue?,
+            givenExtensionReceiverOptions: List<ReceiverValue>,
             scope: FirScope,
             objectsByName: Boolean = false
         )
@@ -73,7 +74,7 @@ abstract class TowerScopeLevel {
 class MemberScopeTowerLevel(
     private val bodyResolveComponents: BodyResolveComponents,
     val dispatchReceiverValue: ReceiverValue,
-    private val extensionReceiver: ReceiverValue? = null,
+    private val givenExtensionReceiverOptions: List<ReceiverValue>,
 ) : TowerScopeLevel() {
     private val scopeSession: ScopeSession get() = bodyResolveComponents.scopeSession
     private val session: FirSession get() = bodyResolveComponents.session
@@ -86,11 +87,11 @@ class MemberScopeTowerLevel(
         var (empty, candidates) = scope.collectCandidates(processScopeMembers)
         consumeCandidates(output, candidates.map { scope to it })
 
-        if (extensionReceiver == null) {
+        if (givenExtensionReceiverOptions.isEmpty()) {
             val withSynthetic = FirSyntheticPropertiesScope(session, scope)
             withSynthetic.processScopeMembers { symbol ->
                 empty = false
-                output.consumeCandidate(symbol, dispatchReceiverValue, null, scope)
+                output.consumeCandidate(symbol, dispatchReceiverValue, givenExtensionReceiverOptions = emptyList(), scope)
             }
         }
         return if (empty) ProcessResult.SCOPE_EMPTY else ProcessResult.FOUND
@@ -103,7 +104,7 @@ class MemberScopeTowerLevel(
         val result = mutableListOf<T>()
         processScopeMembers { candidate ->
             empty = false
-            if (candidate is FirCallableSymbol<*> && candidate.hasConsistentExtensionReceiver(extensionReceiver)) {
+            if (candidate is FirCallableSymbol<*> && candidate.hasConsistentExtensionReceiver(givenExtensionReceiverOptions)) {
                 val fir = candidate.fir
                 if ((fir as? FirConstructor)?.isInner == false) {
                     return@processScopeMembers
@@ -121,14 +122,14 @@ class MemberScopeTowerLevel(
         candidatesWithScope: List<Pair<FirScope, T>>
     ) {
         for ((scope, candidate) in candidatesWithScope) {
-            if (candidate is FirCallableSymbol<*> && candidate.hasConsistentExtensionReceiver(extensionReceiver)) {
+            if (candidate is FirCallableSymbol<*> && candidate.hasConsistentExtensionReceiver(givenExtensionReceiverOptions)) {
                 output.consumeCandidate(
                     candidate, dispatchReceiverValue,
-                    extensionReceiverValue = extensionReceiver,
+                    givenExtensionReceiverOptions,
                     scope
                 )
             } else if (candidate is FirClassLikeSymbol<*>) {
-                output.consumeCandidate(candidate, null, extensionReceiver, scope)
+                output.consumeCandidate(candidate, null, givenExtensionReceiverOptions, scope)
             }
         }
     }
@@ -201,8 +202,30 @@ class MemberScopeTowerLevel(
         }
     }
 
-    private fun FirCallableSymbol<*>.hasConsistentExtensionReceiver(extensionReceiver: Receiver?): Boolean {
+    private fun FirCallableSymbol<*>.hasConsistentExtensionReceiver(givenExtensionReceivers: List<ReceiverValue>): Boolean {
-        return (extensionReceiver != null) == hasExtensionReceiver()
+        return givenExtensionReceivers.isNotEmpty() == hasExtensionReceiver()
+    }
+}
+
+class ContextReceiverGroupMemberScopeTowerLevel(
+    bodyResolveComponents: BodyResolveComponents,
+    contextReceiverGroup: ContextReceiverGroup,
+    givenExtensionReceiverOptions: List<ReceiverValue> = emptyList(),
+) : TowerScopeLevel() {
+    private val memberScopeLevels = contextReceiverGroup.map {
+        MemberScopeTowerLevel(bodyResolveComponents, it, givenExtensionReceiverOptions)
+    }
+
+    override fun processFunctionsByName(info: CallInfo, processor: TowerScopeLevelProcessor<FirFunctionSymbol<*>>): ProcessResult {
+        return memberScopeLevels.minOf { it.processFunctionsByName(info, processor) }
+    }
+
+    override fun processPropertiesByName(info: CallInfo, processor: TowerScopeLevelProcessor<FirVariableSymbol<*>>): ProcessResult {
+        return memberScopeLevels.minOf { it.processPropertiesByName(info, processor) }
+    }
+
+    override fun processObjectsByName(info: CallInfo, processor: TowerScopeLevelProcessor<FirBasedSymbol<*>>): ProcessResult {
+        return memberScopeLevels.minOf { it.processObjectsByName(info, processor) }
     }
 }
 
@@ -216,12 +239,14 @@ class MemberScopeTowerLevel(
 class ScopeTowerLevel(
     private val bodyResolveComponents: BodyResolveComponents,
     val scope: FirScope,
-    val extensionReceiver: ReceiverValue?,
+    private val givenExtensionReceiverOptions: List<ReceiverValue>,
     private val withHideMembersOnly: Boolean,
     private val includeInnerConstructors: Boolean
 ) : TowerScopeLevel() {
     private val session: FirSession get() = bodyResolveComponents.session
 
+    fun areThereExtensionReceiverOptions(): Boolean = givenExtensionReceiverOptions.isNotEmpty()
+
     private fun dispatchReceiverValue(candidate: FirCallableSymbol<*>): ReceiverValue? {
         candidate.fir.importedFromObjectData?.let { data ->
             val objectClassId = data.objectClassId
@@ -257,26 +282,25 @@ class ScopeTowerLevel(
 
     private fun shouldSkipCandidateWithInconsistentExtensionReceiver(candidate: FirCallableSymbol<*>): Boolean {
         // Pre-check explicit extension receiver for default package top-level members
-        if (scope is FirDefaultStarImportingScope && extensionReceiver != null) {
+        if (scope !is FirDefaultStarImportingScope || !areThereExtensionReceiverOptions()) return false
+
+        val declarationReceiverType = candidate.resolvedReceiverTypeRef?.coneType as? ConeClassLikeType ?: return false
+        val startProjectedDeclarationReceiverType = declarationReceiverType.lookupTag.constructClassType(
+            declarationReceiverType.typeArguments.map { ConeStarProjection }.toTypedArray(),
+            isNullable = true
+        )
+
+        return givenExtensionReceiverOptions.none { extensionReceiver ->
             val extensionReceiverType = extensionReceiver.type
-            if (extensionReceiverType is ConeClassLikeType) {
+            // If some receiver is non class like, we should not skip it
-                val declarationReceiverType = candidate.resolvedReceiverTypeRef?.coneType
+            if (extensionReceiverType !is ConeClassLikeType) return@none true
-                if (declarationReceiverType is ConeClassLikeType) {
+
-                    if (!AbstractTypeChecker.isSubtypeOf(
+            AbstractTypeChecker.isSubtypeOf(
-                            session.typeContext,
+                session.typeContext,
-                            extensionReceiverType,
+                extensionReceiverType,
-                            declarationReceiverType.lookupTag.constructClassType(
+                startProjectedDeclarationReceiverType
-                                declarationReceiverType.typeArguments.map { ConeStarProjection }.toTypedArray(),
+            )
-                                isNullable = true
-                            )
-                        )
-                    ) {
-                        return true
-                    }
-                }
-            }
         }
-        return false
     }
 
     private fun <T : FirBasedSymbol<*>> consumeCallableCandidate(
@@ -287,7 +311,7 @@ class ScopeTowerLevel(
         if (withHideMembersOnly && candidate.getAnnotationByClassId(HidesMembers) == null) {
             return
         }
-        val receiverExpected = withHideMembersOnly || extensionReceiver != null
+        val receiverExpected = withHideMembersOnly || areThereExtensionReceiverOptions()
         if (candidateReceiverTypeRef == null == receiverExpected) return
         val dispatchReceiverValue = dispatchReceiverValue(candidate)
         if (dispatchReceiverValue == null && shouldSkipCandidateWithInconsistentExtensionReceiver(candidate)) {
@@ -297,7 +321,7 @@ class ScopeTowerLevel(
         @Suppress("UNCHECKED_CAST")
         processor.consumeCandidate(
             unwrappedCandidate as T, dispatchReceiverValue,
-            extensionReceiverValue = extensionReceiver,
+            givenExtensionReceiverOptions,
             scope
         )
     }
@@ -343,7 +367,7 @@ class ScopeTowerLevel(
             empty = false
             processor.consumeCandidate(
                 it, dispatchReceiverValue = null,
-                extensionReceiverValue = null,
+                givenExtensionReceiverOptions = emptyList(),
                 scope = scope,
                 objectsByName = true
             )

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirBuilderInferenceSession.kt

@@ -64,7 +64,7 @@ class FirBuilderInferenceSession(
     }
 
     private fun Candidate.isSuitableForBuilderInference(): Boolean {
-        val extensionReceiver = extensionReceiverValue
+        val extensionReceiver = chosenExtensionReceiverValue
         val dispatchReceiver = dispatchReceiverValue
         return when {
             extensionReceiver == null && dispatchReceiver == null -> false

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.builder.buildContextReceiver
 import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
@@ -28,6 +29,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.typeFromCallee
 import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.Name
@@ -259,6 +261,7 @@ class FirCallCompleter(
         override fun analyzeAndGetLambdaReturnArguments(
             lambdaAtom: ResolvedLambdaAtom,
             receiverType: ConeKotlinType?,
+            contextReceivers: List<ConeKotlinType>,
             parameters: List<ConeKotlinType>,
             expectedReturnType: ConeKotlinType?,
             stubsForPostponedVariables: Map<TypeVariableMarker, StubTypeMarker>,
@@ -307,6 +310,18 @@ class FirCallCompleter(
                 }
             )
 
+            if (contextReceivers.isNotEmpty()) {
+                lambdaArgument.replaceContextReceivers(
+                    contextReceivers.map { contextReceiverType ->
+                        buildContextReceiver {
+                            typeRef = buildResolvedTypeRef {
+                                type = contextReceiverType
+                            }
+                        }
+                    }
+                )
+            }
+
             val lookupTracker = session.lookupTracker
             val fileSource = components.file.source
             lambdaArgument.valueParameters.forEachIndexed { index, parameter ->

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirDelegatedPropertyInferenceSession.kt

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubsti
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionContext
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
 import org.jetbrains.kotlin.resolve.calls.inference.model.BuilderInferencePosition
-import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintKind
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImpl
 import org.jetbrains.kotlin.resolve.calls.inference.registerTypeVariableIfNotPresent
@@ -74,7 +73,7 @@ class FirDelegatedPropertyInferenceSession(
         if (callee.candidate.system.hasContradiction) return true
 
         val hasStubType =
-            callee.candidate.extensionReceiverValue?.type?.containsStubType() ?: false
+            callee.candidate.chosenExtensionReceiverValue?.type?.containsStubType() ?: false
                     || callee.candidate.dispatchReceiverValue?.type?.containsStubType() ?: false
 
         if (!hasStubType) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt

@@ -51,10 +51,16 @@ fun extractLambdaInfoFromFunctionalType(
     // For lambdas, the existence of the receiver is always implied by the expected type, and a value parameter
     // can never fill its role.
     val receiverType = if (argument.isLambda) expectedType.receiverType(session) else argument.receiverType
+    val contextReceiversNumber =
+        if (argument.isLambda) expectedType.contextReceiversNumberForFunctionType else argument.contextReceivers.size
+
     val valueParametersTypesIncludingReceiver = expectedType.valueParameterTypesIncludingReceiver(session)
     val isExtensionFunctionType = expectedType.isExtensionFunctionType(session)
     val expectedParameters = valueParametersTypesIncludingReceiver.let {
-        if (receiverType != null && isExtensionFunctionType) it.drop(1) else it
+        val forExtension = if (receiverType != null && isExtensionFunctionType) 1 else 0
+        val toDrop = forExtension + contextReceiversNumber
+
+        if (toDrop > 0) it.drop(toDrop) else it
     }
 
     var coerceFirstParameterToExtensionReceiver = false
@@ -84,11 +90,19 @@ fun extractLambdaInfoFromFunctionalType(
         }
     }
 
+    val contextReceivers =
+        when {
+            contextReceiversNumber == 0 -> emptyList()
+            argument.isLambda -> valueParametersTypesIncludingReceiver.subList(0, contextReceiversNumber)
+            else -> argument.contextReceivers.map { it.typeRef.coneType }
+        }
+
     return ResolvedLambdaAtom(
         argument,
         expectedType,
         expectedType.isSuspendFunctionType(session),
         receiverType,
+        contextReceivers,
         parameters,
         returnType,
         typeVariableForLambdaReturnType = returnTypeVariable,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt

@@ -66,7 +66,8 @@ fun Candidate.preprocessLambdaArgument(
             if (resolvedArgument.coerceFirstParameterToExtensionReceiver) parameters.drop(1) else parameters,
             resolvedArgument.receiver,
             resolvedArgument.returnType,
-            isSuspend = resolvedArgument.isSuspend
+            isSuspend = resolvedArgument.isSuspend,
+            contextReceivers = resolvedArgument.contextReceivers,
         )
 
         val position = ConeArgumentConstraintPosition(resolvedArgument.atom)
@@ -124,6 +125,10 @@ private fun extractLambdaInfo(
         it.returnTypeRef.coneTypeSafe<ConeKotlinType>() ?: nothingType
     }
 
+    val contextReceivers = argument.contextReceivers.map {
+        it.typeRef.coneTypeSafe<ConeKotlinType>() ?: nothingType
+    }
+
     val newTypeVariableUsed = returnType == typeVariable.defaultType
     if (newTypeVariableUsed) csBuilder.registerVariable(typeVariable)
 
@@ -132,6 +137,7 @@ private fun extractLambdaInfo(
         expectedType,
         isSuspend,
         receiverType,
+        contextReceivers,
         parameters,
         returnType,
         typeVariable.takeIf { newTypeVariableUsed },

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt

@@ -37,6 +37,7 @@ interface LambdaAnalyzer {
     fun analyzeAndGetLambdaReturnArguments(
         lambdaAtom: ResolvedLambdaAtom,
         receiverType: ConeKotlinType?,
+        contextReceivers: List<ConeKotlinType>,
         parameters: List<ConeKotlinType>,
         expectedReturnType: ConeKotlinType?, // null means, that return type is not proper i.e. it depends on some type variables
         stubsForPostponedVariables: Map<TypeVariableMarker, StubTypeMarker>,
@@ -122,6 +123,7 @@ class PostponedArgumentsAnalyzer(
         fun substitute(type: ConeKotlinType) = currentSubstitutor.safeSubstitute(c, type) as ConeKotlinType
 
         val receiver = lambda.receiver?.let(::substitute)
+        val contextReceivers = lambda.contextReceivers.map(::substitute)
         val parameters = lambda.parameters.map(::substitute)
         val rawReturnType = lambda.returnType
 
@@ -137,6 +139,7 @@ class PostponedArgumentsAnalyzer(
         val results = lambdaAnalyzer.analyzeAndGetLambdaReturnArguments(
             lambda,
             receiver,
+            contextReceivers,
             parameters,
             expectedTypeForReturnArguments,
             stubsForPostponedVariables,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedAtoms.kt

@@ -41,6 +41,7 @@ class ResolvedLambdaAtom(
     expectedType: ConeKotlinType?,
     val isSuspend: Boolean,
     val receiver: ConeKotlinType?,
+    val contextReceivers: List<ConeKotlinType>,
     val parameters: List<ConeKotlinType>,
     var returnType: ConeKotlinType,
     typeVariableForLambdaReturnType: ConeTypeVariableForLambdaReturnType?,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirTypeResolverImpl.kt

@@ -467,8 +467,8 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
 
     private fun createFunctionalType(typeRef: FirFunctionTypeRef): ConeClassLikeType {
         val parameters =
-            listOfNotNull(typeRef.receiverTypeRef?.coneType) +
+            typeRef.contextReceiverTypeRefs.map { it.coneType } +
-                    typeRef.contextReceiverTypeRefs.map { it.coneType } +
+                    listOfNotNull(typeRef.receiverTypeRef?.coneType) +
                     typeRef.valueParameters.map { it.returnTypeRef.coneType.withParameterNameAnnotation(it) } +
                     listOf(typeRef.returnTypeRef.coneType)
         val classId = if (typeRef.isSuspend) {
@@ -550,7 +550,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
         override val dispatchReceiverValue: ReceiverValue?
             get() = null
 
-        override val extensionReceiverValue: ReceiverValue?
+        override val chosenExtensionReceiverValue: ReceiverValue?
             get() = null
 
         override val explicitReceiverKind: ExplicitReceiverKind

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

@@ -131,7 +131,7 @@ class FirCallCompletionResultsWriterTransformer(
         }
 
         var dispatchReceiver = subCandidate.dispatchReceiverExpression()
-        var extensionReceiver = subCandidate.extensionReceiverExpression()
+        var extensionReceiver = subCandidate.chosenExtensionReceiverExpression()
         if (!declaration.isWrappedIntegerOperator()) {
             val expectedDispatchReceiverType = (declaration as? FirCallableDeclaration)?.dispatchReceiverType
             val expectedExtensionReceiverType = (declaration as? FirCallableDeclaration)?.receiverTypeRef?.coneType
@@ -403,7 +403,7 @@ class FirCallCompletionResultsWriterTransformer(
             StoreCalleeReference,
             resolvedReference,
         ).transformDispatchReceiver(StoreReceiver, subCandidate.dispatchReceiverExpression())
-            .transformExtensionReceiver(StoreReceiver, subCandidate.extensionReceiverExpression())
+            .transformExtensionReceiver(StoreReceiver, subCandidate.chosenExtensionReceiverExpression())
     }
 
     override fun transformVariableAssignment(

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

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.utils.isCompanion
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
-import org.jetbrains.kotlin.fir.declarations.primaryConstructorIfAny
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.fir.resolve.*
@@ -228,6 +227,8 @@ class BodyResolveContext(
         holder: SessionHolder,
         f: () -> T
     ): T = withTowerDataCleanup {
+        replaceTowerDataContext(towerDataContext.addContextReceiverGroup(owner.createContextReceiverValues(holder)))
+
         if (type != null) {
             val receiver = ImplicitExtensionReceiverValue(
                 owner.symbol,
@@ -450,6 +451,7 @@ class BodyResolveContext(
         val forMembersResolution =
             staticsAndCompanion
                 .addReceiver(labelName, towerElementsForClass.thisReceiver)
+                .addContextReceiverGroup(towerElementsForClass.contextReceivers)
                 .addNonLocalScopeIfNotNull(typeParameterScope)
 
         val scopeForConstructorHeader =

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fakeElement
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildAnonymousFunctionCopy
+import org.jetbrains.kotlin.fir.declarations.builder.buildContextReceiver
 import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
@@ -34,7 +35,6 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeLocalVariableNoTypeOrIni
 import org.jetbrains.kotlin.fir.resolve.inference.FirStubTypeTransformer
 import org.jetbrains.kotlin.fir.resolve.inference.ResolvedLambdaAtom
 import org.jetbrains.kotlin.fir.resolve.inference.extractLambdaInfoFromFunctionalType
-import org.jetbrains.kotlin.fir.types.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.createTypeSubstitutorByTypeConstructor
 import org.jetbrains.kotlin.fir.resolve.transformers.*
@@ -809,6 +809,19 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         lambda = buildAnonymousFunctionCopy(lambda) {
             receiverTypeRef = lambda.receiverTypeRef?.takeIf { it !is FirImplicitTypeRef }
                 ?: resolvedLambdaAtom?.receiver?.let { lambda.receiverTypeRef?.resolvedTypeFromPrototype(it) }
+
+            contextReceivers.clear()
+            contextReceivers.addAll(
+                lambda.contextReceivers.takeIf { it.isNotEmpty() }
+                    ?: resolvedLambdaAtom?.contextReceivers?.map { receiverType ->
+                        buildContextReceiver {
+                            this.typeRef = buildResolvedTypeRef {
+                                type = receiverType
+                            }
+                        }
+                    }.orEmpty()
+            )
+
             this.valueParameters.clear()
             this.valueParameters.addAll(valueParameters)
             returnTypeRef = (lambda.returnTypeRef as? FirResolvedTypeRef)

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

@@ -8,14 +8,14 @@ package org.jetbrains.kotlin.fir.declarations
 import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.toPersistentList
+import org.jetbrains.kotlin.fir.labelName
 import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.*
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.scopes.impl.wrapNestedClassifierScopeWithSubstitutionForSuperType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addIfNotNull
 
@@ -23,32 +23,39 @@ fun SessionHolder.collectImplicitReceivers(
     type: ConeKotlinType?,
     owner: FirDeclaration
 ): ImplicitReceivers {
-    if (type == null) return ImplicitReceivers(null, emptyList())
-
     val implicitCompanionValues = mutableListOf<ImplicitReceiverValue<*>>()
+    val contextReceiverValues = mutableListOf<ContextReceiverValue<*>>()
     val implicitReceiverValue = when (owner) {
         is FirClass -> {
-            val towerElementsForClass = collectTowerDataElementsForClass(owner, type)
+            val towerElementsForClass = collectTowerDataElementsForClass(owner, type!!)
             implicitCompanionValues.addAll(towerElementsForClass.implicitCompanionValues)
+            contextReceiverValues.addAll(towerElementsForClass.contextReceivers)
 
             towerElementsForClass.thisReceiver
         }
         is FirFunction -> {
-            ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+            contextReceiverValues.addAll(owner.createContextReceiverValues(this))
+            type?.let { ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession) }
         }
         is FirVariable -> {
-            ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
+            contextReceiverValues.addAll(owner.createContextReceiverValues(this))
+            type?.let { ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession) }
         }
         else -> {
-            throw IllegalArgumentException("Incorrect label & receiver owner: ${owner.javaClass}")
+            if (type != null) {
+                throw IllegalArgumentException("Incorrect label & receiver owner: ${owner.javaClass}")
+            }
+
+            null
         }
     }
-    return ImplicitReceivers(implicitReceiverValue, implicitCompanionValues)
+    return ImplicitReceivers(implicitReceiverValue, implicitCompanionValues, contextReceiverValues)
 }
 
 data class ImplicitReceivers(
     val implicitReceiverValue: ImplicitReceiverValue<*>?,
-    val implicitCompanionValues: List<ImplicitReceiverValue<*>>
+    val implicitCompanionValues: List<ImplicitReceiverValue<*>>,
+    val contextReceivers: List<ContextReceiverValue<*>>,
 )
 
 fun SessionHolder.collectTowerDataElementsForClass(owner: FirClass, defaultType: ConeKotlinType): TowerElementsForClass {
@@ -83,9 +90,15 @@ fun SessionHolder.collectTowerDataElementsForClass(owner: FirClass, defaultType:
     }
 
     val thisReceiver = ImplicitDispatchReceiverValue(owner.symbol, defaultType, session, scopeSession)
+    val contextReceivers = (owner as? FirRegularClass)?.contextReceivers?.map {
+        ContextReceiverValueForClass(
+            owner.symbol, it.typeRef.coneType, it.labelName, session, scopeSession,
+        )
+    }.orEmpty()
 
     return TowerElementsForClass(
         thisReceiver,
+        contextReceivers,
         owner.staticScope(this),
         companionReceiver,
         companionObject?.staticScope(this),
@@ -96,6 +109,7 @@ fun SessionHolder.collectTowerDataElementsForClass(owner: FirClass, defaultType:
 
 class TowerElementsForClass(
     val thisReceiver: ImplicitReceiverValue<*>,
+    val contextReceivers: List<ContextReceiverValueForClass>,
     val staticScope: FirScope?,
     val companionReceiver: ImplicitReceiverValue<*>?,
     val companionStaticScope: FirScope?,
@@ -137,7 +151,9 @@ class FirTowerDataContext private constructor(
     fun addNonLocalTowerDataElements(newElements: List<FirTowerDataElement>): FirTowerDataContext {
         return FirTowerDataContext(
             towerDataElements.addAll(newElements),
-            implicitReceiverStack.addAll(newElements.mapNotNull { it.implicitReceiver }),
+            implicitReceiverStack
+                .addAll(newElements.mapNotNull { it.implicitReceiver })
+                .addAllContextReceivers(newElements.flatMap { it.contextReceiverGroup.orEmpty() }),
             localScopes,
             nonLocalTowerDataElements.addAll(newElements)
         )
@@ -162,6 +178,18 @@ class FirTowerDataContext private constructor(
         )
     }
 
+    fun addContextReceiverGroup(contextReceiverGroup: ContextReceiverGroup): FirTowerDataContext {
+        if (contextReceiverGroup.isEmpty()) return this
+        val element = contextReceiverGroup.asTowerDataElement()
+
+        return FirTowerDataContext(
+            towerDataElements.add(element),
+            contextReceiverGroup.fold(implicitReceiverStack, PersistentImplicitReceiverStack::addContextReceiver),
+            localScopes,
+            nonLocalTowerDataElements.add(element)
+        )
+    }
+
     fun addNonLocalScopeIfNotNull(scope: FirScope?): FirTowerDataContext {
         if (scope == null) return this
         return addNonLocalScope(scope)
@@ -179,23 +207,47 @@ class FirTowerDataContext private constructor(
 
     fun createSnapshot(): FirTowerDataContext {
         return FirTowerDataContext(
-            towerDataElements.map { FirTowerDataElement(it.scope, it.implicitReceiver?.createSnapshot(), it.isLocal) }.toPersistentList(),
+            towerDataElements.map(FirTowerDataElement::createSnapshot).toPersistentList(),
             implicitReceiverStack.createSnapshot(),
             localScopes.toPersistentList(),
-            nonLocalTowerDataElements.map { FirTowerDataElement(it.scope, it.implicitReceiver?.createSnapshot(), it.isLocal) }
+            nonLocalTowerDataElements.map(FirTowerDataElement::createSnapshot).toPersistentList()
-                .toPersistentList()
         )
     }
 }
 
-class FirTowerDataElement(val scope: FirScope?, val implicitReceiver: ImplicitReceiverValue<*>?, val isLocal: Boolean)
+class FirTowerDataElement(
+    val scope: FirScope?,
+    val implicitReceiver: ImplicitReceiverValue<*>?,
+    val contextReceiverGroup: ContextReceiverGroup? = null,
+    val isLocal: Boolean,
+) {
+    fun createSnapshot(): FirTowerDataElement =
+        FirTowerDataElement(
+            scope,
+            implicitReceiver?.createSnapshot(),
+            contextReceiverGroup?.map { it.createSnapshot() },
+            isLocal,
+        )
+}
 
 fun ImplicitReceiverValue<*>.asTowerDataElement(): FirTowerDataElement =
-    FirTowerDataElement(scope = null, this, isLocal = false)
+    FirTowerDataElement(scope = null, implicitReceiver = this, isLocal = false)
+
+fun ContextReceiverGroup.asTowerDataElement(): FirTowerDataElement =
+    FirTowerDataElement(scope = null, implicitReceiver = null, contextReceiverGroup = this, isLocal = false)
 
 fun FirScope.asTowerDataElement(isLocal: Boolean): FirTowerDataElement =
-    FirTowerDataElement(this, implicitReceiver = null, isLocal)
+    FirTowerDataElement(scope = this, implicitReceiver = null, isLocal = isLocal)
 
 fun FirClass.staticScope(sessionHolder: SessionHolder) =
     scopeProvider.getStaticScope(this, sessionHolder.session, sessionHolder.scopeSession)
+
+typealias ContextReceiverGroup = List<ContextReceiverValue<*>>
 typealias FirLocalScopes = PersistentList<FirLocalScope>
+
+fun FirCallableDeclaration.createContextReceiverValues(
+    sessionHolder: SessionHolder,
+): List<ContextReceiverValueForCallable> =
+    contextReceivers.map {
+        ContextReceiverValueForCallable(symbol, it.typeRef.coneType, it.labelName, sessionHolder.session, sessionHolder.scopeSession)
+    }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/PersistentImplicitReceiverStack.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import kotlinx.collections.immutable.*
+import org.jetbrains.kotlin.fir.resolve.calls.ContextReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -16,7 +17,7 @@ import org.jetbrains.kotlin.name.Name
 class PersistentImplicitReceiverStack private constructor(
     private val stack: PersistentList<ImplicitReceiverValue<*>>,
     // This multi-map holds indexes of the stack ^
-    private val indexesPerLabel: PersistentSetMultimap<Name, Int>,
+    private val receiversPerLabel: PersistentSetMultimap<Name, ImplicitReceiverValue<*>>,
     private val indexesPerSymbol: PersistentMap<FirBasedSymbol<*>, Int>,
     private val originalTypes: PersistentList<ConeKotlinType>,
 ) : ImplicitReceiverStack(), Iterable<ImplicitReceiverValue<*>> {
@@ -33,16 +34,42 @@ class PersistentImplicitReceiverStack private constructor(
         return receivers.fold(this) { acc, value -> acc.add(name = null, value) }
     }
 
-    fun add(name: Name?, value: ImplicitReceiverValue<*>): PersistentImplicitReceiverStack {
+    fun addAllContextReceivers(receivers: List<ContextReceiverValue<*>>): PersistentImplicitReceiverStack {
+        return receivers.fold(this) { acc, value -> acc.addContextReceiver(value) }
+    }
+
+    fun add(name: Name?, value: ImplicitReceiverValue<*>, aliasLabel: Name? = null): PersistentImplicitReceiverStack {
         val stack = stack.add(value)
         val originalTypes = originalTypes.add(value.originalType)
         val index = stack.size - 1
-        val indexesPerLabel = name?.let { indexesPerLabel.put(it, index) } ?: indexesPerLabel
+        val receiversPerLabel = name?.let { receiversPerLabel.put(it, value) } ?: receiversPerLabel
         val indexesPerSymbol = indexesPerSymbol.put(value.boundSymbol, index)
 
         return PersistentImplicitReceiverStack(
             stack,
-            indexesPerLabel,
+            receiversPerLabel,
+            indexesPerSymbol,
+            originalTypes
+        )
+    }
+
+    fun addContextReceiver(value: ContextReceiverValue<*>): PersistentImplicitReceiverStack {
+        val labelName = value.labelName ?: return this
+
+        val receiversPerLabel = receiversPerLabel.put(labelName, value)
+        return PersistentImplicitReceiverStack(
+            stack,
+            receiversPerLabel,
+            indexesPerSymbol,
+            originalTypes
+        )
+    }
+
+    fun addReceiverLabelAlias(aliasLabel: Name, value: ImplicitReceiverValue<*>): PersistentImplicitReceiverStack {
+        val receiversPerLabel = receiversPerLabel.put(aliasLabel, value)
+        return PersistentImplicitReceiverStack(
+            stack,
+            receiversPerLabel,
             indexesPerSymbol,
             originalTypes
         )
@@ -50,7 +77,7 @@ class PersistentImplicitReceiverStack private constructor(
 
     override operator fun get(name: String?): ImplicitReceiverValue<*>? {
         if (name == null) return stack.lastOrNull()
-        return indexesPerLabel[Name.identifier(name)].lastOrNull()?.let { stack[it] }
+        return receiversPerLabel[Name.identifier(name)].lastOrNull()
     }
 
     override fun lastDispatchReceiver(): ImplicitDispatchReceiverValue? {
@@ -85,7 +112,7 @@ class PersistentImplicitReceiverStack private constructor(
     fun createSnapshot(): PersistentImplicitReceiverStack {
         return PersistentImplicitReceiverStack(
             stack.map { it.createSnapshot() }.toPersistentList(),
-            indexesPerLabel,
+            receiversPerLabel,
             indexesPerSymbol,
             originalTypes
         )

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/calls/AbstractCandidate.kt

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 abstract class AbstractCandidate {
     abstract val symbol: FirBasedSymbol<*>
     abstract val dispatchReceiverValue: ReceiverValue?
-    abstract val extensionReceiverValue: ReceiverValue?
+    abstract val chosenExtensionReceiverValue: ReceiverValue?
     abstract val explicitReceiverKind: ExplicitReceiverKind
     abstract val callInfo: AbstractCallInfo
     abstract val diagnostics: List<ResolutionDiagnostic>

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/calls/ResolutionDiagnostic.kt

@@ -82,7 +82,8 @@ object LowerPriorityToPreserveCompatibilityDiagnostic : ResolutionDiagnostic(RES
 
 object CandidateChosenUsingOverloadResolutionByLambdaAnnotation : ResolutionDiagnostic(RESOLVED)
 
-class UnstableSmartCast(val argument: FirExpressionWithSmartcast, val targetType: ConeKotlinType, val isCastToNotNull: Boolean) : ResolutionDiagnostic(UNSTABLE_SMARTCAST)
+class UnstableSmartCast(val argument: FirExpressionWithSmartcast, val targetType: ConeKotlinType, val isCastToNotNull: Boolean) :
+    ResolutionDiagnostic(UNSTABLE_SMARTCAST)
 
 class ArgumentTypeMismatch(
     val expectedType: ConeKotlinType,
@@ -108,3 +109,13 @@ class Unsupported(val message: String, val source: KtSourceElement? = null) : Re
 object PropertyAsOperator : ResolutionDiagnostic(PROPERTY_AS_OPERATOR)
 
 class DslScopeViolation(val calleeSymbol: FirBasedSymbol<*>) : ResolutionDiagnostic(DSL_SCOPE_VIOLATION)
+
+class MultipleContextReceiversApplicableForExtensionReceivers : ResolutionDiagnostic(INAPPLICABLE)
+
+class NoApplicableValueForContextReceiver(
+    val expectedContextReceiverType: ConeKotlinType
+) : ResolutionDiagnostic(INAPPLICABLE)
+
+class AmbiguousValuesForContextReceiverParameter(
+    val expectedContextReceiverType: ConeKotlinType,
+) : ResolutionDiagnostic(INAPPLICABLE)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirTypeUtils.kt

@@ -81,9 +81,9 @@ val FirTypeRef.isMarkedNullable: Boolean?
 
 val FirFunctionTypeRef.parametersCount: Int
     get() = if (receiverTypeRef != null)
-        valueParameters.size + 1
+        valueParameters.size + contextReceiverTypeRefs.size + 1
     else
-        valueParameters.size
+        valueParameters.size + contextReceiverTypeRefs.size
 
 val EXTENSION_FUNCTION_ANNOTATION = ClassId.fromString("kotlin/ExtensionFunctionType")