[Test] Parallelize test generation

compiler/tests-common/tests/org/jetbrains/kotlin/generators/impl/TestGenerationDSL.kt

@@ -8,26 +8,28 @@ package org.jetbrains.kotlin.generators.impl
 import org.jetbrains.kotlin.generators.InconsistencyChecker
 import org.jetbrains.kotlin.generators.InconsistencyChecker.Companion.inconsistencyChecker
 import org.jetbrains.kotlin.generators.TestGroupSuite
+import org.jetbrains.kotlin.generators.forEachTestClassParallel
 import org.jetbrains.kotlin.generators.testGroupSuite
+import org.jetbrains.kotlin.generators.util.TestGeneratorUtil
 
 fun generateTestGroupSuite(
     args: Array<String>,
+    mainClassName: String? = TestGeneratorUtil.getMainClassName(),
     init: TestGroupSuite.() -> Unit
 ) {
-    generateTestGroupSuite(InconsistencyChecker.hasDryRunArg(args), init)
+    generateTestGroupSuite(InconsistencyChecker.hasDryRunArg(args), mainClassName, init)
 }
 
 fun generateTestGroupSuite(
     dryRun: Boolean = false,
-    init: TestGroupSuite.() -> Unit
+    mainClassName: String? = TestGeneratorUtil.getMainClassName(),
+    init: TestGroupSuite.() -> Unit,
 ) {
     val suite = testGroupSuite(init)
-    for (testGroup in suite.testGroups) {
+    suite.forEachTestClassParallel { testClass ->
-        for (testClass in testGroup.testClasses) {
+        val (changed, testSourceFilePath) = TestGeneratorImpl.generateAndSave(testClass, dryRun, mainClassName)
-            val (changed, testSourceFilePath) = TestGeneratorImpl.generateAndSave(testClass, dryRun)
+        if (changed) {
-            if (changed) {
+            inconsistencyChecker(dryRun).add(testSourceFilePath)
-                inconsistencyChecker(dryRun).add(testSourceFilePath)
-            }
         }
     }
 }

compiler/tests-common/tests/org/jetbrains/kotlin/generators/impl/TestGeneratorImpl.kt

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.generators.TestGenerator
 import org.jetbrains.kotlin.generators.TestGroup
 import org.jetbrains.kotlin.generators.model.*
 import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil
-import org.jetbrains.kotlin.generators.util.TestGeneratorUtil.getMainClassName
 import org.jetbrains.kotlin.test.JUnit3RunnerWithInners
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.TestMetadata
@@ -31,14 +30,15 @@ private val METHOD_GENERATORS = listOf(
 )
 
 object TestGeneratorImpl : TestGenerator(METHOD_GENERATORS) {
-    override fun generateAndSave(testClass: TestGroup.TestClass, dryRun: Boolean): GenerationResult {
+    override fun generateAndSave(testClass: TestGroup.TestClass, dryRun: Boolean, mainClassName: String?): GenerationResult {
         val generatorInstance = TestGeneratorImplInstance(
             testClass.baseDir,
             testClass.suiteTestClassName,
             testClass.baseTestClassName,
             testClass.testModels,
             testClass.useJunit4,
-            methodGenerators
+            methodGenerators,
+            mainClassName,
         )
         return generatorInstance.generateAndSave(dryRun)
     }
@@ -50,7 +50,8 @@ private class TestGeneratorImplInstance(
     baseTestClassFqName: String,
     private val testClassModels: Collection<TestClassModel>,
     private val useJunit4: Boolean,
-    private val methodGenerators: Map<MethodModel.Kind, MethodGenerator<*>>
+    private val methodGenerators: Map<MethodModel.Kind, MethodGenerator<*>>,
+    private val mainClassName: String?
 ) {
     companion object {
         private val GENERATED_FILES = HashSet<String>()
@@ -147,7 +148,7 @@ private class TestGeneratorImplInstance(
         p.println("import java.io.File;")
         p.println("import java.util.regex.Pattern;")
         p.println()
-        p.println("/** This class is generated by {@link ", getMainClassName(), "}. DO NOT MODIFY MANUALLY */")
+        p.println("/** This class is generated by {@link ", mainClassName, "}. DO NOT MODIFY MANUALLY */")
 
         generateSuppressAllWarnings(p)
 

compiler/tests-for-compiler-generator/tests/org/jetbrains/kotlin/test/generators/GenerateCompilerTests.kt

@@ -5,9 +5,15 @@
 
 package org.jetbrains.kotlin.test.generators
 
+import org.jetbrains.kotlin.generators.util.TestGeneratorUtil
+import java.util.stream.Stream
+
 fun main(args: Array<String>) {
     System.setProperty("java.awt.headless", "true")
+    // Determine main class name while still on main thread.
+    val mainClassName = TestGeneratorUtil.getMainClassName()
 
-    generateJUnit3CompilerTests(args)
+    Stream.of(::generateJUnit3CompilerTests, ::generateJUnit5CompilerTests)
-    generateJUnit5CompilerTests(args)
+        .parallel()
+        .forEach { it.invoke(args, mainClassName) }
 }

compiler/tests-for-compiler-generator/tests/org/jetbrains/kotlin/test/generators/GenerateJUnit3CompilerTests.kt

@@ -50,10 +50,10 @@ import org.jetbrains.kotlin.test.TargetBackend
 import org.jetbrains.kotlin.test.utils.CUSTOM_TEST_DATA_EXTENSION_PATTERN
 import org.jetbrains.kotlin.types.AbstractTypeBindingTest
 
-fun generateJUnit3CompilerTests(args: Array<String>) {
+fun generateJUnit3CompilerTests(args: Array<String>, mainClassName: String?) {
     val excludedCustomTestdataPattern = CUSTOM_TEST_DATA_EXTENSION_PATTERN
 
-    generateTestGroupSuite(args) {
+    generateTestGroupSuite(args, mainClassName) {
         testGroup("compiler/tests-gen", "compiler/testData") {
             testClass<AbstractDiagnosticsTestWithJsStdLibAndBackendCompilation> {
                 model("diagnostics/testsWithJsStdLibAndBackendCompilation")

compiler/tests-for-compiler-generator/tests/org/jetbrains/kotlin/test/generators/GenerateJUnit5CompilerTests.kt

@@ -17,10 +17,10 @@ import org.jetbrains.kotlin.test.utils.CUSTOM_TEST_DATA_EXTENSION_PATTERN
 import org.jetbrains.kotlin.visualizer.fir.AbstractFirVisualizerTest
 import org.jetbrains.kotlin.visualizer.psi.AbstractPsiVisualizerTest
 
-fun generateJUnit5CompilerTests(args: Array<String>) {
+fun generateJUnit5CompilerTests(args: Array<String>, mainClassName: String?) {
     val excludedCustomTestdataPattern = CUSTOM_TEST_DATA_EXTENSION_PATTERN
 
-    generateTestGroupSuiteWithJUnit5(args) {
+    generateTestGroupSuiteWithJUnit5(args, mainClassName) {
         testGroup(testsRoot = "compiler/tests-common-new/tests-gen", testDataRoot = "compiler/testData") {
             testClass<AbstractDiagnosticTest> {
                 model("diagnostics/tests", pattern = "^(.*)\\.kts?$", excludedPattern = excludedCustomTestdataPattern)

generators/test-generator/tests/org/jetbrains/kotlin/generators/InconsistencyChecker.kt

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.generators
 
+import java.util.Collections
+
 interface InconsistencyChecker {
     fun add(affectedFile: String)
 
@@ -18,7 +20,7 @@ interface InconsistencyChecker {
 }
 
 object DefaultInconsistencyChecker : InconsistencyChecker {
-    private val files = mutableListOf<String>()
+    private val files = Collections.synchronizedList(mutableListOf<String>())
 
     override fun add(affectedFile: String) {
         files.add(affectedFile)

generators/test-generator/tests/org/jetbrains/kotlin/generators/NewTestGenerationDSL.kt

@@ -5,26 +5,32 @@
 
 package org.jetbrains.kotlin.generators
 
+import org.jetbrains.kotlin.generators.util.TestGeneratorUtil
+
 fun generateTestGroupSuiteWithJUnit5(
     args: Array<String>,
+    // Main class name can only be determined when running on the main thread.
+    // If this call is not made on the main thread, the main class name must be injected.
+    mainClassName: String? = TestGeneratorUtil.getMainClassName(),
     additionalMethodGenerators: List<MethodGenerator<Nothing>> = emptyList(),
-    init: TestGroupSuite.() -> Unit
+    init: TestGroupSuite.() -> Unit,
 ) {
-    generateTestGroupSuiteWithJUnit5(InconsistencyChecker.hasDryRunArg(args), additionalMethodGenerators, init)
+    generateTestGroupSuiteWithJUnit5(InconsistencyChecker.hasDryRunArg(args), mainClassName, additionalMethodGenerators, init)
 }
 
 fun generateTestGroupSuiteWithJUnit5(
     dryRun: Boolean = false,
+    // See above
+    mainClassName: String? = TestGeneratorUtil.getMainClassName(),
     additionalMethodGenerators: List<MethodGenerator<Nothing>> = emptyList(),
-    init: TestGroupSuite.() -> Unit
+    init: TestGroupSuite.() -> Unit,
 ) {
     val suite = TestGroupSuite(ReflectionBasedTargetBackendComputer).apply(init)
-    for (testGroup in suite.testGroups) {
+    suite.forEachTestClassParallel { testClass ->
-        for (testClass in testGroup.testClasses) {
+        val (changed, testSourceFilePath) = NewTestGeneratorImpl(additionalMethodGenerators)
-            val (changed, testSourceFilePath) = NewTestGeneratorImpl(additionalMethodGenerators).generateAndSave(testClass, dryRun)
+            .generateAndSave(testClass, dryRun, mainClassName)
-            if (changed) {
+        if (changed) {
-                InconsistencyChecker.inconsistencyChecker(dryRun).add(testSourceFilePath)
+            InconsistencyChecker.inconsistencyChecker(dryRun).add(testSourceFilePath)
-            }
         }
     }
 }

generators/test-generator/tests/org/jetbrains/kotlin/generators/NewTestGeneratorImpl.kt

@@ -75,13 +75,14 @@ class NewTestGeneratorImpl(
         println("@SuppressWarnings(\"all\")")
     }
 
-    override fun generateAndSave(testClass: TestGroup.TestClass, dryRun: Boolean): GenerationResult {
+    override fun generateAndSave(testClass: TestGroup.TestClass, dryRun: Boolean, mainClassName: String?): GenerationResult {
         val generatorInstance = TestGeneratorInstance(
             testClass.baseDir,
             testClass.suiteTestClassName,
             testClass.baseTestClassName,
             testClass.testModels,
-            methodGenerators
+            methodGenerators,
+            mainClassName,
         )
         return generatorInstance.generateAndSave(dryRun)
     }
@@ -91,7 +92,8 @@ class NewTestGeneratorImpl(
         suiteTestClassFqName: String,
         baseTestClassFqName: String,
         private val testClassModels: Collection<TestClassModel>,
-        private val methodGenerators: Map<MethodModel.Kind, MethodGenerator<*>>
+        private val methodGenerators: Map<MethodModel.Kind, MethodGenerator<*>>,
+        private val mainClassName: String?
     ) {
         private val baseTestClassPackage: String = baseTestClassFqName.substringBeforeLast('.', "")
         private val baseTestClassName: String = baseTestClassFqName.substringAfterLast('.', baseTestClassFqName)
@@ -154,7 +156,7 @@ class NewTestGeneratorImpl(
             p.println("import java.io.File;")
             p.println("import java.util.regex.Pattern;")
             p.println()
-            p.println("/** This class is generated by {@link ", getMainClassName(), "}. DO NOT MODIFY MANUALLY */")
+            p.println("/** This class is generated by {@link ", mainClassName, "}. DO NOT MODIFY MANUALLY */")
 
             p.generateSuppressAllWarnings()
 

generators/test-generator/tests/org/jetbrains/kotlin/generators/TestGenerationDSL.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.generators.util.TestGeneratorUtil
 import org.jetbrains.kotlin.generators.util.extractTagsFromDirectory
 import org.jetbrains.kotlin.test.TargetBackend
 import java.io.File
+import java.util.concurrent.ForkJoinPool
 import java.util.regex.Pattern
 import kotlin.reflect.KClass
 
@@ -19,6 +20,14 @@ fun testGroupSuite(
     return TestGroupSuite(DefaultTargetBackendComputer).apply(init)
 }
 
+fun TestGroupSuite.forEachTestClassParallel(f: (TestGroup.TestClass) -> Unit) {
+    testGroups
+        .parallelStream()
+        .flatMap { it.testClasses.stream() }
+        .sorted(compareByDescending { it.testModels.sumOf { it.methods.size } })
+        .forEach(f)
+}
+
 class TestGroupSuite(val targetBackendComputer: TargetBackendComputer) {
     private val _testGroups = mutableListOf<TestGroup>()
     val testGroups: List<TestGroup>

generators/test-generator/tests/org/jetbrains/kotlin/generators/TestGenerator.kt

@@ -13,7 +13,7 @@ abstract class TestGenerator(
     protected val methodGenerators: Map<MethodModel.Kind, MethodGenerator<*>> =
         methodGenerators.associateBy { it.kind }.withDefault { error("Generator for method with kind $it not found") }
 
-    abstract fun generateAndSave(testClass: TestGroup.TestClass, dryRun: Boolean): GenerationResult
+    abstract fun generateAndSave(testClass: TestGroup.TestClass, dryRun: Boolean, mainClassName: String?): GenerationResult
 
     data class GenerationResult(val newFileGenerated: Boolean, val testSourceFilePath: String)
 }

generators/test-generator/tests/org/jetbrains/kotlin/generators/util/TestGeneratorUtil.kt

@@ -36,6 +36,7 @@ object TestGeneratorUtil {
         return escapeForJavaIdentifier(file.name).replaceFirstChar(Char::uppercaseChar)
     }
 
+    /** Must be called on the main thread, otherwise returns the root class of the worker thread. */
     fun getMainClassName(): String? =
         Throwable().stackTrace.lastOrNull()?.className
 }

libraries/tools/kotlin-project-model/src/testFixtures/kotlin/org/jetbrains/kotlin/project/model/infra/generate/generateTestMethods.kt

@@ -7,24 +7,24 @@ package org.jetbrains.kotlin.project.model.infra.generate
 
 import org.jetbrains.kotlin.generators.*
 import org.jetbrains.kotlin.generators.model.*
+import org.jetbrains.kotlin.generators.util.TestGeneratorUtil
 import org.jetbrains.kotlin.project.model.infra.KpmCoreCasesTestRunner
 import java.io.File
 
 fun generateKpmTestCases(
     dryRun: Boolean = false,
-    init: TestGroupSuite.() -> Unit
+    init: TestGroupSuite.() -> Unit,
 ) {
     val suite = TestGroupSuite(DefaultTargetBackendComputer).apply {
         init()
     }
-    for (testGroup in suite.testGroups) {
+    val mainClassName = TestGeneratorUtil.getMainClassName()
-        for (testClass in testGroup.testClasses) {
+    suite.forEachTestClassParallel { testClass ->
-            val (changed, testSourceFilePath) = NewTestGeneratorImpl(
+        val (changed, testSourceFilePath) = NewTestGeneratorImpl(
-                listOf(KpmCoreCaseTestMethodGenerator)
+            listOf(KpmCoreCaseTestMethodGenerator)
-            ).generateAndSave(testClass, dryRun)
+        ).generateAndSave(testClass, dryRun, mainClassName)
-            if (changed) {
+        if (changed) {
-                InconsistencyChecker.inconsistencyChecker(dryRun).add(testSourceFilePath)
+            InconsistencyChecker.inconsistencyChecker(dryRun).add(testSourceFilePath)
-            }
         }
     }
 }