[FIR] Add proper extension key to FirExtensionPoint

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirClassGenerationExtension.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import kotlin.reflect.KClass
 
 /*
  * TODO:
@@ -26,6 +27,8 @@ abstract class FirClassGenerationExtension(session: FirSession) : FirExtensionPo
 
     final override val mode: Mode = Mode.ANNOTATED_ELEMENT
 
+    final override val extensionType: KClass<out FirExtensionPoint> = FirClassGenerationExtension::class
+
     abstract fun <T> generateClass(
         containingFile: FirFile,
         annotatedDeclaration: T

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.extensions
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import kotlin.reflect.KClass
 
 typealias AnnotationFqn = FqName
 
@@ -19,6 +20,8 @@ abstract class FirExtensionPoint(val session: FirSession) {
 
     abstract val mode: Mode
 
+    internal abstract val extensionType: KClass<out FirExtensionPoint>
+
     fun interface Factory<P : FirExtensionPoint> {
         fun create(session: FirSession): P
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt

@@ -96,7 +96,7 @@ class FirExtensionPointService(
             val extensions = extensionsWithMetaAnnotations[metaAnnotation]
             if (extensions.isEmpty()) return
             for (extension in extensions) {
-                val registeredExtensions = this[extension::class]
+                val registeredExtensions = this[extension.extensionType]
 
                 @Suppress("UNCHECKED_CAST")
                 val map = when (extension.mode) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.FirResolvedDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import kotlin.reflect.KClass
 
 abstract class FirStatusTransformerExtension(session: FirSession) : FirExtensionPoint(session) {
     companion object {
@@ -22,6 +23,8 @@ abstract class FirStatusTransformerExtension(session: FirSession) : FirExtension
     final override val name: FirExtensionPointName
         get() = NAME
 
+    final override val extensionType: KClass<out FirExtensionPoint> = FirStatusTransformerExtension::class
+
     abstract fun transformStatus(declaration: FirDeclaration, status: FirDeclarationStatus): FirDeclarationStatus
 
     fun interface Factory : FirExtensionPoint.Factory<FirStatusTransformerExtension>