From 386f6d93316d6bbc448456eb11754e9595c3138f Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Wed, 9 Jun 2021 10:39:37 +0300
Subject: [PATCH] [FIR] Safe KClass of PSI type in diagnostic factory

---
 .../diagnostics/FirDiagnosticFactory.kt       | 27 +++++---
 .../diagnostics/FirDiagnosticFactoryDsl.kt    | 62 ++++++++++---------
 .../fir/handlers/FirDiagnosticsHandler.kt     |  2 +-
 .../kotlin/fir/AbstractFirDiagnosticTest.kt   |  4 +-
 4 files changed, 54 insertions(+), 41 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt
index 2858a1473da..ee3854906a0 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt
@@ -7,6 +7,7 @@ import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.fir.FirLightSourceElement
 import org.jetbrains.kotlin.fir.FirPsiSourceElement
 import org.jetbrains.kotlin.fir.FirSourceElement
+import kotlin.reflect.KClass
 
 @RequiresOptIn("Please use DiagnosticReporter.reportOn method if possible")
 annotation class InternalDiagnosticFactoryMethod
@@ -15,6 +16,7 @@ sealed class AbstractFirDiagnosticFactory<D : FirDiagnostic<*>, P : PsiElement>(
     val name: String,
     val severity: Severity,
     val defaultPositioningStrategy: SourceElementPositioningStrategy<P>,
+    val psiType: KClass<*>
 ) {
     abstract val firRenderer: FirDiagnosticRenderer<D>
 }
