From 2f63c8a46a4fa74dcab0ea2a36171e3cb77f0b96 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Wed, 25 Mar 2020 14:43:44 +0300
Subject: [PATCH] [FIR] Enhance diagnostic DSL to be able to use concrete
 factories

Before this commit, things like DiagnosticFactory0<KtDeclaration>
were effectively unusable.
---
 .../ErrorNodeDiagnosticCollectorComponent.kt  |  2 +-
 .../fir/analysis/diagnostics/FirDiagnostic.kt | 81 ++++++++++---------
 .../diagnostics/FirDiagnosticFactory.kt       | 42 ++++++----
 .../diagnostics/FirDiagnosticFactoryDsl.kt    | 64 ++++++++-------
 .../FirDiagnosticFactoryToRendererMap.kt      |  8 +-
 .../fir/analysis/diagnostics/FirErrors.kt     | 46 ++++++-----
 .../kotlin/fir/builder/RawFirBuilder.kt       |  6 +-
 .../jetbrains/kotlin/fir/FirSourceElement.kt  | 10 +--
 .../tests/ObjectWithConstructor.fir.kt        |  4 +-
 ...AbstractFirDiagnosticsWithLightTreeTest.kt |  2 +-
 10 files changed, 142 insertions(+), 123 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/ErrorNodeDiagnosticCollectorComponent.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/ErrorNodeDiagnosticCollectorComponent.kt
index 1f7a80ee302..db1dfac7911 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/ErrorNodeDiagnosticCollectorComponent.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/components/ErrorNodeDiagnosticCollectorComponent.kt
@@ -65,7 +65,7 @@ class ErrorNodeDiagnosticCollectorComponent(collector: AbstractDiagnosticCollect
         reporter.report(coneDiagnostic)
     }
 
