[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-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 }
     }