FIR: Introduce ContextFunctionTypeParams type attribute

compiler/fir/analysis-tests/testData/resolve/inference/nestedExtensionFunctionType.fir.txt

@@ -8,7 +8,7 @@ FILE: nestedExtensionFunctionType.kt
          {
             lval <iterator>: R|kotlin/collections/Iterator<@ExtensionFunctionType kotlin/Function1<A, kotlin/Unit>>| = R|<local>/ys|.R|SubstitutionOverride<kotlin/collections/List.iterator: R|kotlin/collections/Iterator<@ExtensionFunctionType kotlin/Function1<A, kotlin/Unit>>|>|()
             while(R|<local>/<iterator>|.R|SubstitutionOverride<kotlin/collections/Iterator.hasNext: R|kotlin/Boolean|>|()) {
-                lval y: R|(A) -> kotlin/Unit| = R|<local>/<iterator>|.R|SubstitutionOverride<kotlin/collections/Iterator.next: R|(A) -> kotlin/Unit|>|()
+                lval y: R|A.() -> kotlin/Unit| = R|<local>/<iterator>|.R|SubstitutionOverride<kotlin/collections/Iterator.next: R|A.() -> kotlin/Unit|>|()
                 R|<local>/y|.R|SubstitutionOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/a|)
             }
 
@@ -19,7 +19,7 @@ FILE: nestedExtensionFunctionType.kt
          {
             lval <iterator>: R|kotlin/collections/Iterator<@ExtensionFunctionType kotlin/Function1<A, kotlin/Unit>>| = R|<local>/zs|.R|SubstitutionOverride<kotlin/Array.iterator: R|kotlin/collections/Iterator<CapturedType(out @ExtensionFunctionType kotlin/Function1<A, kotlin/Unit>)>|>|()
             while(R|<local>/<iterator>|.R|SubstitutionOverride<kotlin/collections/Iterator.hasNext: R|kotlin/Boolean|>|()) {
-                lval z: R|(A) -> kotlin/Unit| = R|<local>/<iterator>|.R|SubstitutionOverride<kotlin/collections/Iterator.next: R|(A) -> kotlin/Unit|>|()
+                lval z: R|A.() -> kotlin/Unit| = R|<local>/<iterator>|.R|SubstitutionOverride<kotlin/collections/Iterator.next: R|A.() -> kotlin/Unit|>|()
                 R|<local>/z|.R|SubstitutionOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/a|)
             }
 

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/CompilerConeAttributes.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -67,6 +68,22 @@ object CompilerConeAttributes {
         override fun toString(): String = "@ExtensionFunctionType"
     }
 
