[AA] KT-55098 Add context receivers to KtFunctionalType

- Context receivers in function types may not be labeled, so the created
  `KtContextReceiver`s have `null` labels. Such labels currently only
  compile due to a bug (see KT-55187).

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/types/KtFe10FunctionalType.kt

@@ -5,7 +5,9 @@
 
 package org.jetbrains.kotlin.analysis.api.descriptors.types
 
+import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.KtTypeProjection
+import org.jetbrains.kotlin.analysis.api.base.KtContextReceiver
 import org.jetbrains.kotlin.analysis.api.descriptors.Fe10AnalysisContext
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.KtFe10DescNamedClassOrObjectSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.ktNullability
@@ -15,16 +17,15 @@ import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.bas
 import org.jetbrains.kotlin.analysis.api.descriptors.types.base.KtFe10Type
 import org.jetbrains.kotlin.analysis.api.descriptors.types.base.asStringForDebugging
 import org.jetbrains.kotlin.analysis.api.descriptors.utils.KtFe10JvmTypeMapperContext
+import org.jetbrains.kotlin.analysis.api.impl.base.KtContextReceiverImpl
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtClassTypeQualifier
 import org.jetbrains.kotlin.analysis.api.types.KtFunctionalType
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeNullability
+import org.jetbrains.kotlin.builtins.*
 import org.jetbrains.kotlin.builtins.functions.FunctionClassDescriptor
-import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
-import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
-import org.jetbrains.kotlin.builtins.getValueParameterTypesFromFunctionType
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.types.SimpleType
 
@@ -55,6 +56,22 @@ internal class KtFe10FunctionalType(
     override val arity: Int
         get() = withValidityAssertion { descriptor.arity }
 
+    override val hasContextReceivers: Boolean
+        get() = withValidityAssertion { type.contextFunctionTypeParamsCount() > 0 }
+
+    @OptIn(KtAnalysisApiInternals::class)
+    override val contextReceivers: List<KtContextReceiver>
+        get() = withValidityAssertion {
+            type.getContextReceiverTypesFromFunctionType().map { receiverType ->
+                // Context receivers in function types may not have labels, hence the `null` label.
+                KtContextReceiverImpl(
+                    receiverType.toKtType(analysisContext),
+                    _label = null,
+                    analysisContext.token,
+                )
+            }
+        }
+
     override val hasReceiver: Boolean
         get() = withValidityAssertion {
             if (descriptor.functionKind.isReflectType) false

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/types/KtFirFunctionalType.kt

@@ -5,13 +5,16 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.types
 
-import org.jetbrains.kotlin.analysis.api.KtTypeProjection
+import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.KtTypeArgumentWithVariance
+import org.jetbrains.kotlin.analysis.api.KtTypeProjection
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationsList
+import org.jetbrains.kotlin.analysis.api.base.KtContextReceiver
 import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.analysis.api.fir.annotations.KtFirAnnotationListForType
 import org.jetbrains.kotlin.analysis.api.fir.types.qualifiers.UsualClassTypeQualifierBuilder
 import org.jetbrains.kotlin.analysis.api.fir.utils.cached
+import org.jetbrains.kotlin.analysis.api.impl.base.KtContextReceiverImpl
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
@@ -20,11 +23,8 @@ import org.jetbrains.kotlin.analysis.api.types.KtFunctionalType
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeNullability
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.fir.types.isExtensionFunctionType
-import org.jetbrains.kotlin.fir.types.isSuspendFunctionType
-import org.jetbrains.kotlin.fir.types.receiverType
-import org.jetbrains.kotlin.fir.types.renderForDebugging
 import org.jetbrains.kotlin.name.ClassId
 
 internal class KtFirFunctionalType(
@@ -56,10 +56,22 @@ internal class KtFirFunctionalType(
             else coneType.typeArguments.size - 1
         }
 
+    @OptIn(KtAnalysisApiInternals::class)
+    override val contextReceivers: List<KtContextReceiver> by cached {
+        ownTypeArguments.subList(0, coneType.contextReceiversNumberForFunctionType).map { typeProjection ->
+            // Context receivers in function types may not have labels, hence the `null` label.
+            KtContextReceiverImpl((typeProjection as KtTypeArgumentWithVariance).type, _label = null, token)
+        }
+    }
+
+    override val hasContextReceivers: Boolean get() = withValidityAssertion { coneType.hasContextReceivers }
+
     override val receiverType: KtType?
         get() = withValidityAssertion {
-            if (coneType.isExtensionFunctionType) (ownTypeArguments.first() as KtTypeArgumentWithVariance).type
-            else null
+            if (coneType.isExtensionFunctionType) {
+                // The extension receiver type follows any context receiver types.
+                (ownTypeArguments[coneType.contextReceiversNumberForFunctionType] as KtTypeArgumentWithVariance).type
+            } else null
         }
 
     override val hasReceiver: Boolean
@@ -68,8 +80,9 @@ internal class KtFirFunctionalType(
         }
 
     override val parameterTypes: List<KtType> by cached {
-        val parameterTypeArgs = if (coneType.isExtensionFunctionType) ownTypeArguments.subList(1, ownTypeArguments.lastIndex)
-        else ownTypeArguments.subList(0, ownTypeArguments.lastIndex)
+        // Parameter types follow context and extension receiver types.
+        val firstIndex = coneType.contextReceiversNumberForFunctionType + (if (coneType.isExtensionFunctionType) 1 else 0)
+        val parameterTypeArgs = ownTypeArguments.subList(firstIndex, ownTypeArguments.lastIndex)
         parameterTypeArgs.map { (it as KtTypeArgumentWithVariance).type }
     }
 

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/base/KtContextReceiver.kt

@@ -37,7 +37,7 @@ public abstract class KtContextReceiver : KtLifetimeOwner {
 }
 
 /**
- * Something which can have a [KtContextReceiver] declared. This may be a callable symbol or a class symbol
+ * Something which can have a [KtContextReceiver] declared. This may be a callable symbol, a class symbol, or a functional type.
  */
 public interface KtContextReceiversOwner : KtLifetimeOwner {
     /**

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/types/KtType.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.api.types
 
 import org.jetbrains.kotlin.analysis.api.KtTypeProjection
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotated
+import org.jetbrains.kotlin.analysis.api.base.KtContextReceiversOwner
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
@@ -52,9 +53,10 @@ public sealed class KtNonErrorClassType : KtClassType() {
     abstract override val qualifiers: List<KtClassTypeQualifier.KtResolvedClassTypeQualifier>
 }
 
-public abstract class KtFunctionalType : KtNonErrorClassType() {
+public abstract class KtFunctionalType : KtNonErrorClassType(), KtContextReceiversOwner {
     public abstract val isSuspend: Boolean
     public abstract val arity: Int
+    public abstract val hasContextReceivers: Boolean
     public abstract val receiverType: KtType?
     public abstract val hasReceiver: Boolean
     public abstract val parameterTypes: List<KtType>