@@ -22,8 +24,9 @@ sealed class AbstractFirDiagnosticFactory<D : FirDiagnostic<*>, P : PsiElement>(
 class FirDiagnosticFactory0<P : PsiElement>(
     name: String,
     severity: Severity,
-    defaultPositioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT,
-) : AbstractFirDiagnosticFactory<FirSimpleDiagnostic<*>, P>(name, severity, defaultPositioningStrategy) {
+    defaultPositioningStrategy: SourceElementPositioningStrategy<P>,
+    psiType: KClass<*>
+) : AbstractFirDiagnosticFactory<FirSimpleDiagnostic<*>, P>(name, severity, defaultPositioningStrategy, psiType) {
     override val firRenderer: FirDiagnosticRenderer<FirSimpleDiagnostic<*>> = SimpleFirDiagnosticRenderer("")
 
     @InternalDiagnosticFactoryMethod
@@ -44,8 +47,9 @@ class FirDiagnosticFactory0<P : PsiElement>(
 class FirDiagnosticFactory1<P : PsiElement, A>(
     name: String,
     severity: Severity,
-    defaultPositioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT,
-) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters1<*, A>, P>(name, severity, defaultPositioningStrategy) {
+    defaultPositioningStrategy: SourceElementPositioningStrategy<P>,
+    psiType: KClass<*>
+) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters1<*, A>, P>(name, severity, defaultPositioningStrategy, psiType) {
     override val firRenderer: FirDiagnosticRenderer<FirDiagnosticWithParameters1<*, A>> = FirDiagnosticWithParameters1Renderer(
         "{0}",
         FirDiagnosticRenderers.TO_STRING
@@ -76,8 +80,9 @@ class FirDiagnosticFactory1<P : PsiElement, A>(
 class FirDiagnosticFactory2<P : PsiElement, A, B>(
     name: String,
     severity: Severity,
-    defaultPositioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT,
-) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters2<*, A, B>, P>(name, severity, defaultPositioningStrategy) {
+    defaultPositioningStrategy: SourceElementPositioningStrategy<P>,
+    psiType: KClass<*>
+) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters2<*, A, B>, P>(name, severity, defaultPositioningStrategy, psiType) {
     override val firRenderer: FirDiagnosticRenderer<FirDiagnosticWithParameters2<*, A, B>> = FirDiagnosticWithParameters2Renderer(
         "{0}, {1}",
         FirDiagnosticRenderers.TO_STRING,
@@ -111,8 +116,9 @@ class FirDiagnosticFactory2<P : PsiElement, A, B>(
 class FirDiagnosticFactory3<P : PsiElement, A, B, C>(
     name: String,
     severity: Severity,
-    defaultPositioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT,
-) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters3<*, A, B, C>, P>(name, severity, defaultPositioningStrategy) {
+    defaultPositioningStrategy: SourceElementPositioningStrategy<P>,
+    psiType: KClass<*>
+) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters3<*, A, B, C>, P>(name, severity, defaultPositioningStrategy, psiType) {
     override val firRenderer: FirDiagnosticRenderer<FirDiagnosticWithParameters3<*, A, B, C>> = FirDiagnosticWithParameters3Renderer(
         "{0}, {1}, {2}",
         FirDiagnosticRenderers.TO_STRING,
@@ -149,8 +155,9 @@ class FirDiagnosticFactory3<P : PsiElement, A, B, C>(
 class FirDiagnosticFactory4<P : PsiElement, A, B, C, D>(
     name: String,
     severity: Severity,
-    defaultPositioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT,
-) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters4<*, A, B, C, D>, P>(name, severity, defaultPositioningStrategy) {
+    defaultPositioningStrategy: SourceElementPositioningStrategy<P>,
+    psiType: KClass<*>
+) : AbstractFirDiagnosticFactory<FirDiagnosticWithParameters4<*, A, B, C, D>, P>(name, severity, defaultPositioningStrategy, psiType) {
     override val firRenderer: FirDiagnosticRenderer<FirDiagnosticWithParameters4<*, A, B, C, D>> = FirDiagnosticWithParameters4Renderer(
         "{0}, {1}, {2}, {3}",
         FirDiagnosticRenderers.TO_STRING,
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryDsl.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryDsl.kt
index 68c9cba9632..c62dd346526 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryDsl.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryDsl.kt
@@ -8,106 +8,112 @@ package org.jetbrains.kotlin.fir.analysis.diagnostics
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Severity
 import kotlin.properties.ReadOnlyProperty
+import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 
-fun <P : PsiElement> warning0(
+inline fun <reified P : PsiElement> warning0(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory0DelegateProvider<P> {
-    return DiagnosticFactory0DelegateProvider(Severity.WARNING, positioningStrategy)
+    return DiagnosticFactory0DelegateProvider(Severity.WARNING, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement, A> warning1(
+inline fun <reified P : PsiElement, A> warning1(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory1DelegateProvider<P, A> {
-    return DiagnosticFactory1DelegateProvider(Severity.WARNING, positioningStrategy)
+    return DiagnosticFactory1DelegateProvider(Severity.WARNING, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement, A, B> warning2(
+inline fun <reified P : PsiElement, A, B> warning2(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory2DelegateProvider<P, A, B> {
-    return DiagnosticFactory2DelegateProvider(Severity.WARNING, positioningStrategy)
+    return DiagnosticFactory2DelegateProvider(Severity.WARNING, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement, A, B, C> warning3(
+inline fun <reified P : PsiElement, A, B, C> warning3(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory3DelegateProvider<P, A, B, C> {
-    return DiagnosticFactory3DelegateProvider(Severity.WARNING, positioningStrategy)
+    return DiagnosticFactory3DelegateProvider(Severity.WARNING, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement> error0(
+inline fun <reified P : PsiElement> error0(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory0DelegateProvider<P> {
-    return DiagnosticFactory0DelegateProvider(Severity.ERROR, positioningStrategy)
+    return DiagnosticFactory0DelegateProvider(Severity.ERROR, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement, A> error1(
+inline fun <reified P : PsiElement, A> error1(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory1DelegateProvider<P, A> {
-    return DiagnosticFactory1DelegateProvider(Severity.ERROR, positioningStrategy)
+    return DiagnosticFactory1DelegateProvider(Severity.ERROR, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement, A, B> error2(
+inline fun <reified P : PsiElement, A, B> error2(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory2DelegateProvider<P, A, B> {
-    return DiagnosticFactory2DelegateProvider(Severity.ERROR, positioningStrategy)
+    return DiagnosticFactory2DelegateProvider(Severity.ERROR, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement, A, B, C> error3(
+inline fun <reified P : PsiElement, A, B, C> error3(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory3DelegateProvider<P, A, B, C> {
-    return DiagnosticFactory3DelegateProvider(Severity.ERROR, positioningStrategy)
+    return DiagnosticFactory3DelegateProvider(Severity.ERROR, positioningStrategy, P::class)
 }
 
-fun <P : PsiElement, A, B, C, D> error4(
+inline fun <reified P : PsiElement, A, B, C, D> error4(
     positioningStrategy: SourceElementPositioningStrategy<P> = SourceElementPositioningStrategy.DEFAULT
 ): DiagnosticFactory4DelegateProvider<P, A, B, C, D> {
-    return DiagnosticFactory4DelegateProvider(Severity.ERROR, positioningStrategy)
+    return DiagnosticFactory4DelegateProvider(Severity.ERROR, positioningStrategy, P::class)
 }
 
 // ------------------------------ Providers ------------------------------
 
 class DiagnosticFactory0DelegateProvider<P : PsiElement>(
     private val severity: Severity,
-    private val positioningStrategy: SourceElementPositioningStrategy<P>
+    private val positioningStrategy: SourceElementPositioningStrategy<P>,
+    private val psiType: KClass<*>
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory0<P>> {
-        return DummyDelegate(FirDiagnosticFactory0(prop.name, severity, positioningStrategy))
+        return DummyDelegate(FirDiagnosticFactory0(prop.name, severity, positioningStrategy, psiType))
     }
 }
 
 class DiagnosticFactory1DelegateProvider<P : PsiElement, A>(
     private val severity: Severity,
-    private val positioningStrategy: SourceElementPositioningStrategy<P>
+    private val positioningStrategy: SourceElementPositioningStrategy<P>,
+    private val psiType: KClass<*>
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory1<P, A>> {
-        return DummyDelegate(FirDiagnosticFactory1(prop.name, severity, positioningStrategy))
+        return DummyDelegate(FirDiagnosticFactory1(prop.name, severity, positioningStrategy, psiType))
     }
 }
 
 class DiagnosticFactory2DelegateProvider<P : PsiElement, A, B>(
     private val severity: Severity,
-    private val positioningStrategy: SourceElementPositioningStrategy<P>
+    private val positioningStrategy: SourceElementPositioningStrategy<P>,
+    private val psiType: KClass<*>
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory2<P, A, B>> {
-        return DummyDelegate(FirDiagnosticFactory2(prop.name, severity, positioningStrategy))
+        return DummyDelegate(FirDiagnosticFactory2(prop.name, severity, positioningStrategy, psiType))
     }
 }
 
 class DiagnosticFactory3DelegateProvider<P : PsiElement, A, B, C>(
     private val severity: Severity,
-    private val positioningStrategy: SourceElementPositioningStrategy<P>
+    private val positioningStrategy: SourceElementPositioningStrategy<P>,
+    private val psiType: KClass<*>
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory3<P, A, B, C>> {
-        return DummyDelegate(FirDiagnosticFactory3(prop.name, severity, positioningStrategy))
+        return DummyDelegate(FirDiagnosticFactory3(prop.name, severity, positioningStrategy, psiType))
     }
 }
 
 class DiagnosticFactory4DelegateProvider<P : PsiElement, A, B, C, D>(
     private val severity: Severity,
-    private val positioningStrategy: SourceElementPositioningStrategy<P>
+    private val positioningStrategy: SourceElementPositioningStrategy<P>,
+    private val psiType: KClass<*>
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory4<P, A, B, C, D>> {
-        return DummyDelegate(FirDiagnosticFactory4(prop.name, severity, positioningStrategy))
+        return DummyDelegate(FirDiagnosticFactory4(prop.name, severity, positioningStrategy, psiType))
     }
 }
 
diff --git a/compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt b/compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt
index c8583cdef39..4909853d96f 100644
--- a/compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt
+++ b/compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt
@@ -245,7 +245,7 @@ class FirDiagnosticsHandler(testServices: TestServices) : FirAnalysisHandler(tes
         }
 
         val argumentText = argument()
-        val factory = FirDiagnosticFactory1<PsiElement, String>(name, severity)
+        val factory = FirDiagnosticFactory1<PsiElement, String>(name, severity, SourceElementPositioningStrategy.DEFAULT, PsiElement::class)
         return when (positionedElement) {
             is FirPsiSourceElement<*> -> FirPsiDiagnosticWithParameters1(
                 positionedElement,
diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt
index 678588a99d4..6ce76ed8d81 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt
@@ -215,14 +215,14 @@ abstract class AbstractFirDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
                 sourceElement,
                 argumentText,
                 severity,
-                FirDiagnosticFactory1(name, severity),
+                FirDiagnosticFactory1(name, severity, SourceElementPositioningStrategy.DEFAULT, PsiElement::class),
                 SourceElementPositioningStrategy.DEFAULT
             )
             is FirLightSourceElement -> FirLightDiagnosticWithParameters1(
                 sourceElement,
                 argumentText,
                 severity,
-                FirDiagnosticFactory1<PsiElement, String>(name, severity),
+                FirDiagnosticFactory1(name, severity, SourceElementPositioningStrategy.DEFAULT, PsiElement::class),
                 SourceElementPositioningStrategy.DEFAULT
             )
         }