[FIR] Add extension for generating new members for existing classes

2020-05-28 14:51:30 +03:00
parent bd194686df
commit 87f0f123a8
10 changed files with 221 additions and 53 deletions
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir
 import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedNestedClassIndex
 import org.jetbrains.kotlin.fir.scopes.impl.FirDeclaredMemberScopeProvider
 import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache

@@ -18,4 +19,5 @@ fun FirSession.registerCommonComponents() {
    register(FirExtensionService::class, FirExtensionService(this))
    register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.create(this))
    register(FirPredicateBasedProvider::class, FirPredicateBasedProvider.create(this))
+    register(GeneratedNestedClassIndex::class, GeneratedNestedClassIndex.create())
 }
@@ -1,44 +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.extensions
-
-import org.jetbrains.kotlin.fir.FirAnnotationContainer
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import kotlin.reflect.KClass
-
-/*
- * TODO:
- *  - add registration of new files
- *  - check that annotations or meta-annotations is not empty
- */
-abstract class FirClassGenerationExtension(session: FirSession) : FirPredicateBasedExtension(session) {
-    companion object {
-        val NAME = FirExtensionPointName("StatusTransformer")
-    }
-
-    final override val name: FirExtensionPointName
-        get() = NAME
-
-    final override val extensionType: KClass<out FirExtension> = FirClassGenerationExtension::class
-
-    abstract fun <T> generateClass(
-        containingFile: FirFile,
-        annotatedDeclaration: T
-    ): List<GeneratedClass> where T : FirDeclaration, T : FirAnnotationContainer
-
-    data class GeneratedClass(val klass: FirRegularClass, val container: FirDeclaration) {
-        init {
-            require(container is FirRegularClass || container is FirFile)
-        }
-    }
-
-    fun interface Factory : FirExtension.Factory<FirClassGenerationExtension>
-}
-
-val FirExtensionService.classGenerationExtensions: List<FirClassGenerationExtension> by FirExtensionService.registeredExtensions()
@@ -0,0 +1,46 @@
+/*
+ * 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.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.GeneratedNestedClass
+import kotlin.reflect.KClass
+
+/*
+ * TODO:
+ *  - check that annotations or meta-annotations is not empty
+ */
+abstract class FirExistingClassModificationExtension(session: FirSession) : FirPredicateBasedExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("ExistingClassModification")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension> = FirExistingClassModificationExtension::class
+
+    abstract fun generateNestedClasses(
+        annotatedDeclaration: FirDeclaration,
+        owners: List<FirAnnotatedDeclaration>
+    ): List<GeneratedDeclaration<FirRegularClass>>
+
+    abstract fun generateMembersForNestedClasses(generatedNestedClass: GeneratedNestedClass): List<FirDeclaration>
+
+    abstract fun generateMembers(
+        annotatedDeclaration: FirDeclaration,
+        owners: List<FirAnnotatedDeclaration>
+    ): List<GeneratedDeclaration<*>>
+
+    data class GeneratedDeclaration<out T : FirDeclaration>(val newDeclaration: T, val owner: FirRegularClass)
+
+    fun interface Factory : FirExtension.Factory<FirExistingClassModificationExtension>
+}
+
+val FirExtensionService.existingClassModifiers: List<FirExistingClassModificationExtension> by FirExtensionService.registeredExtensions()
@@ -16,7 +16,7 @@ abstract class FirExtensionRegistrar {
    ) {
        val AVAILABLE_EXTENSIONS = listOf(
            FirStatusTransformerExtension::class,
-            FirClassGenerationExtension::class,
+            FirExistingClassModificationExtension::class,
            AbstractFirAdditionalCheckersExtension::class
        )
    }
@@ -30,8 +30,8 @@ abstract class FirExtensionRegistrar {
        }

        @JvmName("plusClassGenerationExtension")
-        operator fun ((FirSession) -> FirClassGenerationExtension).unaryPlus() {
-            registerExtension(FirClassGenerationExtension::class, FirClassGenerationExtension.Factory { this.invoke(it) })
+        operator fun ((FirSession) -> FirExistingClassModificationExtension).unaryPlus() {
+            registerExtension(FirExistingClassModificationExtension::class, FirExistingClassModificationExtension.Factory { this.invoke(it) })
        }

        @JvmName("plusAdditionalCheckersExtension")
@@ -25,6 +25,9 @@ abstract class FirPredicateBasedProvider : FirSessionComponent {
    }

    abstract fun getSymbolsByPredicate(predicate: DeclarationPredicate): List<FirAnnotatedDeclaration>
+    abstract fun getSymbolsWithOwnersByPredicate(
+        predicate: DeclarationPredicate
+    ): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>>

    abstract fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>)

