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

2021-10-19 15:26:33 +03:00
parent 75ed11a620
commit 9948ba57c0
12 changed files with 244 additions and 6 deletions
@@ -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
@@ -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
@@ -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" }
@@ -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()
@@ -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)
@@ -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")
+        }
+    }
+}
@@ -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,
@@ -21,5 +21,6 @@ class FirAllOpenComponentRegistrar : FirExtensionRegistrar() {
        +::TopLevelDeclarationsGenerator
        +::ExternalClassGenerator
        +::AdditionalMembersGenerator
+        +::CompanionGenerator
    }
 }
@@ -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)
+    }
+}
@@ -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|())
+    }
@@ -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())
+}
+
+
@@ -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 {