[FIR] Add ability to generate companion object from compiler plugin

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

@@ -167,6 +167,7 @@ object FirSessionFactory {
             }.configure()
 
             val dependenciesSymbolProvider = FirDependenciesSymbolProviderImpl(this)
+            val generatedSymbolsProvider = FirExtensionDeclarationsSymbolProvider.create(this)
             register(
                 FirSymbolProvider::class,
                 FirCompositeSymbolProvider(
@@ -174,13 +175,15 @@ object FirSessionFactory {
                     listOfNotNull(
                         firProvider.symbolProvider,
                         symbolProviderForBinariesFromIncrementalCompilation,
-                        FirExtensionDeclarationsSymbolProvider.create(this),
+                        generatedSymbolsProvider,
                         JavaSymbolProvider(this, projectEnvironment.getFirJavaFacade(this, moduleData, scope)),
                         dependenciesSymbolProvider,
                     )
                 )
             )
 
+            generatedSymbolsProvider?.let { register(FirExtensionDeclarationsSymbolProvider::class, it) }
+
             register(
                 FirDependenciesSymbolProvider::class,
                 dependenciesSymbolProvider

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

@@ -31,7 +31,11 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredi
     abstract fun needToGenerateAdditionalMembersInClass(klass: FirClass): Boolean
     abstract fun needToGenerateNestedClassifiersInClass(klass: FirClass): Boolean
 
-    // Can be called on SUPERTYPES stage
+    /*
+     * Can be called on SUPERTYPES stage
+     *
+     * If classId has `outerClassId.Companion` format then generated class should be a companion object
+     */
     open fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
 
     // Can be called on STATUS stage

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirGeneratedClassDeclaredMemberScope.kt

@@ -10,10 +10,8 @@ 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.caches.getValue
-import org.jetbrains.kotlin.fir.declarations.FirClass
-import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.classId
-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
@@ -133,6 +131,13 @@ class FirGeneratedClassNestedClassifierScope(
         }
 
     private fun generateNestedClassifier(name: Name): FirRegularClassSymbol? {
+        if (klass is FirRegularClass) {
+            val companion = klass.companionObjectSymbol
+            if (companion != null && companion.origin.generated && companion.classId.shortClassName == name) {
+                return companion
+            }
+        }
+
         if (name !in getClassifierNames()) return null
         val generatedClass = useSiteSession.symbolProvider.getClassLikeSymbolByClassId(klass.classId.createNestedClassId(name))
         require(generatedClass is FirRegularClassSymbol?) { "Only regular class are allowed as nested classes" }

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
@@ -14,6 +15,7 @@ import org.jetbrains.kotlin.fir.declarations.validate
 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.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
@@ -24,7 +26,7 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
     session: FirSession,
     cachesFactory: FirCachesFactory,
     private val extensions: List<FirDeclarationGenerationExtension>
-) : FirSymbolProvider(session) {
+) : FirSymbolProvider(session), FirSessionComponent {
     companion object {
         fun create(session: FirSession): FirExtensionDeclarationsSymbolProvider? {
             val extensions = session.extensionService.declarationGenerators
@@ -101,3 +103,5 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
         return fqName.takeIf { packageCache.getValue(fqName, null) }
     }
 }
+
+val FirSession.generatedDeclarationsSymbolProvider: FirExtensionDeclarationsSymbolProvider? by FirSession.nullableSessionComponentAccessor()

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

@@ -66,6 +66,7 @@ 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)
+        COMPANION_GENERATION -> FirCompanionGenerationProcessor(session, scopeSession)
         IMPORTS -> FirImportResolveProcessor(session, scopeSession)
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)

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

@@ -0,0 +1,54 @@
+/*
+ * 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.resolve.transformers.plugin
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.extensions.generatedDeclarationsSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.name.SpecialNames
+
+class FirCompanionGenerationProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirTransformerBasedResolveProcessor(session, scopeSession) {
+    override val transformer: FirTransformer<Nothing?> = FirCompanionGenerationTransformer(session)
+}
+
+private class FirCompanionGenerationTransformer(val session: FirSession) : FirTransformer<Nothing?>() {
+    override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
+        return element
+    }
+
+    override fun transformFile(file: FirFile, data: Nothing?): FirFile {
+        return file.transformDeclarations(this, data)
+    }
+
+    override fun transformRegularClass(regularClass: FirRegularClass, data: Nothing?): FirStatement {
+        generateCompanion(regularClass)
+        return regularClass.transformDeclarations(this, data)
+    }
+
+    private fun generateCompanion(regularClass: FirRegularClass) {
+        // TODO: add proper error reporting
+        val generatedDeclarationProvider = session.generatedDeclarationsSymbolProvider ?: return
+        val companionClassId = regularClass.classId.createNestedClassId(SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT)
+        when (val generatedCompanion = generatedDeclarationProvider.getClassLikeSymbolByClassId(companionClassId)) {
+            null -> {}
+            is FirRegularClassSymbol -> when {
+                regularClass.companionObjectSymbol != null -> error("Plugin generated companion object for class $regularClass, but it is already present in class")
+                else -> regularClass.replaceCompanionObjectSymbol(generatedCompanion)
+            }
+            else -> error("Plugin generated non regular class as companion object")
+        }
+    }
+}

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

@@ -8,6 +8,7 @@ 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),
+    COMPANION_GENERATION(pluginPhase = true),
     IMPORTS,
     SUPER_TYPES,
     SEALED_CLASS_INHERITORS,