[PSI, FE] Support functional types

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?,
-        parameterTypes: List<KotlinType>,
-        parameterNames: List<Name>?,
-        returnType: KotlinType,
-        builtIns: KotlinBuiltIns
+    receiverType: KotlinType?,
+    contextReceiverTypes: List<KotlinType>,
+    parameterTypes: List<KotlinType>,
+    parameterNames: List<Name>?,
+    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 parameterCount = if (receiverType == null) parameterTypes.size else parameterTypes.size + 1
+    val arguments =
+        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.annotations,
-        suspendFunType.getReceiverTypeFromFunctionType(),
-        suspendFunType.getValueParameterTypesFromFunctionType().map(TypeProjection::getType) +
+            suspendFunType.builtIns,
+            suspendFunType.annotations,
+            suspendFunType.getReceiverTypeFromFunctionType(),
+            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,