FIR: add transformer consistency test

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

@@ -20,6 +20,8 @@ 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.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition
 import org.jetbrains.kotlin.psi.KtFile
@@ -89,6 +91,12 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
             it.result
         }
 
+    private fun FirFile.transformChildren(): Set<FirElement> =
+        ConsistencyTransformer().let {
+            this@transformChildren.accept(it, Unit)
+            it.result
+        }
+
     protected fun FirFile.checkChildren() {
         val children = traverseChildren()
         val visitedChildren = visitChildren()
@@ -100,6 +108,17 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         }
     }
 
+    protected fun FirFile.checkTransformedChildren() {
+        val children = traverseChildren()
+        val transformedChildren = transformChildren()
+        children.removeAll(transformedChildren)
+        if (children.isNotEmpty()) {
+            val element = children.first()
+            val elementDump = StringBuilder().also { element.accept(FirRenderer(it)) }.toString()
+            throw AssertionError("FirElement ${element.javaClass} is not transformed: $elementDump")
+        }
+    }
+
     private class ConsistencyVisitor : FirVisitorVoid() {
         var result = hashSetOf<FirElement>()
 
@@ -116,6 +135,22 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         }
     }
 
+    private class ConsistencyTransformer : FirTransformer<Unit>() {
+        var result = hashSetOf<FirElement>()
+
+        override fun <E : FirElement> transformElement(element: E, data: Unit): CompositeTransformResult<E> {
+            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, Unit)
+            }
+            return CompositeTransformResult(element)
+        }
+    }
+
     override fun tearDown() {
         super.tearDown()
         FileTypeRegistry.ourInstanceGetter = Getter<FileTypeRegistry> { FileTypeManager.getInstance() }

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

@@ -46,7 +46,7 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
         println("TIME PER FILE: ${(time / counter) * 1e-6} ms, COUNTER: $counter")
     }
 
-    fun testConsistency() {
+    fun testVisitConsistency() {
         val root = File(testDataPath)
         for (file in root.walkTopDown()) {
             if (file.isDirectory) continue
@@ -63,4 +63,20 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
         }
     }
 
+    fun testTransformConsistency() {
+        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.checkTransformedChildren()
+            } catch (e: Throwable) {
+                println("EXCEPTION in: " + file.toRelativeString(root))
+                throw e
+            }
+        }
+    }
 }