[FIR] Add extension point to metadata serialization

This is a temporary solution for KT-64312
Introduced extension point should be removed (KT-64695)

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtensi
 import org.jetbrains.kotlin.fir.backend.Fir2IrScriptConfiguratorExtension
 import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension
 import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension
+import org.jetbrains.kotlin.fir.serialization.FirMetadataSerializerPlugin
 import kotlin.reflect.KClass
 
 abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
@@ -34,6 +35,8 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirScriptConfiguratorExtension::class,
             Fir2IrScriptConfiguratorExtension::class,
             FirFunctionTypeKindExtension::class,
+            @OptIn(FirExtensionApiInternals::class)
+            FirMetadataSerializerPlugin::class,
         )
 
         internal val ALLOWED_EXTENSIONS_FOR_LIBRARY_SESSION = listOf(
@@ -107,6 +110,12 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             registerExtension(FirFunctionTypeKindExtension::class, this)
         }
 
+        @FirExtensionApiInternals
+        @JvmName("plusMetadataSerializerPlugin")
+        operator fun (FirMetadataSerializerPlugin.Factory).unaryPlus() {
+            registerExtension(FirMetadataSerializerPlugin::class, this)
+        }
+
         // ------------------ reference methods ------------------
 
         @JvmName("plusStatusTransformerExtension")
@@ -169,6 +178,12 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirFunctionTypeKindExtension.Factory { this.invoke(it) }.unaryPlus()
         }
 
+        @FirExtensionApiInternals
+        @JvmName("plusMetadataSerializerPlugin")
+        operator fun ((FirSession) -> FirMetadataSerializerPlugin).unaryPlus() {
+            FirMetadataSerializerPlugin.Factory { this.invoke(it) }.unaryPlus()
+        }
+
         // ------------------ utilities ------------------
 
         @JvmName("bindLeft")

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

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
 import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
 import org.jetbrains.kotlin.fir.expressions.builder.buildConstExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyAnnotationArgumentMapping
+import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.typeAttributeExtensions
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -55,6 +56,7 @@ import org.jetbrains.kotlin.metadata.serialization.Interner
 import org.jetbrains.kotlin.metadata.serialization.MutableTypeTable
 import org.jetbrains.kotlin.metadata.serialization.MutableVersionRequirementTable
 import org.jetbrains.kotlin.name.*
+import org.jetbrains.kotlin.protobuf.GeneratedMessageLite
 import org.jetbrains.kotlin.resolve.RequireKotlinConstants
 import org.jetbrains.kotlin.serialization.deserialization.ProtoEnumFlags
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
@@ -291,6 +293,21 @@ class FirElementSerializer private constructor(
         typeTable.serialize()?.let { builder.typeTable = it }
         versionRequirementTable.serialize()?.let { builder.versionRequirementTable = it }
 
+        if (klass is FirRegularClass) {
+            @OptIn(FirExtensionApiInternals::class)
+            for (plugin in session.extensionService.metadataSerializerPlugins) {
+                val protoRegistrar = object : FirMetadataSerializerPlugin.ProtoRegistrar {
+                    override fun <Type> setExtension(
+                        extension: GeneratedMessageLite.GeneratedExtension<ProtoBuf.Class, Type>,
+                        value: Type,
+                    ) {
+                        builder.setExtension(extension, value)
+                    }
+                }
+                plugin.registerProtoExtensions(klass.symbol, stringTable, protoRegistrar)
+            }
+        }
+
         return builder
     }
 

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.serialization
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.extensions.FirExtension
+import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
+import org.jetbrains.kotlin.fir.extensions.FirExtensionPointName
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.protobuf.GeneratedMessageLite
+import kotlin.reflect.KClass
+
+
+/**
+ * This extension point is temporary and should not be used
+ *
+ * It will be removed soon: KT-64695
+ * DON'T USE IT
+ */
+@FirExtensionApiInternals
+abstract class FirMetadataSerializerPlugin(session: FirSession) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("MetadataSerializerPlugin")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension>
+        get() = FirMetadataSerializerPlugin::class
+
+    abstract fun registerProtoExtensions(
+        symbol: FirRegularClassSymbol,
+        stringTable: FirElementAwareStringTable,
+        protoRegistrar: ProtoRegistrar
+    )
+
+    interface ProtoRegistrar {
+        fun <Type> setExtension(
+            extension: GeneratedMessageLite.GeneratedExtension<ProtoBuf.Class, Type>,
+            value: Type,
+        )
+    }
+
+    fun interface Factory : FirExtension.Factory<FirMetadataSerializerPlugin>
+}
+
+@OptIn(FirExtensionApiInternals::class)
+internal val FirExtensionService.metadataSerializerPlugins: List<FirMetadataSerializerPlugin> by FirExtensionService.registeredExtensions()