Add visitor consistency test

compiler/tests-common/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt

@@ -15,9 +15,12 @@ import com.intellij.psi.tree.IElementType
 import com.intellij.util.PathUtil
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.cli.common.script.CliScriptDefinitionProvider
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.FirSessionBase
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.types.FirType
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtPsiFactory
@@ -25,6 +28,7 @@ import org.jetbrains.kotlin.script.ScriptDefinitionProvider
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.testFramework.KtParsingTestCase
 import java.io.File
+import kotlin.reflect.full.memberProperties
 
 abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
     ".",
@@ -51,17 +55,11 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
     }
 
     protected fun doRawFirTest(filePath: String) {
-        doRawFirTest(filePath, checkNeeded = true)
-    }
-
-    protected fun doRawFirTest(filePath: String, checkNeeded: Boolean) {
         val file = createKtFile(filePath)
         val firFile = file.toFirFile()
-        if (checkNeeded) {
-            val firFileDump = StringBuilder().also { FirRenderer(it).visitFile(firFile) }.toString()
-            val expectedPath = filePath.replace(".kt", ".txt")
-            KotlinTestUtils.assertEqualsToFile(File(expectedPath), firFileDump)
-        }
+        val firFileDump = StringBuilder().also { FirRenderer(it).visitFile(firFile) }.toString()
+        val expectedPath = filePath.replace(".kt", ".txt")
+        KotlinTestUtils.assertEqualsToFile(File(expectedPath), firFileDump)
     }
 
     protected fun createKtFile(filePath: String): KtFile {
@@ -74,6 +72,50 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
     protected fun KtFile.toFirFile(): FirFile =
         RawFirBuilder(object : FirSessionBase() {}).buildFirFile(this)
 
+    private fun FirElement.traverseChildren(result: MutableSet<FirElement> = hashSetOf()): MutableSet<FirElement> {
+        if (!result.add(this)) {
+            return result
+        }
+        for (property in this::class.memberProperties) {
+            val childElement = property.getter.call(this) as? FirElement ?: continue
+            childElement.traverseChildren(result)
+        }
+        return result
+    }
+
+    private fun FirFile.visitChildren(): Set<FirElement> =
+        ConsistencyVisitor().let {
+            this@visitChildren.accept(it)
+            it.result
+        }
+
+    protected fun FirFile.checkChildren() {
+        val children = traverseChildren()
+        val visitedChildren = visitChildren()
+        children.removeAll(visitedChildren)
+        if (children.isNotEmpty()) {
+            val element = children.first()
+            val elementDump = StringBuilder().also { element.accept(FirRenderer(it)) }.toString()
+            throw AssertionError("FirElement ${element.javaClass} is not visited: $elementDump")
+        }
+    }
+
+    private class ConsistencyVisitor : FirVisitorVoid() {
+        var result = hashSetOf<FirElement>()
+
+        override fun visitElement(element: FirElement) {
+            // NB: types are reused sometimes (e.g. in accessors)
+            if (!result.add(element)) {
+                if (element !is FirType) {
+                    val elementDump = StringBuilder().also { element.accept(FirRenderer(it)) }.toString()
+                    throw AssertionError("FirElement ${element.javaClass} is visited twice: $elementDump")
+                }
+            } else {
+                element.acceptChildren(this)
+            }
+        }
+    }
+
     override fun tearDown() {
         super.tearDown()
         FileTypeRegistry.ourInstanceGetter = Getter<FileTypeRegistry> { FileTypeManager.getInstance() }

compiler/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTotalKotlinTestCase.kt

@@ -46,4 +46,21 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
         println("TIME PER FILE: ${(time / counter) * 1e-6} ms, COUNTER: $counter")
     }
 
+    fun testConsistency() {
+        val root = File(testDataPath)
+        for (file in root.walkTopDown()) {
+            if (file.isDirectory) continue
+            if (file.path.contains("testData") || file.path.contains("resources")) continue
+            if (file.extension != "kt") continue
+            val ktFile = createKtFile(file.toRelativeString(root))
+            val firFile = ktFile.toFirFile()
+            try {
+                firFile.checkChildren()
+            } catch (e: Throwable) {
+                println("EXCEPTION in: " + file.toRelativeString(root))
+                throw e
+            }
+        }
+    }
+
 }