[FIR-TEST] Add validation of diagnostics for light tree tests

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir
 import junit.framework.TestCase
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.collectors.FirDiagnosticsCollector
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.resolve.dfa.FirControlFlowGraphReferenceImpl
@@ -77,16 +78,13 @@ abstract class AbstractFirDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
         )
     }
 
-    protected fun checkDiagnostics(file: File, testFiles: List<TestFile>, firFiles: List<FirFile>) {
+    protected open fun checkDiagnostics(file: File, testFiles: List<TestFile>, firFiles: List<FirFile>) {
-        val collectors = mutableMapOf<FirSession, AbstractDiagnosticCollector>()
+        val diagnostics = collectDiagnostics(firFiles)
         val actualText = StringBuilder()
         for (testFile in testFiles) {
             val firFile = firFiles.firstOrNull { it.psi == testFile.ktFile }
             if (firFile != null) {
-                val session = firFile.session
+                testFile.getActualText(diagnostics.getValue(firFile), actualText)
-                val collector = collectors.computeIfAbsent(session) { createCollector(session) }
-                val firDiagnostics = collector.collectDiagnostics(firFile)
-                testFile.getActualText(firDiagnostics, actualText)
             } else {
                 actualText.append(testFile.expectedText)
             }
@@ -94,7 +92,18 @@ abstract class AbstractFirDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
         KotlinTestUtils.assertEqualsToFile(file, actualText.toString())
     }
 
-    protected fun createCollector(session: FirSession): AbstractDiagnosticCollector {
+    protected fun collectDiagnostics(firFiles: List<FirFile>): Map<FirFile, List<FirDiagnostic<*>>> {
+        val collectors = mutableMapOf<FirSession, AbstractDiagnosticCollector>()
+        val result = mutableMapOf<FirFile, List<FirDiagnostic<*>>>()
+        for (firFile in firFiles) {
+            val session = firFile.session
+            val collector = collectors.computeIfAbsent(session) { createCollector(session) }
+            result[firFile] = collector.collectDiagnostics(firFile).toList()
+        }
+        return result
+    }
+
+    private fun createCollector(session: FirSession): AbstractDiagnosticCollector {
         return FirDiagnosticsCollector.create(session)
     }
 

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticsWithLightTreeTest.kt

@@ -5,9 +5,15 @@
 
 package org.jetbrains.kotlin.fir
 
+import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.diagnostics.DiagnosticUtils
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirLightDiagnostic
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import java.io.File
+import kotlin.math.abs
 
 abstract class AbstractFirDiagnosticsWithLightTreeTest : AbstractFirDiagnosticsTest() {
     override fun doTest(filePath: String) {
@@ -20,4 +26,80 @@ abstract class AbstractFirDiagnosticsWithLightTreeTest : AbstractFirDiagnosticsT
 
     override val useLightTree: Boolean
         get() = true
+
+    override fun checkDiagnostics(file: File, testFiles: List<TestFile>, firFiles: List<FirFile>) {
+        val fileToDiagnostics = collectDiagnostics(firFiles)
+        val missingDiagnostics = mutableListOf<MissingDiagnostic>()
+        for (testFile in testFiles) {
+            val firFile = firFiles.firstOrNull { it.name == testFile.name } ?: continue
+            val ktFile = testFile.ktFile!!
+            val diagnostics = fileToDiagnostics[firFile] ?: emptyList()
+
+            val actualDiagnostics = diagnostics.groupBy {
+                require(it is FirLightDiagnostic<*>)
+                it.element.startOffset
+            }.mapValues { (_, diagnostics) -> diagnostics.map { it.factory.name }.countEntries() }
+
+            val existingDiagnostics = testFile.diagnosedRanges.groupBy {
+                it.start
+            }.mapValues { (_, ranges) -> ranges.flatMap { it.getDiagnostics().map { it.name } }.countEntries() }
+
+            for (startOffset in actualDiagnostics.keys + existingDiagnostics.keys) {
+                val expected = existingDiagnostics[startOffset] ?: emptyMap()
+                val actual = actualDiagnostics[startOffset] ?: emptyMap()
+                for (name in expected.keys + actual.keys) {
+                    val expectedCount = expected[name] ?: 0
+                    val actualCount = actual[name] ?: 0
+                    if (expectedCount != actualCount) {
+                        missingDiagnostics += MissingDiagnostic(startOffset, name, expectedCount - actualCount, ktFile)
+                    }
+                }
+            }
+
+            if (missingDiagnostics.isNotEmpty()) {
+                val wasExpected = missingDiagnostics.filter { it.kind == MissingDiagnostic.Kind.WasExpected }
+                val isActual = missingDiagnostics.filter { it.kind == MissingDiagnostic.Kind.IsActual }
+
+                if (wasExpected.sumBy { it.count } == isActual.sumBy { it.count }) return
+
+                val message = buildString {
+                    fun MissingDiagnostic.errorMessage(): String {
+                        val position = DiagnosticUtils.getLineAndColumnInPsiFile(ktFile, TextRange(startOffset, startOffset + 1))
+                        // I don't know why but somehow line is greater for 3 than line in real file
+                        return "  $count of $name at ${position.line - 3}:${position.column}"
+                    }
+
+                    if (wasExpected.isNotEmpty()) {
+                        appendln("Some diagnostics was expected:")
+                        wasExpected.forEach {
+                            appendln(it.errorMessage())
+                        }
+                    }
+                    if (isActual.isNotEmpty()) {
+                        appendln("Some new diagnostics:")
+                        isActual.forEach {
+                            appendln(it.errorMessage())
+                        }
+                    }
+
+                }
+                kotlin.test.fail(message)
+            }
+        }
+
+    }
+
+    private fun <T> List<T>.countEntries(): Map<T, Int> {
+        return groupBy { it }.mapValues { (_, value) -> value.size }
+    }
+
+    private class MissingDiagnostic(val startOffset: Int, val name: String, diff: Int, ktFile: KtFile) {
+        val kind: Kind = if (diff > 0) Kind.WasExpected else Kind.IsActual
+        val count: Int = abs(diff)
+
+        enum class Kind {
+            WasExpected,
+            IsActual
+        }
+    }
 }