From 343928b97e2430195dcbca1e7d8cc3600325a6b3 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Wed, 29 Apr 2020 13:00:22 +0300
Subject: [PATCH] [FIR] Add proper extension key to FirExtensionPoint

---
 .../kotlin/fir/extensions/FirClassGenerationExtension.kt       | 3 +++
 .../org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt   | 3 +++
 .../kotlin/fir/extensions/FirExtensionPointService.kt          | 2 +-
 .../kotlin/fir/extensions/FirStatusTransfomerExtension.kt      | 3 +++
 4 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirClassGenerationExtension.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirClassGenerationExtension.kt
index d3ca9cc9ed9..fc17de1f1aa 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirClassGenerationExtension.kt
+++ b/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
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt
index ce67c2ad451..7cc652976f8 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt
+++ b/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
     }
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt
index 50e823e7a24..c699eda165b 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt
+++ b/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) {
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt
index 93f81429d1e..c811cfc66aa 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt
+++ b/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>