[FIR] Introduce caching service for metadata declarations provided by extensions

This is needed to allow to call for provided declarations multiple times
  with only one real invocation of plugins

compiler/fir/entrypoint/build.gradle.kts

@@ -17,6 +17,7 @@ dependencies {
     api(project(":js:js.frontend"))
 
     implementation(project(":compiler:fir:resolve"))
+    implementation(project(":compiler:fir:fir-serialization"))
     implementation(project(":compiler:fir:fir2ir:jvm-backend"))
     implementation(project(":compiler:backend.jvm"))
     implementation(project(":compiler:ir.serialization.common"))

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt

@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.fir.scopes.FirOverrideService
 import org.jetbrains.kotlin.fir.scopes.FirPlatformClassMapper
 import org.jetbrains.kotlin.fir.scopes.PlatformSpecificOverridabilityRules
 import org.jetbrains.kotlin.fir.scopes.impl.*
+import org.jetbrains.kotlin.fir.serialization.FirProvidedDeclarationsForMetadataService
 import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache
 import org.jetbrains.kotlin.fir.types.FirFunctionTypeKindService
@@ -78,6 +79,7 @@ fun FirSession.registerCommonComponents(languageVersionSettings: LanguageVersion
 @OptIn(SessionConfiguration::class)
 fun FirSession.registerCommonComponentsAfterExtensionsAreConfigured() {
     register(FirFunctionTypeKindService::class, FirFunctionTypeKindServiceImpl(this))
+    register(FirProvidedDeclarationsForMetadataService::class, FirProvidedDeclarationsForMetadataService.create(this))
 }
 
 @OptIn(SessionConfiguration::class)

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

@@ -25,7 +25,6 @@ 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.*
@@ -70,7 +69,7 @@ class FirElementSerializer private constructor(
     private val languageVersionSettings: LanguageVersionSettings,
 ) {
     private val contractSerializer = FirContractSerializer()
-    private val extensionDeclarationProviders = session.extensionService.declarationForMetadataProviders
+    private val providedDeclarationsService = session.providedDeclarationsForMetadataService
 
     fun packagePartProto(
         packageFqName: FqName,
@@ -102,11 +101,9 @@ class FirElementSerializer private constructor(
         }
 
         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()}")
-                }
+        for (declaration in providedDeclarationsService.getProvidedTopLevelDeclarations(packageFqName, scopeSession)) {
+            addDeclaration(declaration) {
+                error("Unsupported top-level declaration type: ${it.render()}")
             }
         }
 
@@ -147,7 +144,8 @@ class FirElementSerializer private constructor(
             builder.addTypeParameter(typeParameterProto(typeParameter))
         }
 
-        val classId = klass.symbol.classId
+        val classSymbol = klass.symbol
+        val classId = classSymbol.classId
         if (classId != StandardClassIds.Any && classId != StandardClassIds.Nothing) {
             // Special classes (Any, Nothing) have no supertypes
             for (superTypeRef in klass.superTypeRefs) {
@@ -159,33 +157,19 @@ 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) {
+
+            for (constructor in providedDeclarationsService.getProvidedConstructors(classSymbol, scopeSession)) {
                 builder.addConstructor(constructorProto(constructor))
             }
         }
 
         val callableMembers =
             extension.customClassMembersProducer?.getCallableMembers(klass)
-                ?: (klass.memberDeclarations() + providedCallables)
+                ?: (klass.memberDeclarations() + providedDeclarationsService.getProvidedCallables(classSymbol, scopeSession))
                     .sortedWith(FirCallableDeclarationComparator)
 
         for (declaration in callableMembers) {
@@ -203,7 +187,7 @@ class FirElementSerializer private constructor(
                 defaultType().scope(session, scopeSession, FakeOverrideTypeCalculator.DoNothing, requiredPhase = null) ?: return emptyList()
             return buildList {
                 scope.getClassifierNames().mapNotNullTo(this) { scope.getSingleClassifier(it) }
-                addAll(providedNestedClassifiers)
+                addAll(providedDeclarationsService.getProvidedNestedClassifiers(classSymbol, scopeSession))
             }
         }
 

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

@@ -0,0 +1,118 @@
+/*
+ * 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.serialization
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationsForMetadataProviderExtension
+import org.jetbrains.kotlin.fir.extensions.declarationForMetadataProviders
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.name.FqName
+
+abstract class FirProvidedDeclarationsForMetadataService : FirSessionComponent {
+    companion object {
+        fun create(session: FirSession): FirProvidedDeclarationsForMetadataService {
+            val extensionProviders = session.extensionService.declarationForMetadataProviders
+            return if (extensionProviders.isEmpty()) Empty else FirProvidedDeclarationsForMetadataServiceImpl(session, extensionProviders)
+        }
+    }
+
+    abstract fun getProvidedTopLevelDeclarations(packageFqName: FqName, scopeSession: ScopeSession): List<FirDeclaration>
+    abstract fun getProvidedConstructors(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirConstructor>
+    abstract fun getProvidedCallables(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirCallableDeclaration>
+    abstract fun getProvidedNestedClassifiers(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirClassLikeSymbol<*>>
+
+    private object Empty : FirProvidedDeclarationsForMetadataService() {
+        override fun getProvidedTopLevelDeclarations(packageFqName: FqName, scopeSession: ScopeSession): List<FirDeclaration> {
+            return emptyList()
+        }
+
+        override fun getProvidedConstructors(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirConstructor> {
+            return emptyList()
+        }
+
+        override fun getProvidedCallables(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirCallableDeclaration> {
+            return emptyList()
+        }
+
+        override fun getProvidedNestedClassifiers(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirClassLikeSymbol<*>> {
+            return emptyList()
+        }
+    }
+}
+
+private class FirProvidedDeclarationsForMetadataServiceImpl(
+    session: FirSession,
+    private val extensionDeclarationProviders: List<FirDeclarationsForMetadataProviderExtension>
+) : FirProvidedDeclarationsForMetadataService() {
+    private val cachesFactory = session.firCachesFactory
+
+    private val topLevelsCache: FirCache<FqName, List<FirDeclaration>, ScopeSession> =
+        cachesFactory.createCache(::computeTopLevelDeclarations)
+
+    private val membersCache: FirCache<FirClassSymbol<*>, ClassDeclarations, ScopeSession> =
+        cachesFactory.createCache(::computeMemberDeclarations)
+
+    private fun computeTopLevelDeclarations(packageFqName: FqName, scopeSession: ScopeSession): List<FirDeclaration> {
+        return buildList {
+            for (extensionProvider in extensionDeclarationProviders) {
+                for (declaration in extensionProvider.provideTopLevelDeclarations(packageFqName, scopeSession)) {
+                    add(declaration)
+                }
+            }
+        }
+    }
+
+    override fun getProvidedTopLevelDeclarations(packageFqName: FqName, scopeSession: ScopeSession): List<FirDeclaration> {
+        return topLevelsCache.getValue(packageFqName, scopeSession)
+    }
+
+    override fun getProvidedConstructors(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirConstructor> {
+        return membersCache.getValue(owner, scopeSession).providedConstructors
+    }
+
+    override fun getProvidedCallables(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirCallableDeclaration> {
+        return membersCache.getValue(owner, scopeSession).providedCallables
+    }
+
+    override fun getProvidedNestedClassifiers(owner: FirClassSymbol<*>, scopeSession: ScopeSession): List<FirClassLikeSymbol<*>> {
+        return membersCache.getValue(owner, scopeSession).providedNestedClasses
+    }
+
+    private data class ClassDeclarations(
+        val providedCallables: List<FirCallableDeclaration>,
+        val providedConstructors: List<FirConstructor>,
+        val providedNestedClasses: List<FirClassLikeSymbol<*>>,
+    )
+
+    private fun computeMemberDeclarations(symbol: FirClassSymbol<*>, scopeSession: ScopeSession): ClassDeclarations {
+        val providedCallables = mutableListOf<FirCallableDeclaration>()
+        val providedConstructors = mutableListOf<FirConstructor>()
+        val providedNestedClassifiers = mutableListOf<FirClassLikeSymbol<*>>()
+
+        for (extensionProvider in extensionDeclarationProviders) {
+            for (declaration in extensionProvider.provideDeclarationsForClass(symbol.fir, scopeSession)) {
+                when (declaration) {
+                    is FirConstructor -> providedConstructors += declaration
+                    is FirCallableDeclaration -> providedCallables += declaration
+                    is FirClassLikeDeclaration -> providedNestedClassifiers += declaration.symbol
+                    else -> error("Unsupported declaration type in: $symbol ${declaration.render()}")
+                }
+            }
+        }
+
+        return ClassDeclarations(providedCallables, providedConstructors, providedNestedClassifiers)
+    }
+}
+
+val FirSession.providedDeclarationsForMetadataService: FirProvidedDeclarationsForMetadataService by FirSession.sessionComponentAccessor()