diff --git a/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/PathWalker.kt b/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/PathWalker.kt
new file mode 100644
index 00000000000..156bcfb30f0
--- /dev/null
+++ b/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/PathWalker.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.lightTree
+
+import java.io.File
+
+fun String.walkTopDown(f: (File) -> Unit) {
+    val root = File(this)
+    for (file in root.walkTopDown()) {
+        if (file.isDirectory) continue
+        if (file.path.contains("testData") || file.path.contains("resources")) continue
+        if (file.extension != "kt") continue
+
+        f(file)
+    }
+}
+
+fun String.walkTopDownWithTestData(f: (File) -> Unit) {
+    val root = File(this)
+    for (file in root.walkTopDown()) {
+        if (file.isDirectory) continue
+        if (file.extension != "kt") continue
+
+        f(file)
+    }
+}
\ No newline at end of file
diff --git a/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/TotalKotlinTest.kt b/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/TotalKotlinTest.kt
index fefd0b54a9f..a979da6bf19 100644
--- a/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/TotalKotlinTest.kt
+++ b/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/TotalKotlinTest.kt
@@ -44,7 +44,6 @@ class TotalKotlinTest : AbstractRawFirBuilderTestCase() {
 
     private fun totalKotlinLight(onlyLightTree: Boolean) {
         val path = System.getProperty("user.dir")
-        val root = File(path)
         var counter = 0
         var time = 0L
 
@@ -54,14 +53,10 @@ class TotalKotlinTest : AbstractRawFirBuilderTestCase() {
 
         if (onlyLightTree) println("LightTree generation") else println("Fir from LightTree converter")
         println("BASE PATH: $path")
-        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 text = FileUtil.loadFile(file, CharsetToolkit.UTF8, true).trim()
+        path.walkTopDown {
+            val text = FileUtil.loadFile(it, CharsetToolkit.UTF8, true).trim()
             time += measureNanoTime {
-                generateFirFromLightTree(onlyLightTree, lightTreeConverter, text, file.name)
+                generateFirFromLightTree(onlyLightTree, lightTreeConverter, text, it.name)
             }
 
             counter++
diff --git a/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/benchmark/AbstractBenchmark.kt b/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/benchmark/AbstractBenchmark.kt
index 08fda50ecde..22ec66dcc5a 100644
--- a/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/benchmark/AbstractBenchmark.kt
+++ b/compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/benchmark/AbstractBenchmark.kt
@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.fir.lightTree.benchmark
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.openapi.vfs.CharsetToolkit
 import org.jetbrains.kotlin.fir.lightTree.benchmark.generators.TreeGenerator
+import org.jetbrains.kotlin.fir.lightTree.walkTopDown
+import org.jetbrains.kotlin.fir.lightTree.walkTopDownWithTestData
 import org.openjdk.jmh.annotations.*
 import java.io.File
 import java.util.concurrent.TimeUnit
@@ -23,17 +25,16 @@ abstract class AbstractBenchmark {
     abstract val generator: TreeGenerator
 
     protected fun readFiles(ignoreTestData: Boolean, path: String) {
-        val root = File(path)
-
         println("BASE PATH: $path")
-        for (file in root.walkTopDown()) {
-            if (file.isDirectory) continue
-            if (ignoreTestData && (file.path.contains("testData") || file.path.contains("resources"))) continue
-            if (file.extension != "kt") continue
-
+        val saveText: (File) -> Unit = { file ->
             val text = FileUtil.loadFile(file, CharsetToolkit.UTF8, true).trim()
             files[file] = text
         }
+        if (ignoreTestData) {
+            path.walkTopDownWithTestData(saveText)
+        } else {
+            path.walkTopDown(saveText)
+        }
     }
 
     protected fun forEachFile(f: (String, File) -> Unit) {