[FIR] Get rid of FirExtension.key

In reality plugin key was used only in two ways:
- found corresponding generator for member scope of generated class
- pass it to declaration origin so it can be read in backend IR plugin

Because of first usage plugin key was required to be an object, which
  reduces functionality of second usage, because there was no way to
  pass any additional info via key to IR (in future this usecase will be
  covered with IR declaration attributes, but right now we don't have
  them)

So now generator which created some class is now saved in its attributes
  and plugins can use any keys they want. And since there is no other
  usages of `FirExtension.key` it removed from FirExtension API at all

compiler/fir/providers/src/org/jetbrains/kotlin/fir/FirGeneratedDeclarationsUtils.kt

@@ -0,0 +1,15 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+
+private object OwnerForGeneratedDeclarationKey : FirDeclarationDataKey()
+
+internal var FirClassLikeDeclaration.ownerGenerator: FirDeclarationGenerationExtension? by FirDeclarationDataRegistry.data(OwnerForGeneratedDeclarationKey)

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.caches.*
 import org.jetbrains.kotlin.fir.declarations.validate
+import org.jetbrains.kotlin.fir.ownerGenerator
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
@@ -82,7 +83,11 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
             else -> extensionsByTopLevelClassId.getValue()[classId]
         } ?: return null
         val generatedClasses = matchedExtensions
-            .mapNotNull { it.generateClassLikeDeclaration(classId) }
+            .mapNotNull { generatorExtension ->
+                generatorExtension.generateClassLikeDeclaration(classId)?.also { symbol ->
+                    symbol.fir.ownerGenerator = generatorExtension
+                }
+            }
             .onEach { it.fir.validate() }
         return when (generatedClasses.size) {
             0 -> null

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

@@ -11,11 +11,11 @@ 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.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.declarations.validate
 import org.jetbrains.kotlin.fir.extensions.*
+import org.jetbrains.kotlin.fir.ownerGenerator
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
@@ -144,8 +144,7 @@ internal inline fun <T, V> FirSession.groupExtensionsByName(
 internal fun FirSession.getExtensionsForClass(klass: FirClass): List<FirDeclarationGenerationExtension> {
     val extensions = extensionService.declarationGenerators
     return if (klass.origin.generated) {
-        val pluginKey = (klass.origin as FirDeclarationOrigin.Plugin).key
+        listOf(klass.ownerGenerator!!)
-        extensions.filter { it.key == pluginKey }
     } else {
         extensions
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirExtension.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -24,8 +23,6 @@ typealias AnnotationFqn = FqName
 abstract class FirExtension(val session: FirSession) {
     abstract val name: FirExtensionPointName
 
-    abstract val key: FirPluginKey
-
     abstract val extensionType: KClass<out FirExtension>
 
     fun interface Factory<out P : FirExtension> {

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

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
-import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
 import org.jetbrains.kotlin.fir.extensions.FirStatusTransformerExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
@@ -38,7 +37,4 @@ class AllOpenStatusTransformer(session: FirSession) : FirStatusTransformerExtens
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {
         register(PREDICATE)
     }
-
-    override val key: FirPluginKey
-        get() = SomePluginKey
 }

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

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
-import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.expressions.arguments
@@ -78,7 +77,4 @@ class AllPublicVisibilityTransformer(session: FirSession) : FirStatusTransformer
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {
         register(PREDICATE)
     }
-
-    override val key: FirPluginKey
-        get() = SomePluginKey
 }

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

@@ -11,14 +11,10 @@ import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSimpleFunctionC
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
-import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.plugin.checkers.DummyNameChecker
 import org.jetbrains.kotlin.fir.plugin.checkers.SignedNumberCallChecker
 
 class PluginAdditionalCheckers(session: FirSession) : FirAdditionalCheckersExtension(session) {
-    override val key: FirPluginKey
-        get() = SomePluginKey
-
     override val declarationCheckers: DeclarationCheckers = object : DeclarationCheckers() {
         override val simpleFunctionCheckers: Set<FirSimpleFunctionChecker>
             get() = setOf(DummyNameChecker)

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.plugin
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
 import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
@@ -43,9 +42,6 @@ class SomeAdditionalSupertypeGenerator(session: FirSession) : FirSupertypeGenera
         )
     }
 
-    override val key: FirPluginKey
-        get() = SomePluginKey
-
     override fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean {
         return session.predicateBasedProvider.matches(PREDICATE, declaration)
     }

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

@@ -53,7 +53,7 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         if (callableId.callableName != MATERIALIZE_NAME) return emptyList()
         val classId = callableId.classId ?: return emptyList()
         val matchedClassSymbol = matchedClasses.firstOrNull { it.classId == classId } ?: return emptyList()
-        return listOf(buildMaterializeFunction(matchedClassSymbol, callableId).symbol)
+        return listOf(buildMaterializeFunction(matchedClassSymbol, callableId, Key).symbol)
     }
 
     override fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
@@ -62,7 +62,7 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         if (matchedClasses.none { it.classId == parentClassId }) return null
         return buildRegularClass {
             moduleData = session.moduleData
-            origin = key.origin
+            origin = Key.origin
             classKind = ClassKind.CLASS
             scopeProvider = session.kotlinScopeProvider
             status = FirResolvedDeclarationStatusImpl(Visibilities.Public, Modality.FINAL, EffectiveVisibility.Public)
@@ -74,7 +74,7 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
 
     override fun generateConstructors(owner: FirClassSymbol<*>): List<FirConstructorSymbol> {
         assert(owner.classId in nestedClassIds)
-        return listOf(buildConstructor(owner.classId, isInner = false).symbol)
+        return listOf(buildConstructor(owner.classId, isInner = false, Key).symbol)
     }
 
     override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
@@ -95,9 +95,6 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         }
     }
 
-    override val key: FirPluginKey
-        get() = Key
-
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {
         register(PREDICATE)
     }

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

@@ -99,9 +99,6 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
         }
     }
 
