FIR: Support (de)serialization of context receivers

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt

@@ -194,6 +194,8 @@ fun deserializeClassToSymbol(
             }
         })
         companionObjectSymbol = (declarations.firstOrNull { it is FirRegularClass && it.isCompanion } as FirRegularClass?)?.symbol
+
+        contextReceivers.addAll(classDeserializer.createContextReceiversForClass(classProto))
     }.also {
         if (isSealed) {
             val inheritors = classProto.sealedSubclassFqNameList.map { nameIndex ->

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt

@@ -19,10 +19,11 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildExpressionStub
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
-import org.jetbrains.kotlin.fir.toEffectiveVisibility
+import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.computeTypeAttributes
+import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitUnitTypeRef
@@ -420,11 +421,44 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
             this.containerSource = c.containerSource
             this.initializer = c.constDeserializer.loadConstant(proto, symbol.callableId, c.nameResolver)
             deprecation = annotations.getDeprecationInfosFromAnnotations(c.session.languageVersionSettings.apiVersion, false)
+
+            contextReceivers.addAll(
+                when {
+                    proto.contextReceiverTypeCount > 0 ->
+                        proto.contextReceiverTypeList.map(::loadContextReceiverFromType)
+                    proto.contextReceiverTypeIdCount > 0 ->
+                        proto.contextReceiverTypeIdList.map(::loadContextReceiverFromTypeId)
+                    else -> emptyList()
+                }
+            )
+
         }.apply {
             versionRequirementsTable = c.versionRequirementTable
         }
     }
 
+    fun loadContextReceiverFromType(proto: ProtoBuf.Type): FirContextReceiver =
+        loadContextReceiverFromTypeRef(proto.toTypeRef(c))
+
+    fun loadContextReceiverFromTypeId(id: Int): FirContextReceiver = loadContextReceiverFromType(c.typeTable[id])
+
+    private fun loadContextReceiverFromTypeRef(typeRef: FirTypeRef): FirContextReceiver =
+        buildContextReceiver {
+            val type = typeRef.coneType
+            this.labelNameFromTypeRef = (type as? ConeLookupTagBasedType)?.lookupTag?.name
+            this.typeRef = typeRef
+        }
+
+    fun createContextReceiversForClass(
+        classProto: ProtoBuf.Class,
+    ): List<FirContextReceiver> = when {
+        classProto.contextReceiverTypeCount > 0 ->
+            classProto.contextReceiverTypeList.map(::loadContextReceiverFromType)
+        classProto.contextReceiverTypeIdCount > 0 ->
+            classProto.contextReceiverTypeIdList.map(::loadContextReceiverFromTypeId)
+        else -> emptyList()
+    }
+
     fun loadFunction(
         proto: ProtoBuf.Function,
         classProto: ProtoBuf.Class? = null,
@@ -485,6 +519,16 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
                 c.annotationDeserializer.loadFunctionAnnotations(c.containerSource, proto, local.nameResolver, local.typeTable)
             deprecation = annotations.getDeprecationInfosFromAnnotations(c.session.languageVersionSettings.apiVersion, false)
             this.containerSource = c.containerSource
+
+            contextReceivers.addAll(
+                when {
+                    proto.contextReceiverTypeCount > 0 ->
+                        proto.contextReceiverTypeList.map(::loadContextReceiverFromType)
+                    proto.contextReceiverTypeIdCount > 0 ->
+                        proto.contextReceiverTypeIdList.map(::loadContextReceiverFromTypeId)
+                    else -> emptyList()
+                }
+            )
         }.apply {
             versionRequirementsTable = c.versionRequirementTable
         }
@@ -561,6 +605,8 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
                 c.annotationDeserializer.loadConstructorAnnotations(c.containerSource, proto, local.nameResolver, local.typeTable)
             containerSource = c.containerSource
             deprecation = annotations.getDeprecationInfosFromAnnotations(c.session.languageVersionSettings.apiVersion, false)
+
+            contextReceivers.addAll(createContextReceiversForClass(classProto))
         }.build().apply {
             containingClassForStaticMemberAttr = c.dispatchReceiver!!.lookupTag
             versionRequirementsTable = c.versionRequirementTable

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

@@ -189,6 +189,17 @@ class FirElementSerializer private constructor(
             }
         }
 
+        if (klass is FirRegularClass) {
+            for (contextReceiver in klass.contextReceivers) {
+                val typeRef = contextReceiver.typeRef
+                if (useTypeTable()) {
+                    builder.addContextReceiverTypeId(typeId(typeRef))
+                } else {
+                    builder.addContextReceiverType(typeProto(contextReceiver.typeRef))
+                }
+            }
+        }
+
         if (versionRequirementTable == null) error("Version requirements must be serialized for classes: ${klass.render()}")
 
         builder.addAllVersionRequirement(versionRequirementTable.serializeVersionRequirements(klass))
@@ -312,6 +323,15 @@ class FirElementSerializer private constructor(
             builder.addTypeParameter(local.typeParameterProto(typeParameter))
         }
 
+        for (contextReceiver in property.contextReceivers) {
+            val typeRef = contextReceiver.typeRef
+            if (useTypeTable()) {
+                builder.addContextReceiverTypeId(typeId(typeRef))
+            } else {
+                builder.addContextReceiverType(typeProto(contextReceiver.typeRef))
+            }
+        }
+
         val receiverTypeRef = property.receiverTypeRef
         if (receiverTypeRef != null) {
             if (useTypeTable()) {
@@ -386,6 +406,15 @@ class FirElementSerializer private constructor(
             builder.addTypeParameter(local.typeParameterProto(typeParameter))
         }
 
+        for (contextReceiver in function.contextReceivers) {
+            val typeRef = contextReceiver.typeRef
+            if (useTypeTable()) {
+                builder.addContextReceiverTypeId(typeId(typeRef))
+            } else {
+                builder.addContextReceiverType(typeProto(contextReceiver.typeRef))
+            }
+        }
+
         val receiverTypeRef = function.receiverTypeRef
         if (receiverTypeRef != null) {
             if (useTypeTable()) {