[FIR] Add ability to provide additional synthetic declarations to metadata from plugins

^KT-55885 Fixed

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

@@ -33,6 +33,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirAssignExpressionAltererExtension::class,
             FirScriptConfiguratorExtension::class,
             FirFunctionTypeKindExtension::class,
+            FirDeclarationsForMetadataProviderExtension::class,
         )
     }
 
@@ -96,6 +97,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             registerExtension(FirFunctionTypeKindExtension::class, this)
         }
 
+        @JvmName("plusDeclarationForMetadataProviderExtension")
+        operator fun (FirDeclarationsForMetadataProviderExtension.Factory).unaryPlus() {
+            registerExtension(FirDeclarationsForMetadataProviderExtension::class, this)
+        }
+
         // ------------------ reference methods ------------------
 
         @JvmName("plusStatusTransformerExtension")
@@ -153,6 +159,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirFunctionTypeKindExtension.Factory { this.invoke(it) }.unaryPlus()
         }
 
+        @JvmName("plusDeclarationForMetadataProviderExtension")
+        operator fun ((FirSession) -> FirDeclarationsForMetadataProviderExtension).unaryPlus() {
+            FirDeclarationsForMetadataProviderExtension.Factory { this.invoke(it) }.unaryPlus()
+        }
+
         // ------------------ utilities ------------------
 
         @JvmName("bindLeft")

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

@@ -25,6 +25,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.declarationForMetadataProviders
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.typeAttributeExtensions
 import org.jetbrains.kotlin.fir.resolve.*
