[FIR] Change API for generating nested classes from plugins

`FirDeclarationGenerationExtension.generateClassLikeDeclaration` was split into two functions: one for generating top level classes, and one for nested classes. Such change reduces verbosity and error-proness of this extension and also allows to smoothly run plugins on local classes ^KT-55248 Fixed
2022-12-21 10:29:54 +02:00
parent faa96ec7c0
commit 6783621eb0
31 changed files with 458 additions and 162 deletions
@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.FirLazyValue
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
-import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassDeclaredMemberScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
@@ -37,7 +36,9 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
     *
     * If classId has `outerClassId.Companion` format then generated class should be a companion object
     */
-    open fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
+    open fun generateTopLevelClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
+
+    open fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? = null

    // Can be called on STATUS stage
    open fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> = emptyList()
@@ -5,16 +5,15 @@

 package org.jetbrains.kotlin.fir.extensions

-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.caches.*
 import org.jetbrains.kotlin.fir.declarations.validate
 import org.jetbrains.kotlin.fir.ownerGenerator
-import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
-import org.jetbrains.kotlin.fir.scopes.impl.groupExtensionsByName
+import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
+import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
@@ -59,17 +58,6 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
            extensions.flatGroupBy { it.topLevelClassIdsCache.getValue() }
        }

