[FIR] Add prototype of first generation phase

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

@@ -0,0 +1,43 @@
+/*
+ * 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
+
+/*
+ * TODO:
+ *  - add registration of new files
+ *  - check that annotations or meta-annotations is not empty
+ */
+abstract class FirClassGenerationExtension(session: FirSession) : FirExtensionPoint(session) {
+    companion object {
+        val NAME = FirExtensionPointName("StatusTransformer")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val mode: Mode = Mode.ANNOTATED_ELEMENT
+
+    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 : FirExtensionPoint.Factory<FirClassGenerationExtension>
+}
+
+val FirExtensionPointService.classGenerationExtensions: FirExtensionPointService.ExtensionsAccessor<FirClassGenerationExtension> by FirExtensionPointService.registeredExtensions()

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

@@ -15,9 +15,15 @@ abstract class FirExtensionRegistrar {
     )
 
     protected inner class ExtensionRegistrarContext {
+        @JvmName("plusStatusTransformerExtension")
         operator fun ((FirSession) -> FirStatusTransformerExtension).unaryPlus() {
             statusTransformerExtensions += FirStatusTransformerExtension.Factory { this.invoke(it) }
         }
+
+        @JvmName("plusClassGenerationExtension")
+        operator fun ((FirSession) -> FirClassGenerationExtension).unaryPlus() {
+            classGenerationExtensions += FirClassGenerationExtension.Factory { this.invoke(it) }
+        }
     }
 
     protected abstract fun ExtensionRegistrarContext.configurePlugin()
@@ -25,21 +31,25 @@ abstract class FirExtensionRegistrar {
     fun configure(): RegisteredExtensions {
         ExtensionRegistrarContext().configurePlugin()
         return RegisteredExtensions(
-            statusTransformerExtensions
+            statusTransformerExtensions,
+            classGenerationExtensions
         )
     }
 
     private val statusTransformerExtensions: MutableList<FirStatusTransformerExtension.Factory> = mutableListOf()
+    private val classGenerationExtensions: MutableList<FirClassGenerationExtension.Factory> = mutableListOf()
 
     class RegisteredExtensions(
-        val statusTransformerExtensions: List<FirStatusTransformerExtension.Factory>
+        val statusTransformerExtensions: List<FirStatusTransformerExtension.Factory>,
+        val classGenerationExtensions: List<FirClassGenerationExtension.Factory>
     ) {
         companion object {
-            val EMPTY = RegisteredExtensions(emptyList())
+            val EMPTY = RegisteredExtensions(emptyList(), emptyList())
         }
     }
 }
 
 fun FirExtensionPointService.registerExtensions(extensions: FirExtensionRegistrar.RegisteredExtensions) {
     registerExtensions(FirStatusTransformerExtension::class, extensions.statusTransformerExtensions)
+    registerExtensions(FirClassGenerationExtension::class, extensions.classGenerationExtensions)
 }

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformerAdapter
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitTypeBodyResolveTransformerAdapter
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveTransformerAdapter
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirFirstGenerationTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirPluginAnnotationsResolveTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 
@@ -20,6 +21,7 @@ fun FirResolvePhase.createTransformerByPhase(scopeSession: ScopeSession): FirTra
     return when (this) {
         RAW_FIR -> throw AssertionError("Raw FIR building phase does not have a transformer")
         ANNOTATIONS_FOR_PLUGINS -> FirPluginAnnotationsResolveTransformer(scopeSession)
+        FIRST_PLUGIN_GENERATION -> FirFirstGenerationTransformer()
         IMPORTS -> FirImportResolveTransformer()
         SUPER_TYPES -> FirSupertypeResolverTransformer(scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsTransformer()

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

@@ -0,0 +1,104 @@
+/*
+ * 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.FirAnnotationContainer
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.extensions.FirClassGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.classGenerationExtensions
+import org.jetbrains.kotlin.fir.extensions.extensionPointService
+import org.jetbrains.kotlin.fir.extensions.hasExtensions
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
+
+class FirFirstGenerationTransformer : FirDefaultTransformer<Nothing?>() {
+    private lateinit var session: FirSession
+    private val generatedDeclarations: MutableList<FirClassGenerationExtension.GeneratedClass> = mutableListOf()
+    private lateinit var file: FirFile
+    private val extensionPointService get() = session.extensionPointService
+
+    override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        session = file.session
+        if (!extensionPointService.hasExtensions) return file.compose()
+        file.replaceResolvePhase(FirResolvePhase.FIRST_PLUGIN_GENERATION)
+        generatedDeclarations.clear()
+        this.file = file
+
+        file.transformDeclarations(this, data)
+        visitAnnotatedDeclaration(file)
+        addGeneratedClasses()
+        return file.compose()
+    }
+
+    private fun addGeneratedClasses() {
+        for ((generatedClass, owner) in generatedDeclarations) {
+            when (owner) {
+                is FirFile -> owner.addDeclaration(generatedClass)
+                is FirRegularClass -> owner.addDeclaration(generatedClass)
+                else -> throw IllegalStateException()
+            }
+        }
+    }
+
+    private fun <T> visitAnnotatedDeclaration(declaration: T) where T : FirDeclaration, T : FirAnnotationContainer {
+        val extensions = extensionPointService.classGenerationExtensions.forDeclaration(declaration, emptyList())
+        extensions.flatMapTo(generatedDeclarations) { it.generateClass(file, declaration) }
+    }
+
+    override fun transformProperty(property: FirProperty, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        visitAnnotatedDeclaration(property)
+        return property.compose()
+    }
+
+    override fun transformRegularClass(
+        regularClass: FirRegularClass,
+        data: Nothing?
+    ): CompositeTransformResult<FirStatement> {
+        regularClass.transformDeclarations(this, data)
+        regularClass.transformCompanionObject(this, data)
+        visitAnnotatedDeclaration(regularClass)
+        return regularClass.compose()
+    }
+
+    override fun transformSimpleFunction(
+        simpleFunction: FirSimpleFunction,
+        data: Nothing?
+    ): CompositeTransformResult<FirDeclaration> {
+        simpleFunction.transformValueParameters(this, data)
+        visitAnnotatedDeclaration(simpleFunction)
+        return simpleFunction.compose()
+    }
+
+    override fun transformConstructor(
+        constructor: FirConstructor,
+        data: Nothing?
+    ): CompositeTransformResult<FirDeclaration> {
+        constructor.transformValueParameters(this, data)
+        visitAnnotatedDeclaration(constructor)
+        return constructor.compose()
+    }
+
+    override fun transformValueParameter(
+        valueParameter: FirValueParameter,
+        data: Nothing?
+    ): CompositeTransformResult<FirStatement> {
+        visitAnnotatedDeclaration(valueParameter)
+        return valueParameter.compose()
+    }
+
+    override fun transformTypeAlias(typeAlias: FirTypeAlias, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        visitAnnotatedDeclaration(typeAlias)
+        return typeAlias.compose()
+    }
+
+    override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+        return element.compose()
+    }
+}

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

@@ -9,6 +9,9 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.declarations.builder.AbstractFirRegularClassBuilder
 import org.jetbrains.kotlin.fir.declarations.builder.FirTypeParameterBuilder
+import org.jetbrains.kotlin.fir.declarations.impl.FirClassImpl
+import org.jetbrains.kotlin.fir.declarations.impl.FirFileImpl
+import org.jetbrains.kotlin.fir.declarations.impl.FirSealedClassImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousObjectSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
@@ -16,6 +19,7 @@ import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeFlexibleType
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import java.lang.IllegalStateException
 
 fun FirTypeParameterBuilder.addDefaultBoundIfNecessary(isFlexible: Boolean = false) {
     if (bounds.isEmpty()) {
@@ -89,4 +93,18 @@ val FirClass<*>.superConeTypes get() = superTypeRefs.mapNotNull { it.coneTypeSaf
 fun FirRegularClass.collectEnumEntries(): Collection<FirEnumEntry> {
     assert(classKind == ClassKind.ENUM_CLASS)
     return declarations.filterIsInstance<FirEnumEntry>()
+}
+
+fun FirFile.addDeclaration(declaration: FirDeclaration) {
+    require(this is FirFileImpl)
+    declarations += declaration
+}
+
+fun FirRegularClass.addDeclaration(declaration: FirDeclaration) {
+    @Suppress("LiftReturnOrAssignment")
+    when (this) {
+        is FirClassImpl -> declarations += declaration
+        is FirSealedClassImpl -> declarations += declaration
+        else -> throw IllegalStateException()
+    }
 }

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

@@ -8,6 +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
     IMPORTS,
     SUPER_TYPES,
     SEALED_CLASS_INHERITORS,