@@ -69,22 +70,34 @@ class FirElementSerializer private constructor(
     private val languageVersionSettings: LanguageVersionSettings,
 ) {
     private val contractSerializer = FirContractSerializer()
+    private val extensionDeclarationProviders = session.extensionService.declarationForMetadataProviders
 
     fun packagePartProto(packageFqName: FqName, files: List<FirFile>): ProtoBuf.Package.Builder {
         val builder = ProtoBuf.Package.newBuilder()
 
+        fun addDeclaration(declaration: FirDeclaration, onUnsupportedDeclaration: (FirDeclaration) -> Unit) {
+            when (declaration) {
+                is FirProperty -> propertyProto(declaration)?.let { builder.addProperty(it) }
+                is FirSimpleFunction -> functionProto(declaration)?.let { builder.addFunction(it) }
+                is FirTypeAlias -> typeAliasProto(declaration)?.let { builder.addTypeAlias(it) }
+                else -> onUnsupportedDeclaration(declaration)
+            }
+        }
+
         for (file in files) {
             for (declaration in file.declarations) {
-                when (declaration) {
-                    is FirProperty -> propertyProto(declaration)?.let { builder.addProperty(it) }
-                    is FirSimpleFunction -> functionProto(declaration)?.let { builder.addFunction(it) }
-                    is FirTypeAlias -> typeAliasProto(declaration)?.let { builder.addTypeAlias(it) }
-                    else -> {}
-                }
+                addDeclaration(declaration) {}
             }
         }
 
         extension.serializePackage(packageFqName, builder)
+        for (extensionProvider in extensionDeclarationProviders) {
+            for (declaration in extensionProvider.provideTopLevelDeclarations(packageFqName, scopeSession)) {
+                addDeclaration(declaration) {
+                    error("Unsupported top-level declaration type: ${it.render()}")
+                }
+            }
+        }
 
         typeTable.serialize()?.let { builder.typeTable = it }
         versionRequirementTable?.serialize()?.let { builder.versionRequirementTable = it }
@@ -135,15 +148,33 @@ class FirElementSerializer private constructor(
             }
         }
 
+        val providedCallables = mutableListOf<FirCallableDeclaration>()
+        val providedConstructors = mutableListOf<FirConstructor>()
+        val providedNestedClassifiers = mutableListOf<FirClassifierSymbol<*>>()
+
+        for (extensionProvider in extensionDeclarationProviders) {
+            for (declaration in extensionProvider.provideDeclarationsForClass(klass, scopeSession)) {
+                when (declaration) {
+                    is FirConstructor -> providedConstructors += declaration
+                    is FirCallableDeclaration -> providedCallables += declaration
+                    is FirClassLikeDeclaration -> providedNestedClassifiers += declaration.symbol
+                    else -> error("Unsupported declaration type in: ${klass.render()} ${declaration.render()}")
+                }
+            }
+        }
+
         if (regularClass != null && regularClass.classKind != ClassKind.ENUM_ENTRY) {
             for (constructor in regularClass.constructors()) {
                 builder.addConstructor(constructorProto(constructor))
             }
+            for (constructor in providedConstructors) {
+                builder.addConstructor(constructorProto(constructor))
+            }
         }
 
         val callableMembers =
             extension.customClassMembersProducer?.getCallableMembers(klass)
-                ?: klass.declarations()
+                ?: (klass.memberDeclarations() + providedCallables)
                     .sortedWith(FirCallableDeclarationComparator)
 
         for (declaration in callableMembers) {
@@ -160,6 +191,7 @@ class FirElementSerializer private constructor(
             val scope = defaultType().scope(session, scopeSession, FakeOverrideTypeCalculator.DoNothing, requiredPhase = null) ?: return emptyList()
             return buildList {
                 scope.getClassifierNames().mapNotNullTo(this) { scope.getSingleClassifier(it) }
+                addAll(providedNestedClassifiers)
             }
         }
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationsForMetadataProviderExtension.kt

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2023 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.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.name.FqName
+import kotlin.reflect.KClass
+
+/**
+ * This extension can be used to provide information about declarations, which should be written in Kotlin metadata (i.e. be part of a module's ABI),
+ *   but for some reason can not (or should not) be generated with [FirDeclarationGenerationExtension]
+ *
+ * Example: assume your plugin generates some constructor in [IrGenerationExtension] which have value parameters matching
+ *   all properties of a class, and this constructor is used only as an implementation detail (so the only actor who accesses it
+ *   is a plugin itself). The constructor should be accessible from another module, so it should be present in the metadata. But you
+ *   can not generate this constructor in [FirDeclarationGenerationExtension], because it depends on types of properties,
+ *   which may be not resolved at the moment of constructor creation.
+ *
+ *
+ *   // MODULE: a
+ *   open class Base {
+ *       val x: Int = 1
+ *       val y = "hello"
+ *
+ *       constructor()
+ *
+ *       // generated
+ *       constructor(x: Int, y: String) { ... } // (1)
+ *   }
+ *
+ *   // MODULE: b(a)
+ *   class Derived : Base {
+ *       val z = 1.0
+ *
+ *       constructor() : super()
+ *
+ *       // generated
+ *       constructor(x: Int, y: String, z: Double) : super(x, y) { ... } // (2)
+ *
+ *       // constructor (1) should be presented in metadata of class Base, so IR plugin
+ *       //  can reference it during generation of constructor (2)
+ *  }
+ *
+ * All declarations provided by this extension should be fully resolved and contain all sub declarations if they exist.
+ *   E.g. if you want to provide some class then this class should contain all declarations of this class you want to be
+ *   present in metadata in `declarations` field of a `FirRegularClass`.
+ *   [FirDeclarationGenerationExtension] won't be called for classes provided by this extension.
+ */
+abstract class FirDeclarationsForMetadataProviderExtension(session: FirSession) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("SyntheticDeclarationsForMetadataGenerationExtension")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension> = FirDeclarationsForMetadataProviderExtension::class
+
+    /**
+     * It's allowed to provide only functions, properties and type aliases
+     */
+    open fun provideTopLevelDeclarations(packageFqName: FqName, scopeSession: ScopeSession): List<FirDeclaration> = emptyList()
+
+    open fun provideDeclarationsForClass(klass: FirClass, scopeSession: ScopeSession): List<FirDeclaration> = emptyList()
+
+    fun interface Factory : FirExtension.Factory<FirDeclarationsForMetadataProviderExtension>
+}
+
+val FirExtensionService.declarationForMetadataProviders: List<FirDeclarationsForMetadataProviderExtension> by FirExtensionService.registeredExtensions()