-    override val key: FirPluginKey
-        get() = Key
-
     object Key : FirPluginKey() {
         override fun toString(): String {
             return "CompanionGeneratorKey"

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

@@ -83,7 +83,7 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
     override fun generateConstructors(owner: FirClassSymbol<*>): List<FirConstructorSymbol> {
         val classId = owner.classId
         if (classId != GENERATED_CLASS_ID && classId !in classIdsForMatchedClasses) return emptyList()
-        return listOf(buildConstructor(classId, isInner = false).symbol)
+        return listOf(buildConstructor(classId, isInner = false, Key).symbol)
     }
 
     private fun generateNestedClass(classId: ClassId, owner: FirClassSymbol<*>): FirClassLikeSymbol<*>? {
@@ -101,13 +101,13 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
         require(owner is FirRegularClassSymbol)
         val matchedClassId = owner.fir.matchedClass ?: return emptyList()
         val matchedClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(matchedClassId) ?: return emptyList()
-        return listOf(buildMaterializeFunction(matchedClassSymbol, callableId).symbol)
+        return listOf(buildMaterializeFunction(matchedClassSymbol, callableId, Key).symbol)
     }
 
     private fun buildClass(classId: ClassId): FirRegularClass {
         return buildRegularClass {
             moduleData = session.moduleData
-            origin = key.origin
+            origin = Key.origin
             classKind = ClassKind.CLASS
             scopeProvider = session.kotlinScopeProvider
             status = FirResolvedDeclarationStatusImpl(Visibilities.Public, Modality.FINAL, EffectiveVisibility.Public)
@@ -141,9 +141,6 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
         return packageFqName == FOO_PACKAGE
     }
 
-    override val key: FirPluginKey
-        get() = Key
-
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {
         register(PREDICATE)
     }

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

@@ -89,7 +89,7 @@ class TopLevelDeclarationsGenerator(session: FirSession) : FirDeclarationGenerat
         }
     }
 
-    override val key: SomePluginKey
+    private val key: SomePluginKey
         get() = SomePluginKey
 
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {

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

@@ -11,12 +11,14 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.containingClassForStaticMemberAttr
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.builder.buildPrimaryConstructor
 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.moduleData
+import org.jetbrains.kotlin.fir.plugin.SomePluginKey
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
@@ -32,7 +34,8 @@ import org.jetbrains.kotlin.name.ClassId
 @OptIn(SymbolInternals::class)
 fun FirDeclarationGenerationExtension.buildMaterializeFunction(
     matchedClassSymbol: FirClassLikeSymbol<*>,
-    callableId: CallableId
+    callableId: CallableId,
+    key: FirPluginKey
 ): FirSimpleFunction {
     return buildSimpleFunction {
         moduleData = session.moduleData
@@ -59,7 +62,7 @@ fun FirDeclarationGenerationExtension.buildMaterializeFunction(
 }
 
 @OptIn(SymbolInternals::class)
-fun FirDeclarationGenerationExtension.buildConstructor(classId: ClassId, isInner: Boolean): FirConstructor {
+fun FirDeclarationGenerationExtension.buildConstructor(classId: ClassId, isInner: Boolean, key: FirPluginKey): FirConstructor {
     val lookupTag = ConeClassLikeLookupTagImpl(classId)
     return buildPrimaryConstructor {
         moduleData = session.moduleData

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

@@ -6,12 +6,10 @@
 package org.jetbrains.kotlin.fir.plugin.types
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyAnnotationArgumentMapping
 import org.jetbrains.kotlin.fir.extensions.FirTypeAttributeExtension
-import org.jetbrains.kotlin.fir.plugin.SomePluginKey
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.ConeAttribute
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
@@ -30,9 +28,6 @@ class FirNumberSignAttributeExtension(session: FirSession) : FirTypeAttributeExt
         private val NegativeClassId = ClassId(PACKAGE_FQN, Name.identifier("Negative"))
     }
 
-    override val key: FirPluginKey
-        get() = SomePluginKey
-
     override fun extractAttributeFromAnnotation(annotation: FirAnnotation): ConeAttribute<*>? {
         val sign = when (annotation.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.classId) {
             PositiveClassId -> ConeNumberSignAttribute.Sign.Positive

plugins/parcelize/parcelize-compiler/src/org/jetbrains/kotlin/parcelize/fir/FirParcelizeCheckersExtension.kt

@@ -10,13 +10,9 @@ import org.jetbrains.kotlin.fir.analysis.checkers.declaration.*
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirAnnotationCallChecker
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
-import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.parcelize.fir.diagnostics.*
 
 class FirParcelizeCheckersExtension(session: FirSession) : FirAdditionalCheckersExtension(session) {
-    override val key: FirPluginKey
-        get() = FirParcelizePluginKey
-
     override val expressionCheckers: ExpressionCheckers = object : ExpressionCheckers() {
         override val annotationCallCheckers: Set<FirAnnotationCallChecker>
             get() = setOf(FirParcelizeAnnotationChecker)

plugins/parcelize/parcelize-compiler/src/org/jetbrains/kotlin/parcelize/fir/FirParcelizeDeclarationGenerator.kt

@@ -176,7 +176,7 @@ class FirParcelizeDeclarationGenerator(session: FirSession) : FirDeclarationGene
         }
     }
 
-    override val key: FirPluginKey
+    private val key: FirPluginKey
         get() = FirParcelizePluginKey
 
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {