[FIR] Rewrite algorithm of generating declarations by plugins

Before there was separate global stage which asks extensions
  to generate new declarations. Now it is changed to provider
  like requests, so extensions will decide do they have declarations
  with specific classId/callableId or not

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

@@ -15,10 +15,10 @@ class FirAllOpenComponentRegistrar : FirExtensionRegistrar() {
         +::AllOpenSupertypeGenerator
 
         // Declaration generators
-        +::AllOpenMemberGenerator
-        +::AllOpenNestedClassGenerator
+//        +::AllOpenMemberGenerator
+//        +::AllOpenNestedClassGenerator
         +::AllOpenAdditionalCheckers
         +::AllOpenTopLevelDeclarationsGenerator
-        +::AllOpenRecursiveNestedClassGenerator
+//        +::AllOpenRecursiveNestedClassGenerator
     }
-}
+}

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

@@ -1,60 +0,0 @@
-/*
- * 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.plugin.generators
-
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
-import org.jetbrains.kotlin.fir.declarations.utils.classId
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
-import org.jetbrains.kotlin.fir.extensions.predicate.and
-import org.jetbrains.kotlin.fir.extensions.predicate.has
-import org.jetbrains.kotlin.fir.extensions.predicate.under
-import org.jetbrains.kotlin.fir.moduleData
-import org.jetbrains.kotlin.fir.plugin.AllOpenPluginKey
-import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.Name
-
-class AllOpenMemberGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        return emptyList()
-    }
-
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        return emptyList()
-    }
-
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        if (annotatedDeclaration !is FirProperty) return emptyList()
-        val owner = owners.last() as? FirRegularClass ?: return emptyList()
-        val propertyName = annotatedDeclaration.name.identifier
-        val function = buildSimpleFunction {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = annotatedDeclaration.returnTypeRef
-            status = annotatedDeclaration.status
-            name = Name.identifier("hello${propertyName.replaceFirstChar(Char::uppercaseChar)}")
-            symbol = FirNamedFunctionSymbol(CallableId(owner.classId, name))
-        }
-        return listOf(GeneratedDeclaration(function, owner))
-    }
-
-    override val predicate: DeclarationPredicate = under("WithGenerated".fqn()) and has("WithHello".fqn())
-
-    override val key: FirPluginKey
-        get() = AllOpenPluginKey
-}

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

@@ -1,105 +0,0 @@
-/*
- * 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.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.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildConstructor
-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.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
-import org.jetbrains.kotlin.fir.extensions.predicate.has
-import org.jetbrains.kotlin.fir.moduleData
-import org.jetbrains.kotlin.fir.plugin.AllOpenPluginKey
-import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
-import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.Name
-
-class AllOpenNestedClassGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        val owner = annotatedDeclaration as? FirRegularClass ?: return emptyList()
-        val newClass = buildRegularClass {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Private,
-                Modality.FINAL,
-                EffectiveVisibility.PrivateInClass
-            ).apply {
-                isInner = true
-            }
-            classKind = ClassKind.CLASS
-            name = Name.identifier("Foo")
-            symbol = FirRegularClassSymbol(owner.symbol.classId.createNestedClassId(name))
-            scopeProvider = owner.scopeProvider
-        }
-        return listOf(GeneratedDeclaration(newClass, owner))
-    }
-
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        val klass = generatedClass.klass
-
-        val classId = klass.symbol.classId
-        val constructor = buildConstructor {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = buildResolvedTypeRef {
-                type = ConeClassLikeTypeImpl(ConeClassLikeLookupTagImpl(classId), emptyArray(), isNullable = false)
-            }
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            symbol = FirConstructorSymbol(CallableId(classId, classId.shortClassName))
-        }
-
-        val function = buildSimpleFunction {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = session.builtinTypes.intType
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            name = Name.identifier("hello")
-            symbol = FirNamedFunctionSymbol(CallableId(classId, name))
-        }
-        return listOf(constructor, function)
-    }
-
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        return emptyList()
-    }
-
-    override val predicate: DeclarationPredicate = has("WithNestedFoo".fqn())
-
-    override val key: FirPluginKey
-        get() = AllOpenPluginKey
-}

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

@@ -1,87 +0,0 @@
-/*
- * 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.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.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
-import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
-import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
-import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyAnnotationArgumentMapping
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
-import org.jetbrains.kotlin.fir.extensions.predicate.has
-import org.jetbrains.kotlin.fir.moduleData
-import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.Name
-
-class AllOpenRecursiveNestedClassGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        if (owners.size > 2) return emptyList()
-        val owner = annotatedDeclaration as? FirRegularClass ?: return emptyList()
-        val newClass = buildRegularClass {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            classKind = ClassKind.CLASS
-            name = Name.identifier("Nested")
-            symbol = FirRegularClassSymbol(owner.symbol.classId.createNestedClassId(name))
-            scopeProvider = owner.scopeProvider
-            annotations += buildAnnotation {
-                annotationTypeRef = buildResolvedTypeRef {
-                    type = ConeClassLikeTypeImpl(
-                        ConeClassLikeLookupTagImpl(annotationClassId),
-                        emptyArray(),
-                        isNullable = false
-                    )
-                }
-                argumentMapping = FirEmptyAnnotationArgumentMapping
-            }
-        }
-        return listOf(GeneratedDeclaration(newClass, owner))
-    }
-
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        return emptyList()
-    }
-
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        return emptyList()
-    }
-
-    override val key: FirPluginKey
-        get() = Key
-
-    override val predicate: DeclarationPredicate
-        get() = has("B".fqn())
-
-    private val annotationClassId = "B".fqn().let {
-        ClassId(it.parent(), it.shortName())
-    }
-
-    private object Key : FirPluginKey()
-}

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

@@ -5,82 +5,74 @@
 
 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.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 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.packageFqName
+import org.jetbrains.kotlin.fir.plugin.AllOpenPluginKey
 import org.jetbrains.kotlin.fir.plugin.fqn
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 class AllOpenTopLevelDeclarationsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
-    override fun generateClasses(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<FirRegularClass>> {
-        val file = owners.first() as FirFile
-        val klass = annotatedDeclaration as? FirRegularClass ?: return emptyList()
-        val newClass = buildRegularClass {
-            moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            classKind = ClassKind.OBJECT
-            name = Name.identifier("TopLevel${klass.name}")
-            symbol = FirRegularClassSymbol(ClassId(file.packageFqName, name))
-            scopeProvider = klass.scopeProvider
-        }
-        return listOf(GeneratedDeclaration(newClass, file))
+    private val predicateBasedProvider = session.predicateBasedProvider
+    private val matchedClasses by lazy {
+        predicateBasedProvider.getSymbolsByPredicate(predicate).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
     }
 
-    override fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration> {
-        val klass = generatedClass.klass
+    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+        if (owner != null) return emptyList()
+        val matchedClassSymbol = findMatchedClassForFunction(callableId) ?: return emptyList()
         val function = buildSimpleFunction {
             moduleData = session.moduleData
-            resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
-            origin = FirDeclarationOrigin.Plugin(key)
-            returnTypeRef = session.builtinTypes.intType
+            origin = key.origin
             status = FirResolvedDeclarationStatusImpl(
                 Visibilities.Public,
                 Modality.FINAL,
                 EffectiveVisibility.Public
             )
-            name = Name.identifier("hello")
-            symbol = FirNamedFunctionSymbol(CallableId(klass.symbol.classId, name))
+            returnTypeRef = session.builtinTypes.stringType
+            valueParameters += buildValueParameter {
+                moduleData = session.moduleData
+                origin = key.origin
+                returnTypeRef = buildResolvedTypeRef {
+                    type = ConeClassLikeTypeImpl(ConeClassLikeLookupTagImpl(matchedClassSymbol.classId), emptyArray(), isNullable = false)
+                }
+                name = Name.identifier("value")
+                symbol = FirValueParameterSymbol(name)
+                isCrossinline = false
+                isNoinline = false
+                isVararg = false
+            }
+            symbol = FirNamedFunctionSymbol(callableId)
+            name = callableId.callableName
         }
-        return listOf(function)
+        return listOf(function.symbol)
     }
 
-    override fun generateMembers(
-        annotatedDeclaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>
-    ): List<GeneratedDeclaration<*>> {
-        return emptyList()
+    private fun findMatchedClassForFunction(callableId: CallableId): FirRegularClassSymbol? {
+        // We generate only top-level functions
+        if (callableId.classId != null) return null
+        return matchedClasses
+            .filter { it.classId.packageFqName == callableId.packageName }
+            .firstOrNull { callableId.callableName.identifier == "dummy${it.classId.shortClassName.identifier}" }
     }
 
-    override val key: FirPluginKey
-        get() = Key
+    override val key: AllOpenPluginKey
+        get() = AllOpenPluginKey
 
     override val predicate: DeclarationPredicate
         get() = has("A".fqn())
 
-    private object Key : FirPluginKey()
 }

plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelCallables.fir.txt

@@ -0,0 +1,14 @@
+FILE: topLevelCallables.kt
+    @R|org/jetbrains/kotlin/fir/plugin/A|() public final class MySuperClass : R|kotlin/Any| {
+        public constructor(): R|foo/MySuperClass| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun test(): R|kotlin/Unit| {
+            lval s: R|kotlin/String| = R|foo/dummyMySuperClass|(this@R|foo/MySuperClass|)
+            R|foo/takeString|(R|<local>/s|)
+        }
+
+    }
+    public final fun takeString(s: R|kotlin/String|): R|kotlin/Unit| {
+    }

plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelCallables.kt

@@ -0,0 +1,17 @@
+package foo
+
+import org.jetbrains.kotlin.fir.plugin.A
+
+/*
+ * Plugin generates `dummyClassName(value: ClassName): String` function for each class annotated with @A
+ */
+
+@A
+class MySuperClass {
+    fun test() {
+        val s = dummyMySuperClass(this)
+        takeString(s)
+    }
+}
+
+fun takeString(s: String) {}

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

@@ -67,6 +67,12 @@ public class FirAllOpenDiagnosticTestGenerated extends AbstractFirAllOpenDiagnos
             runTest("plugins/fir/fir-plugin-prototype/testData/memberGen/recursiveNestedClasses.kt");
         }
 
+        @Test
+        @TestMetadata("topLevelCallables.kt")
+        public void testTopLevelCallables() throws Exception {
+            runTest("plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelCallables.kt");
+        }
+
         @Test
         @TestMetadata("topLevelClass.kt")
         public void testTopLevelClass() throws Exception {