-    private fun ConeSimpleDiagnostic.getFactory(): FirDiagnosticFactory0<FirSourceElement> {
+    private fun ConeSimpleDiagnostic.getFactory(): FirDiagnosticFactory0<FirSourceElement, *> {
         @Suppress("UNCHECKED_CAST")
         return when (kind) {
             DiagnosticKind.Syntax -> FirErrors.SYNTAX_ERROR
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnostic.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnostic.kt
index 9debe271925..14051374c77 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnostic.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnostic.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.analysis.diagnostics
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.fir.FirLightSourceElement
@@ -20,75 +21,75 @@ sealed class FirDiagnostic<out E : FirSourceElement> {
 }
 
 sealed class FirSimpleDiagnostic<out E : FirSourceElement> : FirDiagnostic<E>() {
-    abstract override val factory: FirDiagnosticFactory0<*>
+    abstract override val factory: FirDiagnosticFactory0<*, *>
 }
 
 sealed class FirDiagnosticWithParameters1<out E : FirSourceElement, A> : FirDiagnostic<E>() {
     abstract val a: A
-    abstract override val factory: FirDiagnosticFactory1<*, A>
+    abstract override val factory: FirDiagnosticFactory1<*, *, A>
 }
 
 sealed class FirDiagnosticWithParameters2<out E : FirSourceElement, A, B> : FirDiagnostic<E>() {
     abstract val a: A
     abstract val b: B
-    abstract override val factory: FirDiagnosticFactory2<*, A, B>
+    abstract override val factory: FirDiagnosticFactory2<*, *, A, B>
 }
 
 sealed class FirDiagnosticWithParameters3<out E : FirSourceElement, A, B, C> : FirDiagnostic<E>() {
     abstract val a: A
     abstract val b: B
     abstract val c: C
-    abstract override val factory: FirDiagnosticFactory3<*, A, B, C>
+    abstract override val factory: FirDiagnosticFactory3<*, *, A, B, C>
 }
 
 // ------------------------------ psi diagnostics ------------------------------
 
-interface FirPsiDiagnostic<out E : FirPsiSourceElement> {
+interface FirPsiDiagnostic<P : PsiElement> {
     fun asPsiBasedDiagnostic(): Diagnostic
-    val element: E
+    val element: FirPsiSourceElement<P>
 }
 
-data class FirPsiSimpleDiagnostic<out E : FirPsiSourceElement>(
-    override val element: E,
+data class FirPsiSimpleDiagnostic<P : PsiElement>(
+    override val element: FirPsiSourceElement<P>,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory0<*>
-) : FirSimpleDiagnostic<E>(), FirPsiDiagnostic<E> {
+    override val factory: FirDiagnosticFactory0<FirPsiSourceElement<P>, P>
+) : FirSimpleDiagnostic<FirPsiSourceElement<P>>(), FirPsiDiagnostic<P> {
     override fun asPsiBasedDiagnostic(): Diagnostic {
         return factory.psiDiagnosticFactory.on(element.psi)
     }
 }
 
-data class FirPsiDiagnosticWithParameters1<out E : FirPsiSourceElement, A>(
-    override val element: E,
+data class FirPsiDiagnosticWithParameters1<P : PsiElement, A>(
+    override val element: FirPsiSourceElement<P>,
     override val a: A,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory1<*, A>
-) : FirDiagnosticWithParameters1<E, A>(), FirPsiDiagnostic<E> {
+    override val factory: FirDiagnosticFactory1<FirPsiSourceElement<P>, P, A>
+) : FirDiagnosticWithParameters1<FirPsiSourceElement<P>, A>(), FirPsiDiagnostic<P> {
     override fun asPsiBasedDiagnostic(): Diagnostic {
         return factory.psiDiagnosticFactory.on(element.psi, a)
     }
 }
 
-data class FirPsiDiagnosticWithParameters2<out E : FirPsiSourceElement, A, B>(
-    override val element: E,
+data class FirPsiDiagnosticWithParameters2<P : PsiElement, A, B>(
+    override val element: FirPsiSourceElement<P>,
     override val a: A,
     override val b: B,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory2<*, A, B>
-) : FirDiagnosticWithParameters2<E, A, B>(), FirPsiDiagnostic<E> {
+    override val factory: FirDiagnosticFactory2<FirPsiSourceElement<P>, P, A, B>
+) : FirDiagnosticWithParameters2<FirPsiSourceElement<P>, A, B>(), FirPsiDiagnostic<P> {
     override fun asPsiBasedDiagnostic(): Diagnostic {
         return factory.psiDiagnosticFactory.on(element.psi, a, b)
     }
 }
 
-data class FirPsiDiagnosticWithParameters3<out E : FirPsiSourceElement, A, B, C>(
-    override val element: E,
+data class FirPsiDiagnosticWithParameters3<P : PsiElement, A, B, C>(
+    override val element: FirPsiSourceElement<P>,
     override val a: A,
     override val b: B,
     override val c: C,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory3<*, A, B, C>
-) : FirDiagnosticWithParameters3<E, A, B, C>(), FirPsiDiagnostic<E> {
+    override val factory: FirDiagnosticFactory3<FirPsiSourceElement<P>, P, A, B, C>
+) : FirDiagnosticWithParameters3<FirPsiSourceElement<P>, A, B, C>(), FirPsiDiagnostic<P> {
     override fun asPsiBasedDiagnostic(): Diagnostic {
         return factory.psiDiagnosticFactory.on(element.psi, a, b, c)
     }
@@ -96,36 +97,36 @@ data class FirPsiDiagnosticWithParameters3<out E : FirPsiSourceElement, A, B, C>
 
 // ------------------------------ light tree diagnostics ------------------------------
 
-interface FirLightDiagnostic<out E: FirLightSourceElement> {
-    val element: E
+interface FirLightDiagnostic {
+    val element: FirLightSourceElement
 }
 
-data class FirLightSimpleDiagnostic<out E : FirLightSourceElement>(
-    override val element: E,
+data class FirLightSimpleDiagnostic(
+    override val element: FirLightSourceElement,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory0<*>
-) : FirSimpleDiagnostic<E>(), FirLightDiagnostic<E>
+    override val factory: FirDiagnosticFactory0<*, *>
+) : FirSimpleDiagnostic<FirLightSourceElement>(), FirLightDiagnostic
 
-data class FirLightDiagnosticWithParameters1<out E : FirLightSourceElement, A>(
-    override val element: E,
+data class FirLightDiagnosticWithParameters1<A>(
+    override val element: FirLightSourceElement,
     override val a: A,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory1<*, A>
-) : FirDiagnosticWithParameters1<E, A>(), FirLightDiagnostic<E>
+    override val factory: FirDiagnosticFactory1<*, *, A>
+) : FirDiagnosticWithParameters1<FirLightSourceElement, A>(), FirLightDiagnostic
 
-data class FirLightDiagnosticWithParameters2<out E : FirLightSourceElement, A, B>(
-    override val element: E,
+data class FirLightDiagnosticWithParameters2<A, B>(
+    override val element: FirLightSourceElement,
     override val a: A,
     override val b: B,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory2<*, A, B>
-) : FirDiagnosticWithParameters2<E, A, B>(), FirLightDiagnostic<E>
+    override val factory: FirDiagnosticFactory2<*, *, A, B>
+) : FirDiagnosticWithParameters2<FirLightSourceElement, A, B>(), FirLightDiagnostic
 
-data class FirLightDiagnosticWithParameters3<out E : FirLightSourceElement, A, B, C>(
-    override val element: E,
+data class FirLightDiagnosticWithParameters3<A, B, C>(
+    override val element: FirLightSourceElement,
     override val a: A,
     override val b: B,
     override val c: C,
     override val severity: Severity,
-    override val factory: FirDiagnosticFactory3<*, A, B, C>
-) : FirDiagnosticWithParameters3<E, A, B, C>(), FirLightDiagnostic<E>
\ No newline at end of file
+    override val factory: FirDiagnosticFactory3<*, *, A, B, C>
+) : FirDiagnosticWithParameters3<FirLightSourceElement, A, B, C>(), FirLightDiagnostic
\ No newline at end of file
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 1ef91b7cea3..abc5c668b18 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
@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.fir.FirLightSourceElement
 import org.jetbrains.kotlin.fir.FirPsiSourceElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 
-sealed class AbstractFirDiagnosticFactory<E : FirSourceElement, D : FirDiagnostic<E>>(
+sealed class AbstractFirDiagnosticFactory<out E : FirSourceElement, D : FirDiagnostic<E>>(
     val name: String,
     val severity: Severity,
 ) {
@@ -22,51 +22,59 @@ sealed class AbstractFirDiagnosticFactory<E : FirSourceElement, D : FirDiagnosti
     }
 }
 
-class FirDiagnosticFactory0<E : FirSourceElement>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory0<PsiElement>
+class FirDiagnosticFactory0<E : FirSourceElement, P : PsiElement>(
+    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory0<P>
 ) : AbstractFirDiagnosticFactory<E, FirSimpleDiagnostic<E>>(name, severity) {
     fun on(element: E): FirSimpleDiagnostic<E> {
         return when (element) {
-            is FirPsiSourceElement -> FirPsiSimpleDiagnostic(element, severity, this)
+            is FirPsiSourceElement<*> -> FirPsiSimpleDiagnostic(
+                element as FirPsiSourceElement<P>, severity, this as FirDiagnosticFactory0<FirPsiSourceElement<P>, P>
+            )
             is FirLightSourceElement -> FirLightSimpleDiagnostic(element, severity, this)
             else -> incorrectElement(element)
-        }
+        } as FirSimpleDiagnostic<E>
     }
 }
 
-class FirDiagnosticFactory1<E : FirSourceElement, A>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory1<PsiElement, A>
+class FirDiagnosticFactory1<E : FirSourceElement, P : PsiElement, A>(
+    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory1<P, A>
 ) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters1<E, A>>(name, severity) {
     fun on(element: E, a: A): FirDiagnosticWithParameters1<E, A> {
         return when (element) {
-            is FirPsiSourceElement -> FirPsiDiagnosticWithParameters1(element, a, severity, this)
+            is FirPsiSourceElement<*> -> FirPsiDiagnosticWithParameters1(
+                element as FirPsiSourceElement<P>, a, severity, this as FirDiagnosticFactory1<FirPsiSourceElement<P>, P, A>
+            )
             is FirLightSourceElement -> FirLightDiagnosticWithParameters1(element, a, severity, this)
             else -> incorrectElement(element)
-        }
+        } as FirDiagnosticWithParameters1<E, A>
     }
 }
 
-class FirDiagnosticFactory2<E : FirSourceElement, A, B>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory2<PsiElement, A, B>
+class FirDiagnosticFactory2<E : FirSourceElement, P : PsiElement, A, B>(
+    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory2<P, A, B>
 ) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters2<E, A, B>>(name, severity) {
     fun on(element: E, a: A, b: B): FirDiagnosticWithParameters2<E, A, B> {
         return when (element) {
-            is FirPsiSourceElement -> FirPsiDiagnosticWithParameters2(element, a, b, severity, this)
+            is FirPsiSourceElement<*> -> FirPsiDiagnosticWithParameters2(
+                element as FirPsiSourceElement<P>, a, b, severity, this as FirDiagnosticFactory2<FirPsiSourceElement<P>, P, A, B>
+            )
             is FirLightSourceElement -> FirLightDiagnosticWithParameters2(element, a, b, severity, this)
             else -> incorrectElement(element)
-        }
+        } as FirDiagnosticWithParameters2<E, A, B>
     }
 }
 
-class FirDiagnosticFactory3<E : FirSourceElement, A, B, C>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory3<PsiElement, A, B, C>
+class FirDiagnosticFactory3<E : FirSourceElement, P : PsiElement, A, B, C>(
+    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory3<P, A, B, C>
 ) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters3<E, A, B, C>>(name, severity) {
     fun on(element: E, a: A, b: B, c: C): FirDiagnosticWithParameters3<E, A, B, C> {
         return when (element) {
-            is FirPsiSourceElement -> FirPsiDiagnosticWithParameters3(element, a, b, c, severity, this)
+            is FirPsiSourceElement<*> -> FirPsiDiagnosticWithParameters3(
+                element as FirPsiSourceElement<P>, a, b, c, severity, this as FirDiagnosticFactory3<FirPsiSourceElement<P>, P, A, B, C>
+            )
             is FirLightSourceElement -> FirLightDiagnosticWithParameters3(element, a, b, c, severity, this)
             else -> incorrectElement(element)
-        }
+        } as FirDiagnosticWithParameters3<E, A, B, C>
     }
 }
 
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 63631bba178..043169a6482 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
@@ -11,35 +11,35 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KProperty
 
-fun <E : FirSourceElement> warning0(): DiagnosticFactory0DelegateProvider<E> {
+fun <E : FirSourceElement, P : PsiElement> warning0(): DiagnosticFactory0DelegateProvider<E, P> {
     return DiagnosticFactory0DelegateProvider(Severity.WARNING, null)
 }
 
-fun <E : FirSourceElement, A> warning1(): DiagnosticFactory1DelegateProvider<E, A> {
+fun <E : FirSourceElement, P : PsiElement, A> warning1(): DiagnosticFactory1DelegateProvider<E, P, A> {
     return DiagnosticFactory1DelegateProvider(Severity.WARNING, null)
 }
 
-fun <E : FirSourceElement, A, B> warning2(): DiagnosticFactory2DelegateProvider<E, A, B> {
+fun <E : FirSourceElement, P : PsiElement, A, B> warning2(): DiagnosticFactory2DelegateProvider<E, P, A, B> {
     return DiagnosticFactory2DelegateProvider(Severity.WARNING, null)
 }
 
-fun <E : FirSourceElement, A, B, C> warning3(): DiagnosticFactory3DelegateProvider<E, A, B, C> {
+fun <E : FirSourceElement, P : PsiElement, A, B, C> warning3(): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
     return DiagnosticFactory3DelegateProvider(Severity.WARNING, null)
 }
 
-fun <E : FirSourceElement> error0(): DiagnosticFactory0DelegateProvider<E> {
+fun <E : FirSourceElement, P : PsiElement> error0(): DiagnosticFactory0DelegateProvider<E, P> {
     return DiagnosticFactory0DelegateProvider(Severity.ERROR, null)
 }
 
-fun <E : FirSourceElement, A> error1(): DiagnosticFactory1DelegateProvider<E, A> {
+fun <E : FirSourceElement, P : PsiElement, A> error1(): DiagnosticFactory1DelegateProvider<E, P, A> {
     return DiagnosticFactory1DelegateProvider(Severity.ERROR, null)
 }
 
-fun <E : FirSourceElement, A, B> error2(): DiagnosticFactory2DelegateProvider<E, A, B> {
+fun <E : FirSourceElement, P : PsiElement, A, B> error2(): DiagnosticFactory2DelegateProvider<E, P, A, B> {
     return DiagnosticFactory2DelegateProvider(Severity.ERROR, null)
 }
 
-fun <E : FirSourceElement, A, B, C> error3(): DiagnosticFactory3DelegateProvider<E, A, B, C> {
+fun <E : FirSourceElement, P : PsiElement, A, B, C> error3(): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
     return DiagnosticFactory3DelegateProvider(Severity.ERROR, null)
 }
 
@@ -47,66 +47,74 @@ fun <E : FirSourceElement, A, B, C> error3(): DiagnosticFactory3DelegateProvider
  * Note that those functions can be applicable only for factories
  *   that takes `PsiElement` as first type parameter
  */
-fun <E : FirSourceElement> existing(psiDiagnosticFactory: DiagnosticFactory0<PsiElement>): DiagnosticFactory0DelegateProvider<E> {
+fun <E : FirSourceElement, P : PsiElement> existing(
+    psiDiagnosticFactory: DiagnosticFactory0<P>
+): DiagnosticFactory0DelegateProvider<E, P> {
     return DiagnosticFactory0DelegateProvider(Severity.ERROR, psiDiagnosticFactory)
 }
 
-fun <E : FirSourceElement, A> existing(psiDiagnosticFactory: DiagnosticFactory1<PsiElement, A>): DiagnosticFactory1DelegateProvider<E, A> {
+fun <E : FirSourceElement, P : PsiElement, A> existing(
+    psiDiagnosticFactory: DiagnosticFactory1<P, A>
+): DiagnosticFactory1DelegateProvider<E, P, A> {
     return DiagnosticFactory1DelegateProvider(Severity.ERROR, psiDiagnosticFactory)
 }
 
-fun <E : FirSourceElement, A, B> existing(psiDiagnosticFactory: DiagnosticFactory2<PsiElement, A, B>): DiagnosticFactory2DelegateProvider<E, A, B> {
+fun <E : FirSourceElement, P : PsiElement, A, B> existing(
+    psiDiagnosticFactory: DiagnosticFactory2<P, A, B>
+): DiagnosticFactory2DelegateProvider<E, P, A, B> {
     return DiagnosticFactory2DelegateProvider(Severity.ERROR, psiDiagnosticFactory)
 }
 
-fun <E : FirSourceElement, A, B, C> existing(psiDiagnosticFactory: DiagnosticFactory3<PsiElement, A, B, C>): DiagnosticFactory3DelegateProvider<E, A, B, C> {
+fun <E : FirSourceElement, P : PsiElement, A, B, C> existing(
+    psiDiagnosticFactory: DiagnosticFactory3<P, A, B, C>
+): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
     return DiagnosticFactory3DelegateProvider(Severity.ERROR, psiDiagnosticFactory)
 }
 
 // ------------------------------ Providers ------------------------------
 
-class DiagnosticFactory0DelegateProvider<E : FirSourceElement>(
+class DiagnosticFactory0DelegateProvider<E : FirSourceElement, P : PsiElement>(
     private val severity: Severity,
-    private val psiDiagnosticFactory: DiagnosticFactory0<PsiElement>?
+    private val psiDiagnosticFactory: DiagnosticFactory0<P>?
 ) {
-    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory0<E>> {
-        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory0.create<PsiElement>(severity).apply {
+    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory0<E, P>> {
+        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory0.create<P>(severity).apply {
             initializeName(prop.name)
         }
         return DummyDelegate(FirDiagnosticFactory0(prop.name, severity, psiFactory))
     }
 }
 
-class DiagnosticFactory1DelegateProvider<E : FirSourceElement, A>(
+class DiagnosticFactory1DelegateProvider<E : FirSourceElement, P : PsiElement, A>(
     private val severity: Severity,
-    private val psiDiagnosticFactory: DiagnosticFactory1<PsiElement, A>?
+    private val psiDiagnosticFactory: DiagnosticFactory1<P, A>?
 ) {
-    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory1<E, A>> {
-        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory1.create<PsiElement, A>(severity).apply {
+    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory1<E, P, A>> {
+        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory1.create<P, A>(severity).apply {
             initializeName(prop.name)
         }
         return DummyDelegate(FirDiagnosticFactory1(prop.name, severity, psiFactory))
     }
 }
 
-class DiagnosticFactory2DelegateProvider<E : FirSourceElement, A, B>(
+class DiagnosticFactory2DelegateProvider<E : FirSourceElement, P : PsiElement, A, B>(
     private val severity: Severity,
-    private val psiDiagnosticFactory: DiagnosticFactory2<PsiElement, A, B>?
+    private val psiDiagnosticFactory: DiagnosticFactory2<P, A, B>?
 ) {
-    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory2<E, A, B>> {
-        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory2.create<PsiElement, A, B>(severity).apply {
+    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory2<E, P, A, B>> {
+        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory2.create<P, A, B>(severity).apply {
             initializeName(prop.name)
         }
         return DummyDelegate(FirDiagnosticFactory2(prop.name, severity, psiFactory))
     }
 }
 
-class DiagnosticFactory3DelegateProvider<E : FirSourceElement, A, B, C>(
+class DiagnosticFactory3DelegateProvider<E : FirSourceElement, P : PsiElement, A, B, C>(
     private val severity: Severity,
-    private val psiDiagnosticFactory: DiagnosticFactory3<PsiElement, A, B, C>?
+    private val psiDiagnosticFactory: DiagnosticFactory3<P, A, B, C>?
 ) {
-    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory3<E, A, B, C>> {
-        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory3.create<PsiElement, A, B, C>(severity).apply {
+    operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory3<E, P, A, B, C>> {
+        val psiFactory = psiDiagnosticFactory ?: DiagnosticFactory3.create<P, A, B, C>(severity).apply {
             initializeName(prop.name)
         }
         return DummyDelegate(FirDiagnosticFactory3(prop.name, severity, psiFactory))
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryToRendererMap.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryToRendererMap.kt
index 33d0f8a2618..74909add2e5 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryToRendererMap.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactoryToRendererMap.kt
@@ -16,13 +16,13 @@ class FirDiagnosticFactoryToRendererMap(val name: String) {
 
     operator fun get(factory: AbstractFirDiagnosticFactory<*, *>): DiagnosticRenderer<*>? = diagnosticsMap[factory]
 
-    fun put(factory: FirDiagnosticFactory0<*>, message: String) {
+    fun put(factory: FirDiagnosticFactory0<*, *>, message: String) {
         psiDiagnosticMap.put(factory.psiDiagnosticFactory, message)
         putToFirMap(factory)
     }
 
     fun <A> put(
-        factory: FirDiagnosticFactory1<*, A>,
+        factory: FirDiagnosticFactory1<*, *, A>,
         message: String,
         rendererA: DiagnosticParameterRenderer<A>?
     ) {
@@ -31,7 +31,7 @@ class FirDiagnosticFactoryToRendererMap(val name: String) {
     }
 
     fun <A, B> put(
-        factory: FirDiagnosticFactory2<*, A, B>,
+        factory: FirDiagnosticFactory2<*, *, A, B>,
         message: String,
         rendererA: DiagnosticParameterRenderer<A>?,
         rendererB: DiagnosticParameterRenderer<B>?
@@ -41,7 +41,7 @@ class FirDiagnosticFactoryToRendererMap(val name: String) {
     }
 
     fun <A, B, C> put(
-        factory: FirDiagnosticFactory3<*, A, B, C>,
+        factory: FirDiagnosticFactory3<*, *, A, B, C>,
         message: String,
         rendererA: DiagnosticParameterRenderer<A>?,
         rendererB: DiagnosticParameterRenderer<B>?,
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt
index da12ce7193b..97766c1f92f 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt
@@ -5,36 +5,38 @@
 
 package org.jetbrains.kotlin.fir.analysis.diagnostics
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
 
 object FirErrors {
-    val UNRESOLVED_REFERENCE by error1<FirSourceElement, String?>()
-    val INAPPLICABLE_CANDIDATE by error1<FirSourceElement, Collection<AbstractFirBasedSymbol<*>>>()
-    val AMBIGUITY by error1<FirSourceElement, Collection<AbstractFirBasedSymbol<*>>>()
-    val ASSIGN_OPERATOR_AMBIGUITY by error1<FirSourceElement, Collection<AbstractFirBasedSymbol<*>>>()
-    val SYNTAX_ERROR by error0<FirSourceElement>()
-    val UNRESOLVED_LABEL by error0<FirSourceElement>()
-    val ILLEGAL_CONST_EXPRESSION by error0<FirSourceElement>()
-    val DESERIALIZATION_ERROR by error0<FirSourceElement>()
-    val INFERENCE_ERROR by error0<FirSourceElement>()
-    val NO_SUPERTYPE by error0<FirSourceElement>()
-    val TYPE_PARAMETER_AS_SUPERTYPE by error0<FirSourceElement>()
-    val ENUM_AS_SUPERTYPE by error0<FirSourceElement>()
-    val RECURSION_IN_SUPERTYPES by error0<FirSourceElement>()
-    val RECURSION_IN_IMPLICIT_TYPES by error0<FirSourceElement>()
-    val ERROR_FROM_JAVA_RESOLUTION by error0<FirSourceElement>()
-    val OTHER_ERROR by error0<FirSourceElement>()
-    val TYPE_MISMATCH by error2<FirSourceElement, ConeKotlinType, ConeKotlinType>()
-    val VARIABLE_EXPECTED by error0<FirSourceElement>()
-    val RETURN_NOT_ALLOWED by error0<FirSourceElement>()
-    val CONSTRUCTOR_IN_OBJECT by error0<FirSourceElement>()
-    val SUPER_IS_NOT_AN_EXPRESSION by error0<FirSourceElement>()
+    val UNRESOLVED_REFERENCE by error1<FirSourceElement, PsiElement, String?>()
+    val INAPPLICABLE_CANDIDATE by error1<FirSourceElement, PsiElement, Collection<AbstractFirBasedSymbol<*>>>()
+    val AMBIGUITY by error1<FirSourceElement, PsiElement, Collection<AbstractFirBasedSymbol<*>>>()
+    val ASSIGN_OPERATOR_AMBIGUITY by error1<FirSourceElement, PsiElement, Collection<AbstractFirBasedSymbol<*>>>()
+    val SYNTAX_ERROR by error0<FirSourceElement, PsiElement>()
+    val UNRESOLVED_LABEL by error0<FirSourceElement, PsiElement>()
+    val ILLEGAL_CONST_EXPRESSION by error0<FirSourceElement, PsiElement>()
+    val DESERIALIZATION_ERROR by error0<FirSourceElement, PsiElement>()
+    val INFERENCE_ERROR by error0<FirSourceElement, PsiElement>()
+    val NO_SUPERTYPE by error0<FirSourceElement, PsiElement>()
+    val TYPE_PARAMETER_AS_SUPERTYPE by error0<FirSourceElement, PsiElement>()
+    val ENUM_AS_SUPERTYPE by error0<FirSourceElement, PsiElement>()
+    val RECURSION_IN_SUPERTYPES by error0<FirSourceElement, PsiElement>()
+    val RECURSION_IN_IMPLICIT_TYPES by error0<FirSourceElement, PsiElement>()
+    val ERROR_FROM_JAVA_RESOLUTION by error0<FirSourceElement, PsiElement>()
+    val OTHER_ERROR by error0<FirSourceElement, PsiElement>()
+    val TYPE_MISMATCH by error2<FirSourceElement, PsiElement, ConeKotlinType, ConeKotlinType>()
+    val VARIABLE_EXPECTED by error0<FirSourceElement, PsiElement>()
+    val RETURN_NOT_ALLOWED by error0<FirSourceElement, PsiElement>()
+    val SUPER_IS_NOT_AN_EXPRESSION by error0<FirSourceElement, PsiElement>()
 
-    val INAPPLICABLE_INFIX_MODIFIER by existing<FirSourceElement, String>(Errors.INAPPLICABLE_INFIX_MODIFIER)
+    val INAPPLICABLE_INFIX_MODIFIER by existing<FirSourceElement, PsiElement, String>(Errors.INAPPLICABLE_INFIX_MODIFIER)
+    val CONSTRUCTOR_IN_OBJECT by existing<FirSourceElement, KtDeclaration>(Errors.CONSTRUCTOR_IN_OBJECT)
 
     val REPEATED_MODIFIER by error1<FirSourceElement, KtModifierKeywordToken>()
     val REDUNDANT_MODIFIER by error2<FirSourceElement, KtModifierKeywordToken, KtModifierKeywordToken>()
diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
index e19870fccab..1f17c1b8d52 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
@@ -53,7 +53,7 @@ class RawFirBuilder(
         return file.accept(Visitor(), Unit) as FirFile
     }
 
-    override fun PsiElement.toFirSourceElement(): FirPsiSourceElement {
+    override fun PsiElement.toFirSourceElement(): FirPsiSourceElement<*> {
         return FirPsiSourceElement(this)
     }
 
@@ -1475,7 +1475,7 @@ class RawFirBuilder(
 
         private fun splitToCalleeAndReceiver(
             calleeExpression: KtExpression?,
-            defaultSource: FirPsiSourceElement,
+            defaultSource: FirPsiSourceElement<*>,
         ): Pair<FirNamedReference, FirExpression?> {
             return when (calleeExpression) {
                 is KtSimpleNameExpression -> buildSimpleNamedReference {
@@ -1535,7 +1535,7 @@ class RawFirBuilder(
         override fun visitArrayAccessExpression(expression: KtArrayAccessExpression, data: Unit): FirElement {
             val arrayExpression = expression.arrayExpression
             return buildFunctionCall {
-                val source: FirPsiSourceElement
+                val source: FirPsiSourceElement<*>
                 val getArgument = context.arraySetArgument.remove(expression)
                 if (getArgument != null) {
                     calleeReference = buildSimpleNamedReference {
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt
index 791604219e7..56d06a23138 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt
@@ -21,7 +21,7 @@ sealed class FirSourceElement {
     abstract val endOffset: Int
 }
 
-class FirPsiSourceElement(val psi: PsiElement) : FirSourceElement() {
+class FirPsiSourceElement<P : PsiElement>(val psi: P) : FirSourceElement() {
     override val elementType: IElementType
         get() = psi.node.elementType
 
@@ -42,19 +42,19 @@ class FirLightSourceElement(
         get() = element.tokenType
 }
 
-val FirSourceElement?.psi: PsiElement? get() = (this as? FirPsiSourceElement)?.psi
+val FirSourceElement?.psi: PsiElement? get() = (this as? FirPsiSourceElement<*>)?.psi
 
-val FirElement.psi: PsiElement? get() = (source as? FirPsiSourceElement)?.psi
+val FirElement.psi: PsiElement? get() = (source as? FirPsiSourceElement<*>)?.psi
 
 @Suppress("NOTHING_TO_INLINE")
-inline fun PsiElement.toFirSourceElement(): FirPsiSourceElement = FirPsiSourceElement(this)
+inline fun PsiElement.toFirSourceElement(): FirPsiSourceElement<*> = FirPsiSourceElement(this)
 
 val FirSourceElement?.lightNode: LighterASTNode? get() = (this as? FirLightSourceElement)?.element
 
 fun FirSourceElement?.getModifierList(): FirModifierList? {
     return when (this) {
         null -> null
-        is FirPsiSourceElement -> (psi as? KtModifierListOwner)?.modifierList?.let { FirPsiModifierList(it) }
+        is FirPsiSourceElement<*> -> (psi as? KtModifierListOwner)?.modifierList?.let { FirPsiModifierList(it) }
         is FirLightSourceElement -> {
             val kidsRef = Ref<Array<LighterASTNode?>>()
             tree.getChildren(element, kidsRef)
diff --git a/compiler/testData/diagnostics/tests/ObjectWithConstructor.fir.kt b/compiler/testData/diagnostics/tests/ObjectWithConstructor.fir.kt
index a54a650bf69..7ee1f433ccc 100644
--- a/compiler/testData/diagnostics/tests/ObjectWithConstructor.fir.kt
+++ b/compiler/testData/diagnostics/tests/ObjectWithConstructor.fir.kt
@@ -11,10 +11,10 @@ object A2 public constructor(private val prop: Int) {
 }
 
 val x = object (val prop: Int) {
-    <!CONSTRUCTOR_IN_OBJECT, CONSTRUCTOR_IN_OBJECT!>constructor() : <!UNRESOLVED_REFERENCE!>this<!>(1) {
+    <!CONSTRUCTOR_IN_OBJECT, CONSTRUCTOR_IN_OBJECT!>constructor()<!> : <!UNRESOLVED_REFERENCE!>this<!>(1) {
         val x = 1
         x * x
-    }<!>
+    }
 }
 
 class A3 {
diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticsWithLightTreeTest.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticsWithLightTreeTest.kt
index 5e34087c4b3..de3d7dcea23 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticsWithLightTreeTest.kt
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticsWithLightTreeTest.kt
@@ -36,7 +36,7 @@ abstract class AbstractFirDiagnosticsWithLightTreeTest : AbstractFirDiagnosticsT
             val diagnostics = fileToDiagnostics[firFile] ?: emptyList()
 
             val actualDiagnostics = diagnostics.groupBy {
-                require(it is FirLightDiagnostic<*>)
+                require(it is FirLightDiagnostic)
                 it.element.startOffset
             }.mapValues { (_, diagnostics) -> diagnostics.map { it.factory.name }.countEntries() }