[FIR] Rewrite algorithm of generating declarations by plugins

Before there was separate global stage which asks extensions
  to generate new declarations. Now it is changed to provider
  like requests, so extensions will decide do they have declarations
  with specific classId/callableId or not

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.checkers.registerJvmCheckers
 import org.jetbrains.kotlin.fir.deserialization.ModuleDataProvider
 import org.jetbrains.kotlin.fir.deserialization.SingleModuleDataProvider
 import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
+import org.jetbrains.kotlin.fir.extensions.FirExtensionDeclarationsSymbolProvider
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirCliSession
@@ -157,6 +158,12 @@ object FirSessionFactory {
                 )
             }
 
+            FirSessionConfigurator(this).apply {
+                registerCommonCheckers()
+                registerJvmCheckers()
+                init()
+            }.configure()
+
             val dependenciesSymbolProvider = FirDependenciesSymbolProviderImpl(this)
             register(
                 FirSymbolProvider::class,
@@ -165,6 +172,7 @@ object FirSessionFactory {
                     listOfNotNull(
                         firProvider.symbolProvider,
                         symbolProviderForBinariesFromIncrementalCompilation,
+                        FirExtensionDeclarationsSymbolProvider.create(this),
                         JavaSymbolProvider(this, projectEnvironment.getFirJavaFacade(this, moduleData, scope)),
                         dependenciesSymbolProvider,
                     )
@@ -177,11 +185,6 @@ object FirSessionFactory {
             )
 
             FirJvmDefaultErrorMessages.installJvmErrorMessages()
-            FirSessionConfigurator(this).apply {
-                registerCommonCheckers()
-                registerJvmCheckers()
-                init()
-            }.configure()
             projectEnvironment.registerAsJavaElementFinder(this)
         }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt

@@ -6,10 +6,13 @@
 package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 import kotlin.reflect.KClass
 
 /*
@@ -26,19 +29,10 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredi
 
     final override val extensionType: KClass<out FirExtension> = FirDeclarationGenerationExtension::class
 
-    abstract fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>>
-
-    abstract fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration>
-
-    abstract fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>>
-
-    data class GeneratedDeclaration<out T : FirDeclaration>(val newDeclaration: T, val owner: FirAnnotatedDeclaration)
+    open fun generateClassLikeDeclaration(classId: ClassId, owner: FirClassSymbol<*>?): FirClassLikeSymbol<*>? = null
+    open fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> = emptyList()
+    open fun generateProperties(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirPropertySymbol> = emptyList()
+    open fun hasPackage(packageFqName: FqName): Boolean = false
 
     fun interface Factory : FirExtension.Factory<FirDeclarationGenerationExtension>
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionDeclarationsSymbolProvider.kt

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2010-2021 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.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.FirCachesFactory
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.declarations.validate
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+class FirExtensionDeclarationsSymbolProvider private constructor(
+    session: FirSession,
+    cachesFactory: FirCachesFactory,
+    private val extensions: List<FirDeclarationGenerationExtension>
+) : FirSymbolProvider(session) {
+    companion object {
+        fun create(session: FirSession): FirExtensionDeclarationsSymbolProvider? {
+            val extensions = session.extensionService.declarationGenerators
+            if (extensions.isEmpty()) return null
+            return FirExtensionDeclarationsSymbolProvider(session, session.firCachesFactory, extensions)
+        }
+    }
+
+    // ------------------------------------------ caches ------------------------------------------
+
+    private val classCache: FirCache<ClassId, FirClassLikeSymbol<*>?, Nothing?> = cachesFactory.createCache { classId, _ ->
+        generateClassLikeDeclaration(classId)
+    }
+
+    private val functionCache: FirCache<CallableId, List<FirNamedFunctionSymbol>, Nothing?> = cachesFactory.createCache { callableId, _ ->
+        generateTopLevelFunctions(callableId)
+    }
+
+    private val propertyCache: FirCache<CallableId, List<FirPropertySymbol>, Nothing?> = cachesFactory.createCache { callableId, _ ->
+        generateTopLevelProperties(callableId)
+    }
+
+    private val packageCache: FirCache<FqName, Boolean, Nothing?> = cachesFactory.createCache { packageFqName, _ ->
+        hasPackage(packageFqName)
+    }
+
+    // ------------------------------------------ generators ------------------------------------------
+
+    private fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
+        // TODO: what we should do if multiple extensions want to generate class with same classId?
+        return extensions.firstNotNullOfOrNull { it.generateClassLikeDeclaration(classId, owner = null) }?.also { it.fir.validate() }
+    }
+
+    private fun generateTopLevelFunctions(callableId: CallableId): List<FirNamedFunctionSymbol> {
+        return extensions.flatMap { it.generateFunctions(callableId, owner = null) }.onEach { it.fir.validate() }
+    }
+
+    private fun generateTopLevelProperties(callableId: CallableId): List<FirPropertySymbol> {
+        return extensions.flatMap { it.generateProperties(callableId, owner = null) }.onEach { it.fir.validate() }
+    }
+
+    private fun hasPackage(packageFqName: FqName): Boolean {
+        return extensions.any { it.hasPackage(packageFqName) }
+    }
+
+    // ------------------------------------------ provider methods ------------------------------------------
+
+    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
+        return classCache.getValue(classId, context = null)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
+        destination += functionCache.getValue(CallableId(packageFqName, name), context = null)
+        destination += propertyCache.getValue(CallableId(packageFqName, name), context = null)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
+        destination += functionCache.getValue(CallableId(packageFqName, name), context = null)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
+        destination += propertyCache.getValue(CallableId(packageFqName, name), context = null)
+    }
+
+    override fun getPackage(fqName: FqName): FqName? {
+        return fqName.takeIf { packageCache.getValue(fqName, null) }
+    }
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirTotalResolveProcessor.kt

@@ -66,7 +66,6 @@ fun FirResolvePhase.createCompilerProcessorByPhase(
     return when (this) {
         RAW_FIR -> throw IllegalArgumentException("Raw FIR building phase does not have a transformer")
         ANNOTATIONS_FOR_PLUGINS -> FirPluginAnnotationsResolveProcessor(session, scopeSession)
-        CLASS_GENERATION -> FirGlobalClassGenerationProcessor(session, scopeSession)
         IMPORTS -> FirImportResolveProcessor(session, scopeSession)
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
@@ -74,7 +73,6 @@ fun FirResolvePhase.createCompilerProcessorByPhase(
         STATUS -> FirStatusResolveProcessor(session, scopeSession)
         ARGUMENTS_OF_ANNOTATIONS -> FirAnnotationArgumentsResolveProcessor(session, scopeSession)
         CONTRACTS -> FirContractResolveProcessor(session, scopeSession)
-        NEW_MEMBERS_GENERATION -> FirGlobalNewMemberGenerationProcessor(session, scopeSession)
         IMPLICIT_TYPES_BODY_RESOLVE -> FirImplicitTypeBodyResolveProcessor(session, scopeSession)
         BODY_RESOLVE -> FirBodyResolveProcessor(session, scopeSession)
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalClassGenerationProcessor.kt

@@ -1,70 +0,0 @@
-/*
- * Copyright 2010-2020 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.resolve.transformers.plugin
-
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.utils.addDeclaration
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.declarationGenerators
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.firProvider
-import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
-import org.jetbrains.kotlin.fir.resolve.transformers.FirGlobalResolveProcessor
-
-class FirGlobalClassGenerationProcessor(
-    session: FirSession,
-    scopeSession: ScopeSession
-) : FirGlobalResolveProcessor(session, scopeSession) {
-    override fun process(files: Collection<FirFile>) {
-        val extensions = session.extensionService.declarationGenerators
-        if (extensions.isEmpty()) return
-        val provider = session.predicateBasedProvider
-        for (extension in extensions) {
-            var annotatedDeclarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
-            while (annotatedDeclarations.isNotEmpty()) {
-                val newClasses = generateClasses(annotatedDeclarations, extension)
-                annotatedDeclarations = provider.getSymbolsWithOwnersByPredicate(newClasses, extension.predicate)
-            }
-        }
-    }
-
-    private fun generateClasses(
-        declarations: List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>>,
-        extension: FirDeclarationGenerationExtension,
-    ): List<FirRegularClass> {
-        val newClasses = mutableListOf<FirRegularClass>()
-        for ((declaration, owners) in declarations) {
-            generateClass(extension, declaration, owners, newClasses)
-        }
-        return newClasses
-    }
-
-    @OptIn(FirProviderInternals::class)
-    private fun generateClass(
-        extension: FirDeclarationGenerationExtension,
-        declaration: FirAnnotatedDeclaration,
-        owners: List<FirAnnotatedDeclaration>,
-        newClasses: MutableList<FirRegularClass>
-    ) {
-        val generatedClasses = extension.generateClasses(declaration, owners)
-        for ((klass, owner) in generatedClasses) {
-            when (owner) {
-                is FirRegularClass -> owner.addDeclaration(klass)
-                is FirFile -> owner.addDeclaration(klass)
-                else -> {}
-            }
-            session.generatedClassIndex.registerClass(klass, owner)
-            session.predicateBasedProvider.registerGeneratedDeclaration(klass, owner)
-            session.firProvider.recordGeneratedClass(owner, klass)
-            newClasses += klass
-        }
-    }
-}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalNewMemberGenerationProcessor.kt

@@ -1,67 +0,0 @@
-/*
- * Copyright 2010-2020 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.resolve.transformers.plugin
-
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.utils.addDeclaration
-import org.jetbrains.kotlin.fir.declarations.validate
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.declarationGenerators
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.firProvider
-import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
-import org.jetbrains.kotlin.fir.resolve.transformers.FirGlobalResolveProcessor
-
-class FirGlobalNewMemberGenerationProcessor(
-    session: FirSession,
-    scopeSession: ScopeSession
-) : FirGlobalResolveProcessor(session, scopeSession) {
-    private val index = session.generatedClassIndex
-    private val provider = session.predicateBasedProvider
-
-    override fun process(files: Collection<FirFile>) {
-        val extensions = session.extensionService.declarationGenerators
-        if (extensions.isEmpty()) return
-        for (extension in extensions) {
-            generateNewMembers(extension)
-            fillGeneratedClasses(extension)
-        }
-    }
-
-    @OptIn(FirProviderInternals::class)
-    private fun generateNewMembers(extension: FirDeclarationGenerationExtension) {
-        val declarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
-        for ((declaration, owners) in declarations) {
-            val newMembers = extension.generateMembers(declaration, owners)
-            for ((newMember, owner) in newMembers) {
-                newMember.validate()
-                when (owner) {
-                    is FirRegularClass -> owner.addDeclaration(newMember)
-                    is FirFile -> owner.addDeclaration(newMember)
-                    else -> error("Should not be here")
-                }
-                session.firProvider.recordGeneratedMember(owner, newMember)
-            }
-        }
-    }
-
-    @OptIn(FirProviderInternals::class)
-    private fun fillGeneratedClasses(extension: FirDeclarationGenerationExtension) {
-        for (generatedClass in index[extension.key]) {
-            val klass = generatedClass.klass
-            val newMembers = extension.generateMembersForGeneratedClass(generatedClass)
-            for (newMember in newMembers) {
-                newMember.validate()
-                klass.addDeclaration(newMember)
-                session.firProvider.recordGeneratedMember(klass, newMember)
-            }
-        }
-    }
-}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationOrigin.kt

@@ -25,4 +25,6 @@ sealed class FirDeclarationOrigin(private val displayName: String? = null, val f
     }
 }
 
-abstract class FirPluginKey
+abstract class FirPluginKey {
+    val origin: FirDeclarationOrigin = FirDeclarationOrigin.Plugin(this)
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirResolvePhase.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.declarations
 enum class FirResolvePhase(val pluginPhase: Boolean = false, val noProcessor: Boolean = false) {
     RAW_FIR(noProcessor = true),
     ANNOTATIONS_FOR_PLUGINS(pluginPhase = true),
-    CLASS_GENERATION(pluginPhase = true),
     IMPORTS,
     SUPER_TYPES,
     SEALED_CLASS_INHERITORS,
@@ -16,7 +15,6 @@ enum class FirResolvePhase(val pluginPhase: Boolean = false, val noProcessor: Bo
     STATUS,
     ARGUMENTS_OF_ANNOTATIONS,
     CONTRACTS,
-    NEW_MEMBERS_GENERATION(pluginPhase = true),
     IMPLICIT_TYPES_BODY_RESOLVE,
     BODY_RESOLVE;