[Test] Parallelize test generation

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) {
-        for (testClass in testGroup.testClasses) {
-            val (changed, testSourceFilePath) = NewTestGeneratorImpl(additionalMethodGenerators).generateAndSave(testClass, dryRun)
-            if (changed) {
-                InconsistencyChecker.inconsistencyChecker(dryRun).add(testSourceFilePath)
-            }
+    suite.forEachTestClassParallel { testClass ->
+        val (changed, testSourceFilePath) = NewTestGeneratorImpl(additionalMethodGenerators)
+            .generateAndSave(testClass, dryRun, mainClassName)
+        if (changed) {
+            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
 }