+    class ContextFunctionTypeParams(val contextReceiverNumber: Int) : ConeAttribute<ContextFunctionTypeParams>() {
+        override fun union(other: ContextFunctionTypeParams?): ContextFunctionTypeParams? = other
+        override fun intersect(other: ContextFunctionTypeParams?): ContextFunctionTypeParams = this
+        override fun add(other: ContextFunctionTypeParams?): ContextFunctionTypeParams = this
+
+        override fun isSubtypeOf(other: ContextFunctionTypeParams?): Boolean = true
+
+        override val key: KClass<out ContextFunctionTypeParams> = ContextFunctionTypeParams::class
+
+        override fun toString(): String = "@${StandardNames.FqNames.contextFunctionTypeParams.shortName().asString()}"
+
+        companion object {
+            val ANNOTATION_CLASS_ID = ClassId.topLevel(StandardNames.FqNames.contextFunctionTypeParams)
+        }
+    }
+
     object UnsafeVariance : ConeAttribute<UnsafeVariance>() {
         val ANNOTATION_CLASS_ID = ClassId(FqName("kotlin"), Name.identifier("UnsafeVariance"))
 
@@ -102,6 +119,8 @@ val ConeAttributes.exact: CompilerConeAttributes.Exact? by ConeAttributes.attrib
 val ConeAttributes.noInfer: CompilerConeAttributes.NoInfer? by ConeAttributes.attributeAccessor<CompilerConeAttributes.NoInfer>()
 val ConeAttributes.enhancedNullability: CompilerConeAttributes.EnhancedNullability? by ConeAttributes.attributeAccessor<CompilerConeAttributes.EnhancedNullability>()
 val ConeAttributes.extensionFunctionType: CompilerConeAttributes.ExtensionFunctionType? by ConeAttributes.attributeAccessor<CompilerConeAttributes.ExtensionFunctionType>()
+private val ConeAttributes.contextFunctionTypeParams: CompilerConeAttributes.ContextFunctionTypeParams? by ConeAttributes.attributeAccessor<CompilerConeAttributes.ContextFunctionTypeParams>()
+
 val ConeAttributes.unsafeVarianceType: CompilerConeAttributes.UnsafeVariance? by ConeAttributes.attributeAccessor<CompilerConeAttributes.UnsafeVariance>()
 
 // ------------------------------------------------------------------
@@ -116,3 +135,11 @@ val ConeKotlinType.hasEnhancedNullability: Boolean
 
 val ConeKotlinType.isExtensionFunctionType: Boolean
     get() = attributes.extensionFunctionType != null
+
+val ConeKotlinType.hasContextReceivers: Boolean
+    get() = attributes.contextReceiversNumberForFunctionType > 0
+
+val ConeKotlinType.contextReceiversNumberForFunctionType: Int
+    get() = attributes.contextReceiversNumberForFunctionType
+
+val ConeAttributes.contextReceiversNumberForFunctionType: Int get() = contextFunctionTypeParams?.contextReceiverNumber ?: 0

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/TypeRenderer.kt

@@ -70,9 +70,12 @@ fun ConeTypeProjection.render(): String {
 }
 
 fun ConeKotlinType.renderFunctionType(
-    kind: FunctionClassKind?, isExtension: Boolean, renderType: ConeTypeProjection.() -> String = { render() }
+    kind: FunctionClassKind?, renderType: ConeTypeProjection.() -> String = { render() }
 ): String {
     if (!kind.withPrettyRender()) return renderType()
+
+    val isExtension = isExtensionFunctionType
+
     val renderedType = buildString {
         if (kind == FunctionClassKind.SuspendFunction) {
             append("suspend ")

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirElementSerializer.kt

@@ -662,6 +662,15 @@ class FirElementSerializer private constructor(
                     return functionType
                 }
                 fillFromPossiblyInnerType(builder, type)
+                if (type.hasContextReceivers) {
+                    serializeAnnotationFromAttribute(
+                        correspondingTypeRef?.annotations, CompilerConeAttributes.ContextFunctionTypeParams.ANNOTATION_CLASS_ID, builder,
+                        argumentMapping = buildAnnotationArgumentMapping {
+                            this.mapping[StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME] =
+                                buildConstExpression(source = null, ConstantValueKind.Int, type.contextReceiversNumberForFunctionType)
+                        }
+                    )
+                }
             }
             is ConeTypeParameterType -> {
                 val typeParameter = type.lookupTag.typeParameterSymbol.fir
@@ -721,17 +730,18 @@ class FirElementSerializer private constructor(
     private fun serializeAnnotationFromAttribute(
         existingAnnotations: List<FirAnnotation>?,
         classId: ClassId,
-        builder: ProtoBuf.Type.Builder
+        builder: ProtoBuf.Type.Builder,
+        argumentMapping: FirAnnotationArgumentMapping = FirEmptyAnnotationArgumentMapping,
     ) {
         if (existingAnnotations?.any { it.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.classId == classId } != true) {
             extension.serializeTypeAnnotation(
                 buildAnnotation {
                     annotationTypeRef = buildResolvedTypeRef {
-                        this.type = CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID.constructClassLikeType(
+                        this.type = classId.constructClassLikeType(
                             emptyArray(), isNullable = false
                         )
                     }
-                    argumentMapping = FirEmptyAnnotationArgumentMapping
+                    this.argumentMapping = argumentMapping
                 }, builder
             )
         }

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt

@@ -2161,9 +2161,6 @@ class DeclarationsConverter(
             receiverTypeRef = receiverTypeReference
             returnTypeRef = returnTypeReference
             valueParameters += valueParametersList.map { it.firValueParameter }
-            if (receiverTypeReference != null) {
-                annotations += extensionFunctionAnnotation
-            }
             this.isSuspend = isSuspend
             this.contextReceiverTypeRefs.addAll(
                 functionType.getChildNodeByType(CONTEXT_RECEIVER_LIST)?.getChildNodesByType(CONTEXT_RECEIVER)?.mapNotNull {
@@ -2239,17 +2236,6 @@ class DeclarationsConverter(
         return ValueParameter(isVal, isVar, modifiers, firValueParameter, destructuringDeclaration)
     }
 
-    private val extensionFunctionAnnotation = buildAnnotation {
-        annotationTypeRef = buildResolvedTypeRef {
-            type = ConeClassLikeTypeImpl(
-                ConeClassLikeLookupTagImpl(EXTENSION_FUNCTION_ANNOTATION),
-                emptyArray(),
-                false
-            )
-        }
-        argumentMapping = FirEmptyAnnotationArgumentMapping
-    }
-
     private fun <T> fillDanglingConstraintsTo(
         typeParameters: List<FirTypeParameter>,
         typeConstraints: List<TypeConstraint>,

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1771,9 +1771,6 @@ open class RawFirBuilder(
                         for (valueParameter in unwrappedElement.parameters) {
                             valueParameters += valueParameter.convert<FirValueParameter>()
                         }
-                        if (receiverTypeRef != null) {
-                            annotations += extensionFunctionAnnotation
-                        }
 
                         contextReceiverTypeRefs.addAll(
                             unwrappedElement.contextReceiversTypeReferences.mapNotNull {
@@ -2538,17 +2535,6 @@ open class RawFirBuilder(
             }
         }
     }
-
-    private val extensionFunctionAnnotation = buildAnnotation {
-        annotationTypeRef = buildResolvedTypeRef {
-            type = ConeClassLikeTypeImpl(
-                ConeClassLikeLookupTagImpl(EXTENSION_FUNCTION_ANNOTATION),
-                emptyArray(),
-                isNullable = false
-            )
-        }
-        argumentMapping = FirEmptyAnnotationArgumentMapping
-    }
 }
 
 enum class BodyBuildingMode {

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

@@ -468,6 +468,7 @@ 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.valueParameters.map { it.returnTypeRef.coneType.withParameterNameAnnotation(it) } +
                     listOf(typeRef.returnTypeRef.coneType)
         val classId = if (typeRef.isSuspend) {
@@ -475,7 +476,19 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
         } else {
             StandardClassIds.FunctionN(typeRef.parametersCount)
         }
-        val attributes = typeRef.annotations.computeTypeAttributes(session)
+
+        val attributes = typeRef.annotations.computeTypeAttributes(
+            session,
+            predefined = buildList {
+                if (typeRef.receiverTypeRef != null) {
+                    add(CompilerConeAttributes.ExtensionFunctionType)
+                }
+
+                if (typeRef.contextReceiverTypeRefs.isNotEmpty()) {
+                    add(CompilerConeAttributes.ContextFunctionTypeParams(typeRef.contextReceiverTypeRefs.size))
+                }
+            }
+        )
         val symbol = resolveBuiltInQualified(classId, session)
         return ConeClassLikeTypeImpl(
             symbol.toLookupTag().also {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -1127,9 +1127,7 @@ open class FirRenderer(builder: StringBuilder, protected val mode: RenderMode =
         val kind = resolvedTypeRef.functionTypeKind
         print("R|")
         val coneType = resolvedTypeRef.type
-        print(coneType.renderFunctionType(kind, resolvedTypeRef.annotations.any {
+        print(coneType.renderFunctionType(kind))
-            it.isExtensionFunctionAnnotationCall
-        }))
         print("|")
     }
 

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
@@ -113,9 +114,13 @@ fun ConeClassLikeType.toConstKind(): ConstantValueKind<*>? = when (lookupTag.cla
     else -> null
 }
 
-fun List<FirAnnotation>.computeTypeAttributes(session: FirSession): ConeAttributes {
+fun List<FirAnnotation>.computeTypeAttributes(session: FirSession, predefined: List<ConeAttribute<*>> = emptyList()): ConeAttributes {
-    if (this.isEmpty()) return ConeAttributes.Empty
+    if (this.isEmpty()) {
+        if (predefined.isEmpty()) return ConeAttributes.Empty
+        return ConeAttributes.create(predefined)
+    }
     val attributes = mutableListOf<ConeAttribute<*>>()
+    attributes += predefined
     val customAnnotations = mutableListOf<FirAnnotation>()
     for (annotation in this) {
         val type = annotation.annotationTypeRef.coneTypeSafe<ConeClassLikeType>() ?: continue
@@ -123,6 +128,11 @@ fun List<FirAnnotation>.computeTypeAttributes(session: FirSession): ConeAttribut
             CompilerConeAttributes.Exact.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.Exact
             CompilerConeAttributes.NoInfer.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.NoInfer
             CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.ExtensionFunctionType
+            CompilerConeAttributes.ContextFunctionTypeParams.ANNOTATION_CLASS_ID ->
+                attributes +=
+                    CompilerConeAttributes.ContextFunctionTypeParams(
+                        annotation.extractContextReceiversCount() ?: 0
+                    )
             CompilerConeAttributes.UnsafeVariance.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.UnsafeVariance
             else -> {
                 val attributeFromPlugin = session.extensionService.typeAttributeExtensions.firstNotNullOfOrNull {
@@ -142,6 +152,9 @@ fun List<FirAnnotation>.computeTypeAttributes(session: FirSession): ConeAttribut
     return ConeAttributes.create(attributes)
 }
 
+private fun FirAnnotation.extractContextReceiversCount() =
+    (argumentMapping.mapping[StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME] as? FirConstExpression<*>)?.value as? Int
+
 fun FirTypeProjection.toConeTypeProjection(): ConeTypeProjection =
     when (this) {
         is FirStarProjection -> ConeStarProjection

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -324,7 +324,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
 
         private fun ConeTypeProjection.tryToRenderConeAsFunctionType(): String {
             if (this !is ConeKotlinType) return localTypeRenderer()
-            val functionType = renderFunctionType(functionTypeKind, isExtensionFunctionType) { localTypeRenderer() }
+            val functionType = renderFunctionType(functionTypeKind) { localTypeRenderer() }
             return functionType.removeCurrentFilePackage()
         }
 

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

@@ -30,6 +30,8 @@ object StandardNames {
 
     @JvmField val CHAR_CODE = Name.identifier("code")
 
+    @JvmField val CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME = Name.identifier("count")
+
     @JvmField val COROUTINES_PACKAGE_FQ_NAME = FqName("kotlin.coroutines")
 
     @JvmField val COROUTINES_JVM_INTERNAL_PACKAGE_FQ_NAME = FqName("kotlin.coroutines.jvm.internal")

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

@@ -130,7 +130,7 @@ private fun FqNameUnsafe.getFunctionalClassKind(): FunctionClassKind? {
 
 fun KotlinType.contextFunctionTypeParamsCount(): Int {
     val annotationDescriptor = annotations.findAnnotation(StandardNames.FqNames.contextFunctionTypeParams) ?: return 0
-    val constantValue = annotationDescriptor.allValueArguments.getValue(Name.identifier("count"))
+    val constantValue = annotationDescriptor.allValueArguments.getValue(StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME)
     return (constantValue as IntValue).value
 }
 
@@ -279,7 +279,7 @@ fun Annotations.withContextReceiversFunctionAnnotation(builtIns: KotlinBuiltIns,
         Annotations.create(
             this + BuiltInAnnotationDescriptor(
                 builtIns, StandardNames.FqNames.contextFunctionTypeParams, mapOf(
-                    Name.identifier("count") to IntValue(contextReceiversCount)
+                    StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME to IntValue(contextReceiversCount)
                 )
             )
         )