[PSI, FE] Support functional types

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -10588,6 +10588,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("functionalType.kt")
+                public void testFunctionalType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.kt");
+                }
+
                 @Test
                 @TestMetadata("genericOuterClass.kt")
                 public void testGenericOuterClass() throws Exception {
@@ -10719,6 +10725,28 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 public void testWithExplicitReceiverError() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/withExplicitReceiverError.kt");
                 }
+
+                @Nested
+                @TestMetadata("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP")
+                @TestDataPath("$PROJECT_ROOT")
+                public class FromKEEP {
+                    @Test
+                    public void testAllFilesPresentInFromKEEP() throws Exception {
+                        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                    }
+
+                    @Test
+                    @TestMetadata("clickHandler.kt")
+                    public void testClickHandler() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.kt");
+                    }
+
+                    @Test
+                    @TestMetadata("functionalType.kt")
+                    public void testFunctionalType() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.kt");
+                    }
+                }
             }
         }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/JvmRuntimeTypes.kt

@@ -100,6 +100,7 @@ class JvmRuntimeTypes(
             descriptor.builtIns,
             Annotations.EMPTY,
             actualFunctionDescriptor.extensionReceiverParameter?.type,
+            actualFunctionDescriptor.contextReceiverParameters.map { it.type },
             actualFunctionDescriptor.valueParameters.map { it.type },
             null,
             actualFunctionDescriptor.returnType!!
@@ -146,6 +147,7 @@ class JvmRuntimeTypes(
             Annotations.EMPTY,
             if (isBound) null else referencedFunction.extensionReceiverParameter?.type
                 ?: referencedFunction.dispatchReceiverParameter?.type,
+            referencedFunction.contextReceiverParameters.map { it.type },
             anonymousFunctionDescriptor.valueParameters.drop(receivers).map { it.type },
             null,
             anonymousFunctionDescriptor.returnType!!,

compiler/backend/src/org/jetbrains/kotlin/codegen/callableReferenceUtil.kt

@@ -94,7 +94,8 @@ fun InstructionAdapter.generateClosureFieldsInitializationFromParameters(
 }
 
 fun computeExpectedNumberOfReceivers(referencedFunction: FunctionDescriptor, isBound: Boolean): Int {
-    val receivers = (if (referencedFunction.dispatchReceiverParameter != null) 1 else 0) +
+    val receivers = referencedFunction.contextReceiverParameters.size +
+            (if (referencedFunction.dispatchReceiverParameter != null) 1 else 0) +
             (if (referencedFunction.extensionReceiverParameter != null) 1 else 0) -
             (if (isBound) 1 else 0)
 

compiler/fir/analysis-tests/testData/builtIns/kotlin.txt

@@ -339,6 +339,14 @@ public abstract interface Comparable<in T> : R|kotlin/Any| {
 
 }
 
+@R|kotlin/annotation/Target|(allowedTargets = <implicitArrayOf>(R|kotlin/annotation/AnnotationTarget.TYPE|())) @R|kotlin/annotation/MustBeDocumented|() public final annotation class ContextFunctionTypeParams : R|kotlin/Annotation| {
+    public final val count: R|kotlin/Int|
+        public get(): R|kotlin/Int|
+
+    public constructor(count: R|kotlin/Int|): R|kotlin/ContextFunctionTypeParams|
+
+}
+
 @R|kotlin/annotation/Target|(allowedTargets = <implicitArrayOf>(R|kotlin/annotation/AnnotationTarget.CLASS|(), R|kotlin/annotation/AnnotationTarget.FUNCTION|(), R|kotlin/annotation/AnnotationTarget.PROPERTY|(), R|kotlin/annotation/AnnotationTarget.ANNOTATION_CLASS|(), R|kotlin/annotation/AnnotationTarget.CONSTRUCTOR|(), R|kotlin/annotation/AnnotationTarget.PROPERTY_SETTER|(), R|kotlin/annotation/AnnotationTarget.PROPERTY_GETTER|(), R|kotlin/annotation/AnnotationTarget.TYPEALIAS|())) @R|kotlin/annotation/MustBeDocumented|() public final annotation class Deprecated : R|kotlin/Annotation| {
     public final val level: R|kotlin/DeprecationLevel|
         public get(): R|kotlin/DeprecationLevel|

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -10588,6 +10588,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("functionalType.kt")
+                public void testFunctionalType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.kt");
+                }
+
                 @Test
                 @TestMetadata("genericOuterClass.kt")
                 public void testGenericOuterClass() throws Exception {
@@ -10719,6 +10725,28 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 public void testWithExplicitReceiverError() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/withExplicitReceiverError.kt");
                 }
+
+                @Nested
+                @TestMetadata("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP")
+                @TestDataPath("$PROJECT_ROOT")
+                public class FromKEEP {
+                    @Test
+                    public void testAllFilesPresentInFromKEEP() throws Exception {
+                        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                    }
+
+                    @Test
+                    @TestMetadata("clickHandler.kt")
+                    public void testClickHandler() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.kt");
+                    }
+
+                    @Test
+                    @TestMetadata("functionalType.kt")
+                    public void testFunctionalType() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.kt");
+                    }
+                }
             }
         }
 

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -10588,6 +10588,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("functionalType.kt")
+                public void testFunctionalType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.kt");
+                }
+
                 @Test
                 @TestMetadata("genericOuterClass.kt")
                 public void testGenericOuterClass() throws Exception {
@@ -10719,6 +10725,28 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 public void testWithExplicitReceiverError() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/withExplicitReceiverError.kt");
                 }
+
+                @Nested
+                @TestMetadata("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP")
+                @TestDataPath("$PROJECT_ROOT")
+                public class FromKEEP {
+                    @Test
+                    public void testAllFilesPresentInFromKEEP() throws Exception {
+                        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                    }
+
+                    @Test
+                    @TestMetadata("clickHandler.kt")
+                    public void testClickHandler() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.kt");
+                    }
+
+                    @Test
+                    @TestMetadata("functionalType.kt")
+                    public void testFunctionalType() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.kt");
+                    }
+                }
             }
         }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/TypeResolver.kt