-    private val extensionsByNestedClassifierClassId: FirCache<ClassId, Map<ClassId, List<FirDeclarationGenerationExtension>>, Nothing?> =
-        session.firCachesFactory.createCache cache@{ outerClassId, _ ->
-            val outerClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(outerClassId) as? FirClassSymbol<*>
-                ?: return@cache emptyMap()
-            session.groupExtensionsByName(
-                outerClassSymbol.fir,
-                nameExtractor = { nestedClassifierNamesCache.getValue(outerClassSymbol) },
-                nameTransformer = { outerClassId.createNestedClassId(it) }
-            )
-        }
-
    private val extensionsByTopLevelCallableId: FirLazyValue<Map<CallableId, List<FirDeclarationGenerationExtension>>, Nothing?> =
        session.firCachesFactory.createLazyValue {
            extensions.flatGroupBy { it.topLevelCallableIdsCache.getValue() }
@@ -78,21 +66,34 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
    // ------------------------------------------ generators ------------------------------------------

    private fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        val matchedExtensions = when {
-            classId.isNestedClass -> extensionsByNestedClassifierClassId.getValue(classId.outerClassId!!)[classId]
-            else -> extensionsByTopLevelClassId.getValue()[classId]
-        } ?: return null
-        val generatedClasses = matchedExtensions
-            .mapNotNull { generatorExtension ->
-                generatorExtension.generateClassLikeDeclaration(classId)?.also { symbol ->
-                    symbol.fir.ownerGenerator = generatorExtension
+        return when {
+            classId.isLocal -> null
+            classId.isNestedClass -> {
+                val owner = session.symbolProvider.getClassLikeSymbolByClassId(classId.outerClassId!!) as? FirClassSymbol<*> ?: return null
+                val nestedClassifierScope = session.nestedClassifierScope(owner.fir) ?: return null
+                var result: FirClassLikeSymbol<*>? = null
+                nestedClassifierScope.processClassifiersByName(classId.shortClassName) {
+                    if (it is FirClassLikeSymbol<*>) {
+                        result = it
+                    }
+                }
+                result
+            }
+            else -> {
+                val matchedExtensions = extensionsByTopLevelClassId.getValue()[classId] ?: return null
+                val generatedClasses = matchedExtensions
+                    .mapNotNull { generatorExtension ->
+                        generatorExtension.generateTopLevelClassLikeDeclaration(classId)?.also { symbol ->
+                            symbol.fir.ownerGenerator = generatorExtension
+                        }
+                    }
+                    .onEach { it.fir.validate() }
+                when (generatedClasses.size) {
+                    0 -> null
+                    1 -> generatedClasses.first()
+                    else -> error("Multiple plugins generated classes with same classId $classId\n${generatedClasses.joinToString("\n") { it.fir.render() }}")
                }
            }
-            .onEach { it.fir.validate() }
-        return when (generatedClasses.size) {
-            0 -> null
-            1 -> generatedClasses.first()
-            else -> error("Multiple plugins generated classes with same classId $classId\n${generatedClasses.joinToString("\n") { it.fir.render() }}")
        }
    }

@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.fir.extensions
 import org.jetbrains.kotlin.fir.FirSession
 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.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -24,7 +21,7 @@ import org.jetbrains.kotlin.name.Name
 * This is also legal, because plugins can not generate annotation classes which can influence other plugins or this plugin itself
 */
 class FirSwitchableExtensionDeclarationsSymbolProvider private constructor(
-    private val delegate: FirSymbolProvider
+    private val delegate: FirExtensionDeclarationsSymbolProvider
 ) : FirSymbolProvider(delegate.session) {
    companion object {
        fun create(session: FirSession): FirSwitchableExtensionDeclarationsSymbolProvider? {
@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.declarations.validate
 import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.ownerGenerator
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
@@ -160,16 +161,18 @@ internal fun FirSession.getExtensionsForClass(klass: FirClass): List<FirDeclarat
 class FirGeneratedClassNestedClassifierScope private constructor(
    useSiteSession: FirSession,
    klass: FirClass,
-    private val nestedClassifierNames: Set<Name>
+    private val extensionsByName: Map<Name, List<FirDeclarationGenerationExtension>>,
 ) : FirNestedClassifierScope(klass, useSiteSession) {
    companion object {
        @OptIn(FirExtensionApiInternals::class)
        fun create(useSiteSession: FirSession, klass: FirClass): FirGeneratedClassNestedClassifierScope? {
-            val extensions = useSiteSession.getExtensionsForClass(klass)
            val symbol = klass.symbol
-            val classifierNames = extensions.flatMapTo(mutableSetOf()) { it.nestedClassifierNamesCache.getValue(symbol) }
-            if (classifierNames.isEmpty()) return null
-            return FirGeneratedClassNestedClassifierScope(useSiteSession, klass, classifierNames)
+
+            val extensionsByName = useSiteSession.getExtensionsForClass(klass).flatGroupBy {
+                it.nestedClassifierNamesCache.getValue(symbol)
+            }
+            if (extensionsByName.isEmpty()) return null
+            return FirGeneratedClassNestedClassifierScope(useSiteSession, klass, extensionsByName)
        }
    }

@@ -186,10 +189,25 @@ class FirGeneratedClassNestedClassifierScope private constructor(
            }
        }

-        if (name !in nestedClassifierNames) return null
-        val generatedClass = useSiteSession.generatedDeclarationsSymbolProvider
-            ?.getClassLikeSymbolByClassId(klass.classId.createNestedClassId(name))
-        require(generatedClass is FirRegularClassSymbol?) { "Only regular class are allowed as nested classes" }
+        val extensions = extensionsByName[name] ?: return null
+
+        val generatedClasses = extensions.mapNotNull { extension ->
+            extension.generateNestedClassLikeDeclaration(klass.symbol, name)?.also { symbol ->
+                symbol.fir.ownerGenerator = extension
+            }
+        }
+
+        val generatedClass = when (generatedClasses.size) {
+            0 -> return null
+            1 -> generatedClasses.first()
+            else -> error(
+                """
+                     Multiple plugins generated nested class with same name $name for class ${klass.classId}:
+                    ${generatedClasses.joinToString("\n") { it.fir.render() }}
+                """.trimIndent()
+            )
+        }
+        require(generatedClass is FirRegularClassSymbol) { "Only regular class are allowed as nested classes" }
        return generatedClass
    }

@@ -198,10 +216,10 @@ class FirGeneratedClassNestedClassifierScope private constructor(
    }

    override fun isEmpty(): Boolean {
-        return nestedClassifierNames.isEmpty()
+        return extensionsByName.isEmpty()
    }

    override fun getClassifierNames(): Set<Name> {
-        return nestedClassifierNames
+        return extensionsByName.keys
    }
 }