@@ -43,6 +46,10 @@ private class FirPredicateBasedProviderImpl(private val session: FirSession) : F
        }
    }

+    override fun getSymbolsWithOwnersByPredicate(predicate: DeclarationPredicate): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>> {
+        return getSymbolsByPredicate(predicate).map { it to cache.ownersForDeclaration.getValue(it) }
+    }
+
    override fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>) {
        cache.ownersForDeclaration[declaration] = owners
        registerOwnersDeclarations(declaration, owners)
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.fir.resolve.transformers

+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase.*
@@ -12,9 +13,10 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitTypeBodyResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirGlobalExtensionStatusProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirPluginAnnotationsResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirTransformerBasedExtensionStatusProcessor
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.*
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose

 fun FirResolvePhase.createCompilerProcessorByPhase(
    session: FirSession,
@@ -23,7 +25,7 @@ fun FirResolvePhase.createCompilerProcessorByPhase(
    return when (this) {
        RAW_FIR -> throw IllegalStateException("Raw FIR building phase does not have a transformer")
        ANNOTATIONS_FOR_PLUGINS -> FirPluginAnnotationsResolveProcessor(session, scopeSession)
-//        FIRST_PLUGIN_GENERATION -> FirFirstGenerationTransformer()
+        NESTED_CLASS_GENERATION -> FirGlobalNestedClassGenerationProcessor(session, scopeSession)
        IMPORTS -> FirImportResolveProcessor(session, scopeSession)
        SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
        SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
@@ -31,15 +33,20 @@ fun FirResolvePhase.createCompilerProcessorByPhase(
        EXTENSION_STATUS_UPDATE -> FirGlobalExtensionStatusProcessor(session, scopeSession)
        STATUS -> FirStatusResolveProcessor(session, scopeSession)
        CONTRACTS -> FirContractResolveProcessor(session, scopeSession)
+        NEW_MEMBERS_GENERATION -> FirGlobalNewMemberGenerationProcessor(session, scopeSession)
        IMPLICIT_TYPES_BODY_RESOLVE -> FirImplicitTypeBodyResolveProcessor(session, scopeSession)
        BODY_RESOLVE -> FirBodyResolveProcessor(session, scopeSession)
    }
 }

-fun FirResolvePhase.createTransformerBasedProcessorByPhase(session: FirSession, scopeSession: ScopeSession): FirTransformerBasedResolveProcessor {
+fun FirResolvePhase.createTransformerBasedProcessorByPhase(
+    session: FirSession,
+    scopeSession: ScopeSession
+): FirTransformerBasedResolveProcessor {
    return when (this) {
        RAW_FIR -> throw IllegalStateException("Raw FIR building phase does not have a transformer")
        ANNOTATIONS_FOR_PLUGINS -> FirPluginAnnotationsResolveProcessor(session, scopeSession)
+        NESTED_CLASS_GENERATION -> FirDummyTransformerBasedProcessor(session, scopeSession) // TODO: remove
        IMPORTS -> FirImportResolveProcessor(session, scopeSession)
        SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
        SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
@@ -47,7 +54,22 @@ fun FirResolvePhase.createTransformerBasedProcessorByPhase(session: FirSession,
        EXTENSION_STATUS_UPDATE -> FirTransformerBasedExtensionStatusProcessor(session, scopeSession)
        STATUS -> FirStatusResolveProcessor(session, scopeSession)
        CONTRACTS -> FirContractResolveProcessor(session, scopeSession)
+        NEW_MEMBERS_GENERATION -> FirDummyTransformerBasedProcessor(session, scopeSession) // TODO: remove
        IMPLICIT_TYPES_BODY_RESOLVE -> FirImplicitTypeBodyResolveProcessor(session, scopeSession)
        BODY_RESOLVE -> FirBodyResolveProcessor(session, scopeSession)
    }
 }
+
+private class FirDummyTransformerBasedProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirTransformerBasedResolveProcessor(session, scopeSession) {
+    override val transformer: FirTransformer<Nothing?>
+        get() = DummyTransformer
+
+    private object DummyTransformer : FirTransformer<Nothing?>() {
+        override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+            return element.compose()
+        }
+    }
+}
@@ -0,0 +1,38 @@
+/*
+ * 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.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.addDeclaration
+import org.jetbrains.kotlin.fir.extensions.existingClassModifiers
+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.transformers.FirGlobalResolveProcessor
+
+class FirGlobalNestedClassGenerationProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirGlobalResolveProcessor(session, scopeSession) {
+    override fun process() {
+        val extensions = session.extensionService.existingClassModifiers
+        if (extensions.isEmpty()) return
+        val provider = session.predicateBasedProvider
+        val index = session.generatedNestedClassIndex
+        for (extension in extensions) {
+            val declarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
+            for ((declaration, owners) in declarations) {
+                val nestedClasses = extension.generateNestedClasses(declaration, owners)
+                for ((nestedClass, owner) in nestedClasses) {
+                    owner.addDeclaration(nestedClass)
+                    index.registerNestedClass(nestedClass, owner)
+                }
+            }
+        }
+    }
+}
+
@@ -0,0 +1,55 @@
+/*
+ * 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.addDeclaration
+import org.jetbrains.kotlin.fir.declarations.validate
+import org.jetbrains.kotlin.fir.extensions.FirExistingClassModificationExtension
+import org.jetbrains.kotlin.fir.extensions.existingClassModifiers
+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.transformers.FirGlobalResolveProcessor
+
+class FirGlobalNewMemberGenerationProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirGlobalResolveProcessor(session, scopeSession) {
+    private val index = session.generatedNestedClassIndex
+    private val provider = session.predicateBasedProvider
+
+    override fun process() {
+        val extensions = session.extensionService.existingClassModifiers
+        if (extensions.isEmpty()) return
+        for (extension in extensions) {
+            generateNewMembers(extension)
+            fillGeneratedNestedClasses(extension)
+        }
+    }
+
+    private fun generateNewMembers(extension: FirExistingClassModificationExtension) {
+        val declarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
+        for ((declaration, owners) in declarations) {
+            val newMembers = extension.generateMembers(declaration, owners)
+            for ((newMember, owner) in newMembers) {
+                newMember.validate()
+                owner.addDeclaration(newMember)
+            }
+        }
+    }
+
+    private fun fillGeneratedNestedClasses(extension: FirExistingClassModificationExtension) {
+        for (generatedNestedClass in index[extension.key]) {
+            val nestedClass = generatedNestedClass.nestedClass
+            val newMembers = extension.generateMembersForNestedClasses(generatedNestedClass)
+            for (newMember in newMembers) {
+                newMember.validate()
+                nestedClass.addDeclaration(newMember)
+            }
+        }
+    }
+}
@@ -0,0 +1,41 @@
+/*
+ * 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 com.google.common.collect.ArrayListMultimap
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+
+data class GeneratedNestedClass(val nestedClass: FirRegularClass, val owner: FirRegularClass)
+
+abstract class GeneratedNestedClassIndex : FirSessionComponent {
+    companion object {
+        fun create(): GeneratedNestedClassIndex {
+            return GeneratedNestedClassIndexImpl()
+        }
+    }
+
+    abstract fun registerNestedClass(klass: FirRegularClass, owner: FirRegularClass)
+    abstract operator fun get(key: FirPluginKey): List<GeneratedNestedClass>
+}
+
+val FirSession.generatedNestedClassIndex: GeneratedNestedClassIndex by FirSession.sessionComponentAccessor()
+
+private class GeneratedNestedClassIndexImpl : GeneratedNestedClassIndex() {
+    private val index: ArrayListMultimap<FirPluginKey, GeneratedNestedClass> = ArrayListMultimap.create()
+
+    override fun registerNestedClass(klass: FirRegularClass, owner: FirRegularClass) {
+        val key = (klass.origin as FirDeclarationOrigin.Plugin).key
+        index.put(key, GeneratedNestedClass(klass, owner))
+    }
+
+    override fun get(key: FirPluginKey): List<GeneratedNestedClass> {
+        return index.get(key)
+    }
+}
@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir.declarations
 enum class FirResolvePhase {
    RAW_FIR,
    ANNOTATIONS_FOR_PLUGINS, // run only if some extensions are registered
-//    FIRST_PLUGIN_GENERATION, // plugin phase
+    NESTED_CLASS_GENERATION, // plugin phase
    IMPORTS,
    SUPER_TYPES,
    SEALED_CLASS_INHERITORS,
@@ -16,6 +16,7 @@ enum class FirResolvePhase {
    EXTENSION_STATUS_UPDATE,
    STATUS,
    CONTRACTS,
+    NEW_MEMBERS_GENERATION, // plugin
    IMPLICIT_TYPES_BODY_RESOLVE,
    BODY_RESOLVE;