From 8dce8c1549692c851e3a2ad473e3f7900377b1e9 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Fri, 18 Jan 2019 17:58:24 +0300
Subject: [PATCH] Make correct traverse in both FIR consistency tests

---
 .../builder/AbstractRawFirBuilderTestCase.kt  | 22 ++++++++++++------
 .../RawFirBuilderTotalKotlinTestCase.kt       | 23 ++++++-------------
 2 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt b/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt
index cf79dc9d116..a946c2e01db 100644
--- a/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt
+++ b/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt
@@ -19,6 +19,7 @@ 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.render
 import org.jetbrains.kotlin.fir.types.FirType
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
@@ -31,6 +32,7 @@ import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.testFramework.KtParsingTestCase
 import java.io.File
 import kotlin.reflect.full.memberProperties
+import kotlin.reflect.jvm.isAccessible
 
 abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
     "",
@@ -78,9 +80,15 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         if (!result.add(this)) {
             return result
         }
-        for (property in this::class.memberProperties) {
-            val childElement = property.getter.call(this) as? FirElement ?: continue
-            childElement.traverseChildren(result)
+        propertyLoop@ for (property in this::class.memberProperties) {
+            val childElement = property.getter.apply { isAccessible = true }.call(this)
+
+            when (childElement) {
+                is FirElement -> childElement.traverseChildren(result)
+                is List<*> -> childElement.filterIsInstance<FirElement>().forEach { it.traverseChildren(result) }
+                else -> continue@propertyLoop
+            }
+
         }
         return result
     }
@@ -93,7 +101,7 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
 
     private fun FirFile.transformChildren(): Set<FirElement> =
         ConsistencyTransformer().let {
-            this@transformChildren.accept(it, Unit)
+            this@transformChildren.transform<FirFile, Unit>(it, Unit)
             it.result
         }
 
@@ -103,7 +111,7 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         children.removeAll(visitedChildren)
         if (children.isNotEmpty()) {
             val element = children.first()
-            val elementDump = StringBuilder().also { element.accept(FirRenderer(it)) }.toString()
+            val elementDump = element.render()
             throw AssertionError("FirElement ${element.javaClass} is not visited: $elementDump")
         }
     }
@@ -114,7 +122,7 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         children.removeAll(transformedChildren)
         if (children.isNotEmpty()) {
             val element = children.first()
-            val elementDump = StringBuilder().also { element.accept(FirRenderer(it)) }.toString()
+            val elementDump = element.render()
             throw AssertionError("FirElement ${element.javaClass} is not transformed: $elementDump")
         }
     }
@@ -145,7 +153,7 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
                     throw AssertionError("FirElement ${element.javaClass} is visited twice: $elementDump")
                 }
             } else {
-                element.acceptChildren(this, Unit)
+                element.transformChildren(this, Unit)
             }
             return CompositeTransformResult(element)
         }
diff --git a/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTotalKotlinTestCase.kt b/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTotalKotlinTestCase.kt
index 30dc8a5910f..52f52d06a08 100644
--- a/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTotalKotlinTestCase.kt
+++ b/compiler/fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTotalKotlinTestCase.kt
@@ -141,7 +141,7 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
         testTotalKotlinWithGivenMode(stubMode = true)
     }
 
-    fun testVisitConsistency() {
+    private fun testConsistency(checkConsistency: FirFile.() -> Unit) {
         val root = File(testDataPath)
         for (file in root.walkTopDown()) {
             if (file.isDirectory) continue
@@ -150,7 +150,7 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
             val ktFile = createKtFile(file.toRelativeString(root))
             val firFile = ktFile.toFirFile(stubMode = false)
             try {
-                firFile.checkChildren()
+                firFile.checkConsistency()
             } catch (e: Throwable) {
                 println("EXCEPTION in: " + file.toRelativeString(root))
                 throw e
@@ -158,20 +158,11 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
         }
     }
 
+    fun testVisitConsistency() {
+        testConsistency { checkChildren() }
+    }
+
     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(stubMode = false)
-            try {
-                firFile.checkTransformedChildren()
-            } catch (e: Throwable) {
-                println("EXCEPTION in: " + file.toRelativeString(root))
-                throw e
-            }
-        }
+        testConsistency { checkTransformedChildren() }
     }
 }
\ No newline at end of file