@@ -346,7 +346,12 @@ class TypeResolver(
 
             override fun visitFunctionType(type: KtFunctionType) {
                 val receiverTypeRef = type.receiverTypeReference
-                val receiverType = if (receiverTypeRef == null) null else resolveType(c.noBareTypes(), receiverTypeRef)
+                val receiverType = if (receiverTypeRef?.typeElement == null) null else resolveType(c.noBareTypes(), receiverTypeRef)
+
+                val contextReceiversTypeRefs = type.contextReceiversTypeReferences
+                val contextReceiversTypes = contextReceiversTypeRefs?.mapNotNull {
+                    resolveType(c.noBareTypes(), it)
+                } ?: emptyList()
 
                 val parameterDescriptors = resolveParametersOfFunctionType(type.parameters)
                 checkParametersOfFunctionType(parameterDescriptors)
@@ -362,7 +367,7 @@ class TypeResolver(
 
                 result = type(
                     createFunctionType(
-                        moduleDescriptor.builtIns, annotations, receiverType,
+                        moduleDescriptor.builtIns, annotations, receiverType, contextReceiversTypes,
                         parameterDescriptors.map { it.type },
                         parameterDescriptors.map { it.name },
                         returnType,

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/ArgumentTypeResolver.java

@@ -44,6 +44,7 @@ import javax.inject.Inject;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import static org.jetbrains.kotlin.psi.KtPsiUtil.getLastElementDeparenthesized;
 import static org.jetbrains.kotlin.resolve.BindingContextUtils.getRecordedTypeInfo;
@@ -334,7 +335,7 @@ public class ArgumentTypeResolver {
         }
 
         return FunctionTypesKt.createFunctionType(
-                builtIns, Annotations.Companion.getEMPTY(), null, Collections.emptyList(), null, TypeUtils.DONT_CARE
+                builtIns, Annotations.Companion.getEMPTY(), null, Collections.emptyList(), Collections.emptyList(), null, TypeUtils.DONT_CARE
         );
     }
 
@@ -372,7 +373,7 @@ public class ArgumentTypeResolver {
             return expectedTypeIsUnknown
                    ? functionPlaceholders.createFunctionPlaceholderType(Collections.emptyList(), /* hasDeclaredArguments = */ false)
                    : FunctionTypesKt.createFunctionType(
-                           builtIns, Annotations.Companion.getEMPTY(), null, Collections.emptyList(), null, DONT_CARE
+                           builtIns, Annotations.Companion.getEMPTY(), null, Collections.emptyList(), Collections.emptyList(), null, DONT_CARE
                    );
         }
         List<KtParameter> valueParameters = function.getValueParameters();
@@ -391,11 +392,14 @@ public class ArgumentTypeResolver {
         KotlinType returnType = resolveTypeRefWithDefault(function.getTypeReference(), scope, temporaryTrace, DONT_CARE);
         assert returnType != null;
         KotlinType receiverType = resolveTypeRefWithDefault(function.getReceiverTypeReference(), scope, temporaryTrace, null);
+        List<KotlinType> contextReceiversTypes = function.getContextReceivers().stream().map(contextReceiver ->
+            resolveTypeRefWithDefault(contextReceiver.typeReference(), scope, temporaryTrace, null)
+        ).collect(Collectors.toList());
 
         return expectedTypeIsUnknown && isFunctionLiteral
                ? functionPlaceholders.createFunctionPlaceholderType(parameterTypes, /* hasDeclaredArguments = */ true)
                : FunctionTypesKt.createFunctionType(
-                       builtIns, Annotations.Companion.getEMPTY(), receiverType, parameterTypes, parameterNames, returnType, suspendFunctionTypeExpected
+                       builtIns, Annotations.Companion.getEMPTY(), receiverType, contextReceiversTypes, parameterTypes, parameterNames, returnType, suspendFunctionTypeExpected
                );
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallResolver.java

@@ -357,7 +357,7 @@ public class CallResolver {
                 parameterTypes.add(NO_EXPECTED_TYPE);
             }
             expectedType = FunctionTypesKt.createFunctionType(
-                    builtIns, Annotations.Companion.getEMPTY(), null, parameterTypes, null, context.expectedType
+                    builtIns, Annotations.Companion.getEMPTY(), null, Collections.emptyList(), parameterTypes, null, context.expectedType
             );
         }
         KotlinType calleeType = expressionTypingServices.safeGetType(

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/BuilderInferenceUtil.kt

@@ -177,7 +177,7 @@ class BuilderInferenceSupport(
             approximationSubstitutor.buildSubstitutor().substitute(lambdaExpectedType, Variance.IN_VARIANCE) ?: return
 
         val newExpectedType = createFunctionType(
-            newReceiverType.builtIns, approximatedLambdaType.annotations, newReceiverType,
+            newReceiverType.builtIns, approximatedLambdaType.annotations, newReceiverType, emptyList(), // TODO: Context receivers?
             approximatedLambdaType.getValueParameterTypesFromFunctionType().map(TypeProjection::getType),
             parameterNames = null, // TODO: parameterNames
             returnType = approximatedLambdaType.getReturnTypeFromFunctionType(),

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/ConstraintSystemBuilderImpl.kt

@@ -16,9 +16,7 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference
 
-import org.jetbrains.kotlin.builtins.createFunctionType
+import org.jetbrains.kotlin.builtins.*
-import org.jetbrains.kotlin.builtins.isBuiltinExtensionFunctionalType
-import org.jetbrains.kotlin.builtins.isSuspendFunctionType
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilderImpl.ConstraintKind.EQUAL
@@ -471,8 +469,9 @@ internal fun createTypeForFunctionPlaceholder(
         functionPlaceholderTypeConstructor.argumentTypes
     }
     val receiverType = if (isExtension) DONT_CARE else null
+    val contextReceiverTypes = (0 until expectedType.contextFunctionTypeParamsCount()).map { DONT_CARE }
     return createFunctionType(
-        functionPlaceholder.builtIns, Annotations.EMPTY, receiverType, newArgumentTypes, null, DONT_CARE,
+        functionPlaceholder.builtIns, Annotations.EMPTY, receiverType, contextReceiverTypes, newArgumentTypes, null, DONT_CARE,
         suspendFunction = expectedType.isSuspendFunctionType
     )
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tasks/dynamicCalls.kt

@@ -186,9 +186,11 @@ class DynamicCallableDescriptors(private val storageManager: StorageManager, bui
             val funLiteral = funLiteralExpr.functionLiteral
 
             val receiverType = funLiteral.receiverTypeReference?.let { dynamicType }
+            val contextReceiversTypes = funLiteral.contextReceivers.map { dynamicType }
+
             val parameterTypes = funLiteral.valueParameters.map { dynamicType }
 
-            return createFunctionType(owner.builtIns, Annotations.EMPTY, receiverType, parameterTypes, null, dynamicType)
+            return createFunctionType(owner.builtIns, Annotations.EMPTY, receiverType, contextReceiversTypes, parameterTypes, null, dynamicType)
         }
 
         for (arg in call.valueArguments) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt

@@ -171,8 +171,9 @@ class KotlinResolutionCallbacksImpl(
             @OptIn(TypeRefinement::class) callComponents.kotlinTypeChecker.kotlinTypeRefiner.refineType(it)
         }
 
+        // TODO: Context receivers?
         val expectedType = createFunctionType(
-            builtIns, annotations, refinedReceiverType, parameters, null,
+            builtIns, annotations, refinedReceiverType, emptyList(), parameters, null,
             lambdaInfo.expectedType, isSuspend
         )
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt

@@ -334,6 +334,7 @@ class ResolvedAtomCompleter(
             builtIns,
             existingLambdaType.annotations,
             substitutedLambdaTypes.receiverType?.substitutedType,
+            emptyList(),
             substitutedLambdaTypes.parameterTypes.map { it.substitutedType },
             null, // parameter names transforms to special annotations, so they are already taken from parameter types
             substitutedLambdaTypes.returnType.substitutedType,

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/DoubleColonExpressionResolver.kt

@@ -830,6 +830,9 @@ class DoubleColonExpressionResolver(
     }
 
     companion object {
+        private fun contextReceiverTypesFor(descriptor: CallableDescriptor): List<KotlinType> =
+            descriptor.contextReceiverParameters.map { it.type }
+
         private fun receiverTypeFor(descriptor: CallableDescriptor, lhs: DoubleColonLHS?): KotlinType? =
             (descriptor.extensionReceiverParameter ?: descriptor.dispatchReceiverParameter)?.let { (lhs as? DoubleColonLHS.Type)?.type }
 
@@ -849,6 +852,7 @@ class DoubleColonExpressionResolver(
             reflectionTypes: ReflectionTypes,
             scopeOwnerDescriptor: DeclarationDescriptor
         ): KotlinType? {
+            val contextReceiverTypes = contextReceiverTypesFor(descriptor)
             val receiverType = receiverTypeFor(descriptor, lhs)
             return when (descriptor) {
                 is FunctionDescriptor -> {
@@ -857,7 +861,7 @@ class DoubleColonExpressionResolver(
                     val parametersNames = descriptor.valueParameters.map { it.name }
                     return reflectionTypes.getKFunctionType(
                         Annotations.EMPTY, receiverType,
-                        parametersTypes, parametersNames, returnType, descriptor.builtIns, descriptor.isSuspend
+                        contextReceiverTypes, parametersTypes, parametersNames, returnType, descriptor.builtIns, descriptor.isSuspend
                     )
                 }
                 is PropertyDescriptor -> {

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt

@@ -397,6 +397,7 @@ fun SimpleFunctionDescriptor.createFunctionType(
         builtIns,
         Annotations.EMPTY,
         extensionReceiverParameter?.type,
+        contextReceiverParameters.map { it.type },
         if (shouldUseVarargType) valueParameters.map { it.varargElementType ?: it.type } else valueParameters.map { it.type },
         null,
         returnType ?: return null,

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ReflectionReferencesGenerator.kt

@@ -540,6 +540,7 @@ class ReflectionReferencesGenerator(statementGenerator: StatementGenerator) : St
         (statementGenerator.context.irBuiltIns as IrBuiltInsOverDescriptors).builtIns,
         annotations,
         null,
+        emptyList(),
         arguments.dropLast(1).map { it.type },
         null,
         arguments.last().type,

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinParsing.java

@@ -2172,6 +2172,12 @@ public class KotlinParsing extends AbstractKotlinParsing {
         IElementType lookahead = lookahead(1);
         IElementType lookahead2 = lookahead(2);
         boolean typeBeforeDot = true;
+        boolean withContextReceiver = at(CONTEXT_KEYWORD) && lookahead == LPAR;
+
+        if (withContextReceiver) {
+            parseContextReceiverList();
+        }
+
         if (at(IDENTIFIER) && !(lookahead == DOT && lookahead2 == IDENTIFIER) && lookahead != LT && at(DYNAMIC_KEYWORD)) {
             PsiBuilder.Marker dynamicType = mark();
             advance(); // DYNAMIC_KEYWORD
@@ -2180,7 +2186,7 @@ public class KotlinParsing extends AbstractKotlinParsing {
         else if (at(IDENTIFIER) || at(PACKAGE_KEYWORD) || atParenthesizedMutableForPlatformTypes(0)) {
             parseUserType(allowSimpleIntersectionTypes);
         }
-        else if (at(LPAR)) {
+        else if (at(LPAR) && !withContextReceiver) {
             PsiBuilder.Marker functionOrParenthesizedType = mark();
 
             // This may be a function parameter list or just a parenthesized type
@@ -2211,7 +2217,7 @@ public class KotlinParsing extends AbstractKotlinParsing {
             }
 
         }
-        else {
+        else if (!withContextReceiver) {
             errorWithRecovery("Type expected",
                     TokenSet.orSet(TOP_LEVEL_DECLARATION_FIRST,
                                    TokenSet.create(EQ, COMMA, GT, RBRACKET, DOT, RPAR, RBRACE, LBRACE, SEMICOLON),
@@ -2240,8 +2246,10 @@ public class KotlinParsing extends AbstractKotlinParsing {
             intersectionType.done(INTERSECTION_TYPE);
             wasIntersection = true;
         }
+        boolean withExtensionReceiver = typeBeforeDot && at(DOT);
 
-        if (typeBeforeDot && !wasIntersection && at(DOT)) {
+
+        if (withExtensionReceiver && !wasIntersection || withContextReceiver) {
             // This is a receiver for a function type
             //  A.(B) -> C
             //   ^
@@ -2253,7 +2261,9 @@ public class KotlinParsing extends AbstractKotlinParsing {
             receiverTypeRef.done(TYPE_REFERENCE);
             receiverType.done(FUNCTION_TYPE_RECEIVER);
 
-            advance(); // DOT
+            if (withExtensionReceiver) {
+                advance(); // DOT
+            }
 
             if (at(LPAR)) {
                 parseFunctionTypeContents().drop();

compiler/psi/src/org/jetbrains/kotlin/psi/KtFunctionType.java

@@ -45,6 +45,10 @@ public class KtFunctionType extends KtElementImplStub<KotlinPlaceHolderStub<KtFu
     @Override
     public List<KtTypeReference> getTypeArgumentsAsTypes() {
         ArrayList<KtTypeReference> result = Lists.newArrayList();
+        List<KtTypeReference> contextReceiversTypeRefs = getContextReceiversTypeReferences();
+        if (contextReceiversTypeRefs != null) {
+            result.addAll(contextReceiversTypeRefs);
+        }
         KtTypeReference receiverTypeRef = getReceiverTypeReference();
         if (receiverTypeRef != null) {
             result.add(receiverTypeRef);
@@ -89,6 +93,20 @@ public class KtFunctionType extends KtElementImplStub<KotlinPlaceHolderStub<KtFu
         return receiverDeclaration.getTypeReference();
     }
 
+    public List<KtTypeReference> getContextReceiversTypeReferences() {
+        KtFunctionTypeReceiver receiverDeclaration = getReceiver();
+        if (receiverDeclaration == null) {
+            return null;
+        }
+        KtTypeReference receiverTypeRef = receiverDeclaration.getTypeReference();
+        KtContextReceiverList contextReceiverList = receiverTypeRef.getStubOrPsiChild(KtStubElementTypes.CONTEXT_RECEIVER_LIST);
+        if (contextReceiverList != null) {
+            return contextReceiverList.typeReferences();
+        } else {
+            return Collections.emptyList();
+        }
+    }
+
     @Nullable
     public KtTypeReference getReturnTypeReference() {
         return getStubOrPsiChild(KtStubElementTypes.TYPE_REFERENCE);

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponeArgumentsChecks.kt

@@ -86,7 +86,7 @@ private fun preprocessLambdaArgument(
 
     if (expectedType != null) {
         val lambdaType = createFunctionType(
-            csBuilder.builtIns, Annotations.EMPTY, resolvedArgument.receiver,
+            csBuilder.builtIns, Annotations.EMPTY, resolvedArgument.receiver, emptyList(), // TODO: Context receivers?
             resolvedArgument.parameters, null, resolvedArgument.returnType, resolvedArgument.isSuspend
         )
         csBuilder.addSubtypeConstraint(lambdaType, expectedType, ArgumentConstraintPositionImpl(argument))

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/SuspendConversionUtils.kt

@@ -43,6 +43,7 @@ object SuspendTypeConversions : ParameterTypeConversion {
             candidate.callComponents.builtIns,
             expectedParameterType.annotations,
             expectedParameterType.getReceiverTypeFromFunctionType(),
+            expectedParameterType.getContextReceiverTypesFromFunctionType(),
             expectedParameterType.getValueParameterTypesFromFunctionType().map { it.type },
             parameterNames = null,
             expectedParameterType.getReturnTypeFromFunctionType(),

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/UnitTypeConversions.kt

@@ -78,6 +78,7 @@ object UnitTypeConversions : ParameterTypeConversion {
             candidate.callComponents.builtIns,
             expectedParameterType.annotations,
             expectedParameterType.getReceiverTypeFromFunctionType(),
+            expectedParameterType.getContextReceiverTypesFromFunctionType(),
             expectedParameterType.getValueParameterTypesFromFunctionType().map { it.type },
             parameterNames = null,
             candidate.callComponents.builtIns.nullableAnyType,

compiler/testData/builtin-classes/default/kotlin.txt

@@ -195,6 +195,12 @@ public interface Comparable</*0*/ in T> {
     public abstract operator fun compareTo(/*0*/ other: T): kotlin.Int
 }
 
+@kotlin.annotation.Target(allowedTargets = {AnnotationTarget.TYPE}) @kotlin.annotation.MustBeDocumented public final annotation class ContextFunctionTypeParams : kotlin.Annotation {
+    /*primary*/ public constructor ContextFunctionTypeParams(/*0*/ count: kotlin.Int)
+    public final val count: kotlin.Int
+        public final fun <get-count>(): kotlin.Int
+}
+
 @kotlin.annotation.Target(allowedTargets = {AnnotationTarget.CLASS, AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY, AnnotationTarget.ANNOTATION_CLASS, AnnotationTarget.CONSTRUCTOR, AnnotationTarget.PROPERTY_SETTER, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.TYPEALIAS}) @kotlin.annotation.MustBeDocumented public final annotation class Deprecated : kotlin.Annotation {
     /*primary*/ public constructor Deprecated(/*0*/ message: kotlin.String, /*1*/ replaceWith: kotlin.ReplaceWith = ..., /*2*/ level: kotlin.DeprecationLevel = ...)
     public final val level: kotlin.DeprecationLevel

compiler/testData/builtin-classes/java6/kotlin.txt

@@ -203,6 +203,12 @@ public interface Comparable</*0*/ in T> {
     public abstract operator fun compareTo(/*0*/ other: T): kotlin.Int
 }
 
+@kotlin.annotation.Target(allowedTargets = {AnnotationTarget.TYPE}) @kotlin.annotation.MustBeDocumented public final annotation class ContextFunctionTypeParams : kotlin.Annotation {
+    /*primary*/ public constructor ContextFunctionTypeParams(/*0*/ count: kotlin.Int)
+    public final val count: kotlin.Int
+        public final fun <get-count>(): kotlin.Int
+}
+
 @kotlin.annotation.Target(allowedTargets = {AnnotationTarget.CLASS, AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY, AnnotationTarget.ANNOTATION_CLASS, AnnotationTarget.CONSTRUCTOR, AnnotationTarget.PROPERTY_SETTER, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.TYPEALIAS}) @kotlin.annotation.MustBeDocumented public final annotation class Deprecated : kotlin.Annotation {
     /*primary*/ public constructor Deprecated(/*0*/ message: kotlin.String, /*1*/ replaceWith: kotlin.ReplaceWith = ..., /*2*/ level: kotlin.DeprecationLevel = ...)
     public final val level: kotlin.DeprecationLevel

compiler/testData/builtin-classes/java8/kotlin.txt

@@ -205,6 +205,12 @@ public interface Comparable</*0*/ in T> {
     public abstract operator fun compareTo(/*0*/ other: T): kotlin.Int
 }
 
+@kotlin.annotation.Target(allowedTargets = {AnnotationTarget.TYPE}) @kotlin.annotation.MustBeDocumented public final annotation class ContextFunctionTypeParams : kotlin.Annotation {
+    /*primary*/ public constructor ContextFunctionTypeParams(/*0*/ count: kotlin.Int)
+    public final val count: kotlin.Int
+        public final fun <get-count>(): kotlin.Int
+}
+
 @kotlin.annotation.Target(allowedTargets = {AnnotationTarget.CLASS, AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY, AnnotationTarget.ANNOTATION_CLASS, AnnotationTarget.CONSTRUCTOR, AnnotationTarget.PROPERTY_SETTER, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.TYPEALIAS}) @kotlin.annotation.MustBeDocumented public final annotation class Deprecated : kotlin.Annotation {
     /*primary*/ public constructor Deprecated(/*0*/ message: kotlin.String, /*1*/ replaceWith: kotlin.ReplaceWith = ..., /*2*/ level: kotlin.DeprecationLevel = ...)
     public final val level: kotlin.DeprecationLevel

compiler/testData/builtin-classes/newMethods/kotlin.txt

@@ -203,6 +203,12 @@ public interface Comparable</*0*/ in T> {
     public abstract operator fun compareTo(/*0*/ other: T): kotlin.Int
 }
 
+@kotlin.annotation.Target(allowedTargets = {AnnotationTarget.TYPE}) @kotlin.annotation.MustBeDocumented public final annotation class ContextFunctionTypeParams : kotlin.Annotation {
+    /*primary*/ public constructor ContextFunctionTypeParams(/*0*/ count: kotlin.Int)
+    public final val count: kotlin.Int
+        public final fun <get-count>(): kotlin.Int
+}
+
 @kotlin.annotation.Target(allowedTargets = {AnnotationTarget.CLASS, AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY, AnnotationTarget.ANNOTATION_CLASS, AnnotationTarget.CONSTRUCTOR, AnnotationTarget.PROPERTY_SETTER, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.TYPEALIAS}) @kotlin.annotation.MustBeDocumented public final annotation class Deprecated : kotlin.Annotation {
     /*primary*/ public constructor Deprecated(/*0*/ message: kotlin.String, /*1*/ replaceWith: kotlin.ReplaceWith = ..., /*2*/ level: kotlin.DeprecationLevel = ...)
     public final val level: kotlin.DeprecationLevel

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.kt

@@ -0,0 +1,10 @@
+// FIR_IDENTICAL
+
+class Button
+class ClickEvent
+
+typealias ClickHandler = context(Button) (ClickEvent) -> Unit
+
+fun handleClick(clickHandler: ClickHandler) {
+    clickHandler(Button(), ClickEvent())
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.txt

@@ -0,0 +1,18 @@
+package
+
+public fun handleClick(/*0*/ clickHandler: ClickHandler /* = @kotlin.ContextFunctionTypeParams(count = 1) (ClickEvent) -> kotlin.Unit */): kotlin.Unit
+
+public final class Button {
+    public constructor Button()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class ClickEvent {
+    public constructor ClickEvent()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+public typealias ClickHandler = @kotlin.ContextFunctionTypeParams(count = 1) (ClickEvent) -> kotlin.Unit

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.fir.kt

@@ -0,0 +1,21 @@
+// !LANGUAGE: +ContextReceivers
+// !DIAGNOSTICS: -ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE
+
+class Context
+class Receiver
+class Param
+
+fun foo(context: Context, receiver: Receiver, p: Param) {}
+
+context(Context)
+fun bar(receiver: Receiver, p: Param) {}
+
+context(Context)
+fun Receiver.baz(p: Param) {}
+
+fun main() {
+    var g: context(Context) Receiver.(Param) -> Unit
+    g = ::<!UNRESOLVED_REFERENCE!>foo<!>         // OK
+    g = ::bar         // OK
+    g = Receiver::baz // OK
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.kt

@@ -0,0 +1,20 @@
+// !DIAGNOSTICS: -ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE
+
+class Context
+class Receiver
+class Param
+
+fun foo(context: Context, receiver: Receiver, p: Param) {}
+
+context(Context)
+fun bar(receiver: Receiver, p: Param) {}
+
+context(Context)
+fun Receiver.baz(p: Param) {}
+
+fun main() {
+    var g: context(Context) Receiver.(Param) -> Unit
+    g = ::foo         // OK
+    g = ::bar         // OK
+    g = Receiver::baz // OK
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.txt

@@ -0,0 +1,27 @@
+package
+
+public fun bar(/*0*/ receiver: Receiver, /*1*/ p: Param): kotlin.Unit
+public fun foo(/*0*/ context: Context, /*1*/ receiver: Receiver, /*2*/ p: Param): kotlin.Unit
+public fun main(): kotlin.Unit
+public fun Receiver.baz(/*0*/ p: Param): kotlin.Unit
+
+public final class Context {
+    public constructor Context()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Param {
+    public constructor Param()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Receiver {
+    public constructor Receiver()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.fir.kt

@@ -0,0 +1,39 @@
+class Param
+class C
+class R
+
+context(C)
+fun R.f1(g: context(C) R.(Param) -> Unit) {
+    g(this@C, this@R, Param())
+}
+
+context(C)
+fun f2(g: context(C) (Param) -> Unit) {
+    g(this@C, Param())
+}
+
+context(C)
+fun R.f3(g: context(C) R.() -> Unit) {
+    g(this@C, this@R)
+}
+
+context(C)
+fun f4(g: context(C) () -> Unit) {
+    g(this@C)
+}
+
+fun test() {
+    val lf1: context(C) R.(Param) -> Unit = { _ -> }
+    val lf2: context(C) (Param) -> Unit = { _ -> }
+    val lf3: context(C) R.() -> Unit = { }
+    val lf4: context(C) () -> Unit = { }
+
+    with(C()) {
+        with(R()) {
+            f1(lf1)
+            f2(lf2)
+            f3(lf3)
+            f4(lf4)
+        }
+    }
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.kt

@@ -0,0 +1,39 @@
+class Param
+class C
+class R
+
+context(C)
+fun R.f1(g: context(C) R.(Param) -> Unit) {
+    g(this@C, this@R, Param())
+}
+
+context(C)
+fun f2(g: context(C) (Param) -> Unit) {
+    g(this@C, Param())
+}
+
+context(C)
+fun R.f3(g: context(C) R.() -> Unit) {
+    g(this@C, this@R)
+}
+
+context(C)
+fun f4(g: context(C) () -> Unit) {
+    g(this@C)
+}
+
+fun test() {
+    val lf1: context(C) R.(Param) -> Unit = { _ -> }
+    val lf2: context(C) (Param) -> Unit = { _ -> }
+    val lf3: context(C) R.() -> Unit = { }
+    val lf4: context(C) () -> Unit = { }
+
+    with(C()) {
+        with(R()) {
+            f1(lf1)
+            f2(lf2)
+            f3(lf3)
+            f4(lf4)
+        }
+    }
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.txt

@@ -0,0 +1,28 @@
+package
+
+public fun f2(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) (Param) -> kotlin.Unit): kotlin.Unit
+public fun f4(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) () -> kotlin.Unit): kotlin.Unit
+public fun test(): kotlin.Unit
+public fun R.f1(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) (R.(Param) -> kotlin.Unit)): kotlin.Unit
+public fun R.f3(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) (R.() -> kotlin.Unit)): kotlin.Unit
+
+public final class C {
+    public constructor C()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Param {
+    public constructor Param()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class R {
+    public constructor R()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -10594,6 +10594,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/dp.kt");
                 }
 
+                @Test
+                @TestMetadata("functionalType.kt")
+                public void testFunctionalType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.kt");
+                }
+
                 @Test
                 @TestMetadata("genericOuterClass.kt")
                 public void testGenericOuterClass() throws Exception {
@@ -10725,6 +10731,28 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 public void testWithExplicitReceiverError() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/withExplicitReceiverError.kt");
                 }
+
+                @Nested
+                @TestMetadata("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP")
+                @TestDataPath("$PROJECT_ROOT")
+                public class FromKEEP {
+                    @Test
+                    public void testAllFilesPresentInFromKEEP() throws Exception {
+                        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                    }
+
+                    @Test
+                    @TestMetadata("clickHandler.kt")
+                    public void testClickHandler() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.kt");
+                    }
+
+                    @Test
+                    @TestMetadata("functionalType.kt")
+                    public void testFunctionalType() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.kt");
+                    }
+                }
             }
         }
 

core/builtins/src/kotlin/Annotations.kt

@@ -97,6 +97,11 @@ public enum class DeprecationLevel {
 @MustBeDocumented
 public annotation class ExtensionFunctionType
 
+// TODO: How to make it require opt-in? @RequiresOptIn doesn't support TYPE target
+@Target(TYPE)
+@MustBeDocumented
+public annotation class ContextFunctionTypeParams(val count: Int)
+
 /**
  * Annotates type arguments of functional type and holds corresponding parameter name specified by the user in type declaration (if any).
  */

core/compiler.common/src/org/jetbrains/kotlin/builtins/StandardNames.kt

@@ -114,6 +114,7 @@ object StandardNames {
         @JvmField val deprecationLevel: FqName = fqName("DeprecationLevel")
         @JvmField val replaceWith: FqName = fqName("ReplaceWith")
         @JvmField val extensionFunctionType: FqName = fqName("ExtensionFunctionType")
+        @JvmField val contextFunctionTypeParams: FqName = fqName("ContextFunctionTypeParams")
         @JvmField val parameterName: FqName = fqName("ParameterName")
         @JvmField val parameterNameClassId: ClassId = ClassId.topLevel(parameterName)
         @JvmField val annotation: FqName = fqName("Annotation")

core/descriptors/src/org/jetbrains/kotlin/builtins/ReflectionTypes.kt

@@ -58,13 +58,15 @@ class ReflectionTypes(module: ModuleDescriptor, private val notFoundClasses: Not
     fun getKFunctionType(
         annotations: Annotations,
         receiverType: KotlinType?,
+        contextReceiverTypes: List<KotlinType>,
         parameterTypes: List<KotlinType>,
         parameterNames: List<Name>?,
         returnType: KotlinType,
         builtIns: KotlinBuiltIns,
         isSuspend: Boolean
     ): SimpleType {
-        val arguments = getFunctionTypeArgumentProjections(receiverType, parameterTypes, parameterNames, returnType, builtIns)
+        val arguments =
+            getFunctionTypeArgumentProjections(receiverType, contextReceiverTypes, parameterTypes, parameterNames, returnType, builtIns)
         val classDescriptor =
             if (isSuspend) getKSuspendFunction(arguments.size - 1 /* return type */) else getKFunction(arguments.size - 1 /* return type */)
         return KotlinTypeFactory.simpleNotNullType(annotations, classDescriptor, arguments)

core/descriptors/src/org/jetbrains/kotlin/builtins/functionTypes.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.builtins
 
 import org.jetbrains.kotlin.builtins.StandardNames.BUILT_INS_PACKAGE_NAME
-import org.jetbrains.kotlin.builtins.functions.BuiltInFictitiousFunctionClassFactory
 import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
@@ -16,6 +15,7 @@ import org.jetbrains.kotlin.descriptors.annotations.FilteredAnnotations
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqNameUnsafe
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.constants.IntValue
 import org.jetbrains.kotlin.resolve.constants.StringValue
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 import org.jetbrains.kotlin.types.*
@@ -128,10 +128,29 @@ private fun FqNameUnsafe.getFunctionalClassKind(): FunctionClassKind? {
     return FunctionClassKind.getFunctionalClassKind(shortName().asString(), toSafe().parent())
 }
 
+fun KotlinType.contextFunctionTypeParamsCount(): Int {
+    val annotationDescriptor = annotations.findAnnotation(StandardNames.FqNames.contextFunctionTypeParams) ?: return 0
+    val constantValue = annotationDescriptor.allValueArguments.getValue(Name.identifier("count"))
+    return (constantValue as IntValue).value
+}
 
 fun KotlinType.getReceiverTypeFromFunctionType(): KotlinType? {
     assert(isBuiltinFunctionalType) { "Not a function type: $this" }
-    return if (isTypeAnnotatedWithExtensionFunctionType) arguments.first().type else null
+    if (!isTypeAnnotatedWithExtensionFunctionType) {
+        return null
+    }
+    val index = contextFunctionTypeParamsCount()
+    return arguments[index].type
+}
+
+fun KotlinType.getContextReceiverTypesFromFunctionType(): List<KotlinType> {
+    assert(isBuiltinFunctionalType) { "Not a function type: $this" }
+    val contextReceiversCount = contextFunctionTypeParamsCount()
+    return if (contextReceiversCount == 0) {
+        emptyList()
+    } else {
+        arguments.subList(0, contextReceiversCount).map { it.type }
+    }
 }
 
 fun KotlinType.getReturnTypeFromFunctionType(): KotlinType {
@@ -148,7 +167,7 @@ fun KotlinType.replaceReturnType(newReturnType: KotlinType): KotlinType {
 fun KotlinType.getValueParameterTypesFromFunctionType(): List<TypeProjection> {
     assert(isBuiltinFunctionalType) { "Not a function type: $this" }
     val arguments = arguments
-    val first = if (isBuiltinExtensionFunctionalType) 1 else 0
+    val first = contextFunctionTypeParamsCount() + if (isBuiltinExtensionFunctionalType) 1 else 0
     val last = arguments.size - 1
     assert(first <= last) { "Not an exact function type: $this" }
     return arguments.subList(first, last)
@@ -183,14 +202,16 @@ fun KotlinType.extractParameterNameFromFunctionTypeArgument(): Name? {
 }
 
 fun getFunctionTypeArgumentProjections(
-        receiverType: KotlinType?,
+    receiverType: KotlinType?,
-        parameterTypes: List<KotlinType>,
+    contextReceiverTypes: List<KotlinType>,
-        parameterNames: List<Name>?,
+    parameterTypes: List<KotlinType>,
-        returnType: KotlinType,
+    parameterNames: List<Name>?,
-        builtIns: KotlinBuiltIns
+    returnType: KotlinType,
+    builtIns: KotlinBuiltIns
 ): List<TypeProjection> {
-    val arguments = ArrayList<TypeProjection>(parameterTypes.size + (if (receiverType != null) 1 else 0) + 1)
+    val arguments = ArrayList<TypeProjection>(parameterTypes.size + contextReceiverTypes.size + (if (receiverType != null) 1 else 0) + 1)
 
+    arguments.addAll(contextReceiverTypes.map { it.asTypeProjection() })
     arguments.addIfNotNull(receiverType?.asTypeProjection())
 
     parameterTypes.mapIndexedTo(arguments) { index, type ->
@@ -219,17 +240,22 @@ fun createFunctionType(
     builtIns: KotlinBuiltIns,
     annotations: Annotations,
     receiverType: KotlinType?,
+    contextReceiverTypes: List<KotlinType>,
     parameterTypes: List<KotlinType>,
     parameterNames: List<Name>?,
     returnType: KotlinType,
     suspendFunction: Boolean = false
 ): SimpleType {
-    val arguments = getFunctionTypeArgumentProjections(receiverType, parameterTypes, parameterNames, returnType, builtIns)
+    val arguments =
-    val parameterCount = if (receiverType == null) parameterTypes.size else parameterTypes.size + 1
+        getFunctionTypeArgumentProjections(receiverType, contextReceiverTypes, parameterTypes, parameterNames, returnType, builtIns)
+    val parameterCount = parameterTypes.size + contextReceiverTypes.size + if (receiverType == null) 0 else 1
     val classDescriptor = getFunctionDescriptor(builtIns, parameterCount, suspendFunction)
 
     // TODO: preserve laziness of given annotations
-    val typeAnnotations = if (receiverType != null) annotations.withExtensionFunctionAnnotation(builtIns) else annotations
+    var typeAnnotations = annotations
+    if (receiverType != null) typeAnnotations = typeAnnotations.withExtensionFunctionAnnotation(builtIns)
+    if (contextReceiverTypes.isNotEmpty()) typeAnnotations =
+        typeAnnotations.withContextReceiversFunctionAnnotation(builtIns, contextReceiverTypes.size)
 
     return KotlinTypeFactory.simpleNotNullType(typeAnnotations, classDescriptor, arguments)
 }
@@ -246,6 +272,19 @@ fun Annotations.withExtensionFunctionAnnotation(builtIns: KotlinBuiltIns) =
         Annotations.create(this + BuiltInAnnotationDescriptor(builtIns, StandardNames.FqNames.extensionFunctionType, emptyMap()))
     }
 
+fun Annotations.withContextReceiversFunctionAnnotation(builtIns: KotlinBuiltIns, contextReceiversCount: Int) =
+    if (hasAnnotation(StandardNames.FqNames.contextFunctionTypeParams)) {
+        this
+    } else {
+        Annotations.create(
+            this + BuiltInAnnotationDescriptor(
+                builtIns, StandardNames.FqNames.contextFunctionTypeParams, mapOf(
+                    Name.identifier("count") to IntValue(contextReceiversCount)
+                )
+            )
+        )
+    }
+
 fun getFunctionDescriptor(builtIns: KotlinBuiltIns, parameterCount: Int, isSuspendFunction: Boolean) =
     if (isSuspendFunction) builtIns.getSuspendFunction(parameterCount) else builtIns.getFunction(parameterCount)
 

core/descriptors/src/org/jetbrains/kotlin/builtins/suspendFunctionTypes.kt

@@ -43,10 +43,11 @@ fun transformSuspendFunctionToRuntimeFunctionType(suspendFunType: KotlinType): S
     }
 
     return createFunctionType(
-        suspendFunType.builtIns,
+            suspendFunType.builtIns,
-        suspendFunType.annotations,
+            suspendFunType.annotations,
-        suspendFunType.getReceiverTypeFromFunctionType(),
+            suspendFunType.getReceiverTypeFromFunctionType(),
-        suspendFunType.getValueParameterTypesFromFunctionType().map(TypeProjection::getType) +
+            suspendFunType.getContextReceiverTypesFromFunctionType(),
+            suspendFunType.getValueParameterTypesFromFunctionType().map(TypeProjection::getType) +
                 KotlinTypeFactory.simpleType(
                     Annotations.EMPTY,
                     // Continuation interface is not a part of built-ins anymore, it has been moved to stdlib.

core/descriptors/src/org/jetbrains/kotlin/resolve/sam/SamConversionResolverImpl.kt

@@ -76,14 +76,15 @@ fun getFunctionTypeForAbstractMethod(
     val parameterTypes = ArrayList<KotlinType>(valueParameters.size)
     val parameterNames = ArrayList<Name>(valueParameters.size)
 
-    var startIndex = 0
+    val contextReceiversTypes = function.contextReceiverParameters.map { it.type }
+    var startIndex = contextReceiversTypes.size
     var receiverType: KotlinType? = null
     val extensionReceiver = function.extensionReceiverParameter
     if (extensionReceiver != null) {
         receiverType = extensionReceiver.type
     } else if (shouldConvertFirstParameterToDescriptor && function.valueParameters.isNotEmpty()) {
         receiverType = valueParameters[0].type
-        startIndex = 1
+        startIndex += 1
     }
 
     for (i in startIndex until valueParameters.size) {
@@ -93,7 +94,7 @@ fun getFunctionTypeForAbstractMethod(
     }
 
     return createFunctionType(
-        function.builtIns, EMPTY, receiverType, parameterTypes,
+        function.builtIns, EMPTY, receiverType, contextReceiversTypes, parameterTypes,
         parameterNames, returnType, function.isSuspend
     )
 }

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/TypeDeserializer.kt

@@ -226,6 +226,7 @@ class TypeDeserializer(
             funType.builtIns,
             funType.annotations,
             funType.getReceiverTypeFromFunctionType(),
+            funType.getContextReceiverTypesFromFunctionType(),
             funType.getValueParameterTypesFromFunctionType().dropLast(1).map(TypeProjection::getType),
             // TODO: names
             null,

libraries/stdlib/api/js-v1/kotlin.kt

@@ -1187,6 +1187,14 @@ public open class ConcurrentModificationException : kotlin.RuntimeException {
     public constructor ConcurrentModificationException(cause: kotlin.Throwable?)
 }
 
+@kotlin.annotation.Target(allowedTargets = {AnnotationTarget.TYPE})
+@kotlin.annotation.MustBeDocumented
+public final annotation class ContextFunctionTypeParams : kotlin.Annotation {
+    public constructor ContextFunctionTypeParams(count: kotlin.Int)
+
+    public final val count: kotlin.Int { get; }
+}
+
 @kotlin.SinceKotlin(version = "1.4")
 @kotlin.ExperimentalStdlibApi
 public final class DeepRecursiveFunction<T, R> {

libraries/stdlib/api/js/kotlin.kt

@@ -1158,6 +1158,14 @@ public open class ConcurrentModificationException : kotlin.RuntimeException {
     public constructor ConcurrentModificationException(cause: kotlin.Throwable?)
 }
 
+@kotlin.annotation.Target(allowedTargets = {AnnotationTarget.TYPE})
+@kotlin.annotation.MustBeDocumented
+public final annotation class ContextFunctionTypeParams : kotlin.Annotation {
+    public constructor ContextFunctionTypeParams(count: kotlin.Int)
+
+    public final val count: kotlin.Int { get; }
+}
+
 @kotlin.SinceKotlin(version = "1.4")
 @kotlin.ExperimentalStdlibApi
 public final class DeepRecursiveFunction<T, R> {

libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt

@@ -1,6 +1,10 @@
 public abstract interface annotation class kotlin/BuilderInference : java/lang/annotation/Annotation {
 }
 
+public abstract interface annotation class kotlin/ContextFunctionTypeParams : java/lang/annotation/Annotation {
+	public abstract fun count ()I
+}
+
 public final class kotlin/DeepRecursiveFunction {
 	public fun <init> (Lkotlin/jvm/functions/Function3;)V
 }