[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;
 

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirAllOpenComponentRegistrar.kt

@@ -15,10 +15,10 @@ class FirAllOpenComponentRegistrar : FirExtensionRegistrar() {
         +::AllOpenSupertypeGenerator
 
         // Declaration generators
-        +::AllOpenMemberGenerator
-        +::AllOpenNestedClassGenerator
+//        +::AllOpenMemberGenerator
+//        +::AllOpenNestedClassGenerator
         +::AllOpenAdditionalCheckers
         +::AllOpenTopLevelDeclarationsGenerator
-        +::AllOpenRecursiveNestedClassGenerator
+//        +::AllOpenRecursiveNestedClassGenerator
     }
-}
+}

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenMemberGenerator.kt

@@ -1,60 +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.plugin.generators
-
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
-import org.jetbrains.kotlin.fir.declarations.utils.classId
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
-import org.jetbrains.kotlin.fir.extensions.predicate.and
-import org.jetbrains.kotlin.fir.extensions.predicate.has
-import org.jetbrains.kotlin.fir.extensions.predicate.under
-import org.jetbrains.kotlin.fir.moduleData
-import org.jetbrains.kotlin.fir.plugin.AllOpenPluginKey
-import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.Name
-
-class AllOpenMemberGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        return emptyList()
-    }
-
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        return emptyList()
-    }
-
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        if (annotatedDeclaration !is FirProperty) return emptyList()
-        val owner = owners.last() as? FirRegularClass ?: return emptyList()
-        val propertyName = annotatedDeclaration.name.identifier
-        val function = buildSimpleFunction {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = annotatedDeclaration.returnTypeRef
-            status = annotatedDeclaration.status
-            name = Name.identifier("hello${propertyName.replaceFirstChar(Char::uppercaseChar)}")
-            symbol = FirNamedFunctionSymbol(CallableId(owner.classId, name))
-        }
-        return listOf(GeneratedDeclaration(function, owner))
-    }
-
-    override val predicate: DeclarationPredicate = under("WithGenerated".fqn()) and has("WithHello".fqn())
-
-    override val key: FirPluginKey
-        get() = AllOpenPluginKey
-}

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenNestedClassGenerator.kt

@@ -1,105 +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.plugin.generators
-
-import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.descriptors.EffectiveVisibility
-import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.descriptors.Visibilities
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildConstructor
-import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
-import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
-import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
-import org.jetbrains.kotlin.fir.extensions.predicate.has
-import org.jetbrains.kotlin.fir.moduleData
-import org.jetbrains.kotlin.fir.plugin.AllOpenPluginKey
-import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
-import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.Name
-
-class AllOpenNestedClassGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        val owner = annotatedDeclaration as? FirRegularClass ?: return emptyList()
-        val newClass = buildRegularClass {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Private,
-                Modality.FINAL,
-                EffectiveVisibility.PrivateInClass
-            ).apply {
-                isInner = true
-            }
-            classKind = ClassKind.CLASS
-            name = Name.identifier("Foo")
-            symbol = FirRegularClassSymbol(owner.symbol.classId.createNestedClassId(name))
-            scopeProvider = owner.scopeProvider
-        }
-        return listOf(GeneratedDeclaration(newClass, owner))
-    }
-
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        val klass = generatedClass.klass
-
-        val classId = klass.symbol.classId
-        val constructor = buildConstructor {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = buildResolvedTypeRef {
-                type = ConeClassLikeTypeImpl(ConeClassLikeLookupTagImpl(classId), emptyArray(), isNullable = false)
-            }
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            symbol = FirConstructorSymbol(CallableId(classId, classId.shortClassName))
-        }
-
-        val function = buildSimpleFunction {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = session.builtinTypes.intType
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            name = Name.identifier("hello")
-            symbol = FirNamedFunctionSymbol(CallableId(classId, name))
-        }
-        return listOf(constructor, function)
-    }
-
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        return emptyList()
-    }
-
-    override val predicate: DeclarationPredicate = has("WithNestedFoo".fqn())
-
-    override val key: FirPluginKey
-        get() = AllOpenPluginKey
-}

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenRecursiveNestedClassGenerator.kt

@@ -1,87 +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.plugin.generators
-
-import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.descriptors.EffectiveVisibility
-import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.descriptors.Visibilities
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
-import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
-import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
-import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyAnnotationArgumentMapping
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
-import org.jetbrains.kotlin.fir.extensions.predicate.has
-import org.jetbrains.kotlin.fir.moduleData
-import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.Name
-
-class AllOpenRecursiveNestedClassGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        if (owners.size > 2) return emptyList()
-        val owner = annotatedDeclaration as? FirRegularClass ?: return emptyList()
-        val newClass = buildRegularClass {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            classKind = ClassKind.CLASS
-            name = Name.identifier("Nested")
-            symbol = FirRegularClassSymbol(owner.symbol.classId.createNestedClassId(name))
-            scopeProvider = owner.scopeProvider
-            annotations += buildAnnotation {
-                annotationTypeRef = buildResolvedTypeRef {
-                    type = ConeClassLikeTypeImpl(
-                        ConeClassLikeLookupTagImpl(annotationClassId),
-                        emptyArray(),
-                        isNullable = false
-                    )
-                }
-                argumentMapping = FirEmptyAnnotationArgumentMapping
-            }
-        }
-        return listOf(GeneratedDeclaration(newClass, owner))
-    }
-
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        return emptyList()
-    }
-
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        return emptyList()
-    }
-
-    override val key: FirPluginKey
-        get() = Key
-
-    override val predicate: DeclarationPredicate
-        get() = has("B".fqn())
-
-    private val annotationClassId = "B".fqn().let {
-        ClassId(it.parent(), it.shortName())
-    }
-
-    private object Key : FirPluginKey()
-}

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenTopLevelDeclarationsGenerator.kt

@@ -5,82 +5,74 @@
 
 package org.jetbrains.kotlin.fir.plugin.generators
 
-import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.EffectiveVisibility
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.has
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.moduleData
-import org.jetbrains.kotlin.fir.packageFqName
+import org.jetbrains.kotlin.fir.plugin.AllOpenPluginKey
 import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 class AllOpenTopLevelDeclarationsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        val file = owners.first() as FirFile
-        val klass = annotatedDeclaration as? FirRegularClass ?: return emptyList()
-        val newClass = buildRegularClass {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            classKind = ClassKind.OBJECT
-            name = Name.identifier("TopLevel${klass.name}")
-            symbol = FirRegularClassSymbol(ClassId(file.packageFqName, name))
-            scopeProvider = klass.scopeProvider
-        }
-        return listOf(GeneratedDeclaration(newClass, file))
+    private val predicateBasedProvider = session.predicateBasedProvider
+    private val matchedClasses by lazy {
+        predicateBasedProvider.getSymbolsByPredicate(predicate).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
     }
 
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        val klass = generatedClass.klass
+    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+        if (owner != null) return emptyList()
+        val matchedClassSymbol = findMatchedClassForFunction(callableId) ?: return emptyList()
         val function = buildSimpleFunction {
             moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = session.builtinTypes.intType
+            origin = key.origin
             status = FirResolvedDeclarationStatusImpl(
                 Visibilities.Public,
                 Modality.FINAL,
                 EffectiveVisibility.Public
             )
-            name = Name.identifier("hello")
-            symbol = FirNamedFunctionSymbol(CallableId(klass.symbol.classId, name))
+            returnTypeRef = session.builtinTypes.stringType
+            valueParameters += buildValueParameter {
+                moduleData = session.moduleData
+                origin = key.origin
+                returnTypeRef = buildResolvedTypeRef {
+                    type = ConeClassLikeTypeImpl(ConeClassLikeLookupTagImpl(matchedClassSymbol.classId), emptyArray(), isNullable = false)
+                }
+                name = Name.identifier("value")
+                symbol = FirValueParameterSymbol(name)
+                isCrossinline = false
+                isNoinline = false
+                isVararg = false
+            }
+            symbol = FirNamedFunctionSymbol(callableId)
+            name = callableId.callableName
         }
-        return listOf(function)
+        return listOf(function.symbol)
     }
 
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        return emptyList()
+    private fun findMatchedClassForFunction(callableId: CallableId): FirRegularClassSymbol? {
+        // We generate only top-level functions
+        if (callableId.classId != null) return null
+        return matchedClasses
+            .filter { it.classId.packageFqName == callableId.packageName }
+            .firstOrNull { callableId.callableName.identifier == "dummy${it.classId.shortClassName.identifier}" }
     }
 
-    override val key: FirPluginKey
-        get() = Key
+    override val key: AllOpenPluginKey
+        get() = AllOpenPluginKey
 
     override val predicate: DeclarationPredicate
         get() = has("A".fqn())
 
-    private object Key : FirPluginKey()
 }

plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelCallables.fir.txt

@@ -0,0 +1,14 @@
+FILE: topLevelCallables.kt
+    @R|org/jetbrains/kotlin/fir/plugin/A|() public final class MySuperClass : R|kotlin/Any| {
+        public constructor(): R|foo/MySuperClass| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun test(): R|kotlin/Unit| {
+            lval s: R|kotlin/String| = R|foo/dummyMySuperClass|(this@R|foo/MySuperClass|)
+            R|foo/takeString|(R|<local>/s|)
+        }
+
+    }
+    public final fun takeString(s: R|kotlin/String|): R|kotlin/Unit| {
+    }

plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelCallables.kt

@@ -0,0 +1,17 @@
+package foo
+
+import org.jetbrains.kotlin.fir.plugin.A
+
+/*
+ * Plugin generates `dummyClassName(value: ClassName): String` function for each class annotated with @A
+ */
+
+@A
+class MySuperClass {
+    fun test() {
+        val s = dummyMySuperClass(this)
+        takeString(s)
+    }
+}
+
+fun takeString(s: String) {}

plugins/fir/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/FirAllOpenDiagnosticTestGenerated.java

@@ -67,6 +67,12 @@ public class FirAllOpenDiagnosticTestGenerated extends AbstractFirAllOpenDiagnos
             runTest("plugins/fir/fir-plugin-prototype/testData/memberGen/recursiveNestedClasses.kt");
         }
 
+        @Test
+        @TestMetadata("topLevelCallables.kt")
+        public void testTopLevelCallables() throws Exception {
+            runTest("plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelCallables.kt");
+        }
+
         @Test
         @TestMetadata("topLevelClass.kt")
         public void testTopLevelClass() throws Exception {