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

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirAllOpenComponentRegistrar.kt

@@ -21,5 +21,6 @@ class FirAllOpenComponentRegistrar : FirExtensionRegistrar() {
         +::TopLevelDeclarationsGenerator
         +::ExternalClassGenerator
         +::AdditionalMembersGenerator
+        +::CompanionGenerator
     }
 }

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/CompanionGenerator.kt

@@ -0,0 +1,128 @@
+/*
+ * 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.plugin.generators
+
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.descriptors.EffectiveVisibility
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
+import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.fir.declarations.utils.isCompanion
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.predicate.has
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.plugin.fqn
+import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.SpecialNames
+
+/*
+ * Generates companion object with fun foo(): Int for each class annotated with @D
+ */
+class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
+    companion object {
+        private val PREDICATE = has("D".fqn())
+        private val FOO_NAME = Name.identifier("foo")
+    }
+
+    private val matchedClasses by lazy {
+        session.predicateBasedProvider.getSymbolsByPredicate(PREDICATE)
+            .filterIsInstance<FirRegularClassSymbol>()
+    }
+
+    override fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
+        if (classId.shortClassName != SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT) return null
+        val owner = matchedClasses.firstOrNull { it.classId == classId.outerClassId } ?: return null
+        if (owner.companionObjectSymbol != null) return null
+        val regularClass = buildRegularClass {
+            moduleData = session.moduleData
+            origin = Key.origin
+            classKind = ClassKind.OBJECT
+            scopeProvider = session.kotlinScopeProvider
+            status = FirResolvedDeclarationStatusImpl(
+                Visibilities.Public,
+                Modality.FINAL,
+                EffectiveVisibility.Public
+            )
+            name = SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
+            symbol = FirRegularClassSymbol(classId)
+            superTypeRefs += session.builtinTypes.anyType
+        }
+        return regularClass.symbol
+    }
+
+    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+        if (owner == null || owner.origin != Key.origin) return emptyList()
+        if (callableId.callableName != FOO_NAME) return emptyList()
+        val function = buildSimpleFunction {
+            moduleData = session.moduleData
+            origin = Key.origin
+            status = FirResolvedDeclarationStatusImpl(
+                Visibilities.Public,
+                Modality.FINAL,
+                EffectiveVisibility.Public
+            )
+            name = FOO_NAME
+            symbol = FirNamedFunctionSymbol(callableId)
+            returnTypeRef = session.builtinTypes.intType
+            dispatchReceiverType = owner.defaultType()
+        }
+        return listOf(function.symbol)
+    }
+
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+        return setOf(FOO_NAME)
+    }
+
+    @OptIn(SymbolInternals::class)
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
+        return if (session.predicateBasedProvider.matches(PREDICATE, classSymbol.fir)) {
+            setOf(SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT)
+        } else {
+            emptySet()
+        }
+    }
+
+    override fun needToGenerateAdditionalMembersInClass(klass: FirClass): Boolean {
+        if (klass !is FirRegularClass) return false
+        if (matchedClasses.none { it.classId == klass.classId.outerClassId }) return false
+        return klass.isCompanion
+    }
+
+    override fun needToGenerateNestedClassifiersInClass(klass: FirClass): Boolean {
+        return session.predicateBasedProvider.matches(PREDICATE, klass)
+    }
+
+    override val key: FirPluginKey
+        get() = Key
+
+    object Key : FirPluginKey() {
+        override fun toString(): String {
+            return "CompanionGeneratorKey"
+        }
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
+}

plugins/fir/fir-plugin-prototype/testData/diagnostics/memberGen/classWithCompanionObject.fir.txt

@@ -0,0 +1,18 @@
+FILE: classWithCompanionObject.kt
+    @R|org/jetbrains/kotlin/fir/plugin/D|() public final class SomeClass : R|kotlin/Any|, R|foo/MyInterface| {
+        public constructor(): R|SomeClass| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final object Companion : R|kotlin/Any| {
+            public final fun foo(): R|kotlin/Int|
+
+        }
+
+    }
+    public final fun takeInt(x: R|kotlin/Int|): R|kotlin/Unit| {
+    }
+    public final fun test(): R|kotlin/Unit| {
+        R|/takeInt|(Q|SomeClass|.R|/SomeClass.Companion.foo|())
+        R|/takeInt|(Q|SomeClass.Companion|.R|/SomeClass.Companion.foo|())
+    }

plugins/fir/fir-plugin-prototype/testData/diagnostics/memberGen/classWithCompanionObject.kt

@@ -0,0 +1,13 @@
+import org.jetbrains.kotlin.fir.plugin.D
+
+@D
+class SomeClass
+
+fun takeInt(x: Int) {}
+
+fun test() {
+    takeInt(SomeClass.foo())
+    takeInt(SomeClass.Companion.foo())
+}
+
+

plugins/fir/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirAllOpenDiagnosticTestGenerated.java

@@ -55,6 +55,12 @@ public class FirAllOpenDiagnosticTestGenerated extends AbstractFirAllOpenDiagnos
             KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("plugins/fir/fir-plugin-prototype/testData/diagnostics/memberGen"), Pattern.compile("^(.+)\\.kt$"), null, true);
         }
 
+        @Test
+        @TestMetadata("classWithCompanionObject.kt")
+        public void testClassWithCompanionObject() throws Exception {
+            runTest("plugins/fir/fir-plugin-prototype/testData/diagnostics/memberGen/classWithCompanionObject.kt");
+        }
+
         @Test
         @TestMetadata("classWithGeneratedMembersAndNestedClass.kt")
         public void testClassWithGeneratedMembersAndNestedClass() throws Exception {