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

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")
+        }
+    }
+}