[Test] Add ability to freely transform module structure which was parsed from directive

Service ModuleStructureExtractor may break readabilty of test so it
  should be used very carefully

Originally it is created for box inline test where we have different
  runners on same testdata where one runner assumes that all files
  in test in same module and other that they in different ones

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/TestInfrastructureInternals.kt

@@ -0,0 +1,14 @@
+/*
+ * Copyright 2010-2021 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.test
+
+/**
+ * This annotation used for marking low-level services of test infrastructure which
+ *   normally should not be used. Please think twice before using something
+ *   marked with this annotation
+ */
+@RequiresOptIn
+annotation class TestInfrastructureInternals

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/ModuleStructureExtractor.kt

@@ -5,11 +5,13 @@
 
 package org.jetbrains.kotlin.test.services
 
+import org.jetbrains.kotlin.test.TestInfrastructureInternals
 import org.jetbrains.kotlin.test.directives.model.DirectivesContainer
 
-abstract class ModuleStructureExtractor(
+abstract class ModuleStructureExtractor @OptIn(TestInfrastructureInternals::class) constructor(
     protected val testServices: TestServices,
-    protected val additionalSourceProviders: List<AdditionalSourceProvider>
+    protected val additionalSourceProviders: List<AdditionalSourceProvider>,
+    protected val moduleStructureTransformers: List<ModuleStructureTransformer>
 ) {
     abstract fun splitTestDataByModules(
         testDataFileName: String,

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/ModuleStructureTransformer.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2021 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.test.services
+
+import org.jetbrains.kotlin.test.TestInfrastructureInternals
+
+@TestInfrastructureInternals
+abstract class ModuleStructureTransformer {
+    abstract fun transformModuleStructure(moduleStructure: TestModuleStructure): TestModuleStructure
+}

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/builders/TestConfigurationBuilder.kt

@@ -7,12 +7,14 @@ package org.jetbrains.kotlin.test.builders
 
 import org.jetbrains.kotlin.test.Constructor
 import org.jetbrains.kotlin.test.TestConfiguration
+import org.jetbrains.kotlin.test.TestInfrastructureInternals
 import org.jetbrains.kotlin.test.directives.model.DirectivesContainer
 import org.jetbrains.kotlin.test.impl.TestConfigurationImpl
 import org.jetbrains.kotlin.test.model.*
 import org.jetbrains.kotlin.test.services.*
 
 @DefaultsDsl
+@OptIn(TestInfrastructureInternals::class)
 class TestConfigurationBuilder {
     val defaultsProviderBuilder: DefaultsProviderBuilder = DefaultsProviderBuilder()
     lateinit var assertions: AssertionsService
@@ -26,6 +28,7 @@ class TestConfigurationBuilder {
     private val environmentConfigurators: MutableList<Constructor<EnvironmentConfigurator>> = mutableListOf()
 
     private val additionalSourceProviders: MutableList<Constructor<AdditionalSourceProvider>> = mutableListOf()
+    private val moduleStructureTransformers: MutableList<ModuleStructureTransformer> = mutableListOf()
 
     private val metaTestConfigurators: MutableList<Constructor<MetaTestConfigurator>> = mutableListOf()
     private val afterAnalysisCheckers: MutableList<Constructor<AfterAnalysisChecker>> = mutableListOf()
@@ -115,6 +118,11 @@ class TestConfigurationBuilder {
         additionalSourceProviders += providers
     }
 
+    @TestInfrastructureInternals
+    fun useModuleStructureTransformers(vararg transformers: ModuleStructureTransformer) {
+        moduleStructureTransformers += transformers
+    }
+
     fun useMetaTestConfigurators(vararg configurators: Constructor<MetaTestConfigurator>) {
         metaTestConfigurators += configurators
     }
@@ -147,6 +155,7 @@ class TestConfigurationBuilder {
             additionalMetaInfoProcessors,
             environmentConfigurators,
             additionalSourceProviders,
+            moduleStructureTransformers,
             metaTestConfigurators,
             afterAnalysisCheckers,
             metaInfoHandlerEnabled,

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/impl/TestConfigurationImpl.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.test.impl
 import com.intellij.openapi.Disposable
 import org.jetbrains.kotlin.test.Constructor
 import org.jetbrains.kotlin.test.TestConfiguration
+import org.jetbrains.kotlin.test.TestInfrastructureInternals
 import org.jetbrains.kotlin.test.directives.model.ComposedDirectivesContainer
 import org.jetbrains.kotlin.test.directives.model.DirectivesContainer
 import org.jetbrains.kotlin.test.directives.model.RegisteredDirectives
@@ -16,6 +17,7 @@ import org.jetbrains.kotlin.test.services.*
 import org.jetbrains.kotlin.test.services.impl.ModuleStructureExtractorImpl
 import org.jetbrains.kotlin.test.utils.TestDisposable
 
+@OptIn(TestInfrastructureInternals::class)
 class TestConfigurationImpl(
     testInfo: KotlinTestInfo,
 
@@ -31,6 +33,7 @@ class TestConfigurationImpl(
     environmentConfigurators: List<Constructor<EnvironmentConfigurator>>,
 
     additionalSourceProviders: List<Constructor<AdditionalSourceProvider>>,
+    moduleStructureTransformers: List<ModuleStructureTransformer>,
     metaTestConfigurators: List<Constructor<MetaTestConfigurator>>,
     afterAnalysisCheckers: List<Constructor<AfterAnalysisChecker>>,
 
@@ -70,6 +73,7 @@ class TestConfigurationImpl(
         additionalSourceProviders.map { it.invoke(testServices) }.also {
             it.flatMapTo(allDirectives) { provider -> provider.directives }
         },
+        moduleStructureTransformers,
         this.environmentConfigurators
     )
 

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/services/impl/ModuleStructureExtractorImpl.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.platform.konan.NativePlatforms
 import org.jetbrains.kotlin.test.Assertions
 import org.jetbrains.kotlin.test.TargetBackend
+import org.jetbrains.kotlin.test.TestInfrastructureInternals
 import org.jetbrains.kotlin.test.builders.LanguageVersionSettingsBuilder
 import org.jetbrains.kotlin.test.directives.AdditionalFilesDirectives
 import org.jetbrains.kotlin.test.directives.ModuleStructureDirectives
@@ -35,11 +36,13 @@ import java.io.File
  * - All directives between `MODULE` and `FILE` directives belongs to module
  * - All directives before first `MODULE` are global and belongs to each declared module
  */
+@OptIn(TestInfrastructureInternals::class)
 class ModuleStructureExtractorImpl(
     testServices: TestServices,
     additionalSourceProviders: List<AdditionalSourceProvider>,
+    moduleStructureTransformers: List<ModuleStructureTransformer>,
     private val environmentConfigurators: List<EnvironmentConfigurator>
-) : ModuleStructureExtractor(testServices, additionalSourceProviders) {
+) : ModuleStructureExtractor(testServices, additionalSourceProviders, moduleStructureTransformers) {
     companion object {
         private val allowedExtensionsForFiles = listOf(".kt", ".kts", ".java")
 
@@ -57,7 +60,11 @@ class ModuleStructureExtractorImpl(
     ): TestModuleStructure {
         val testDataFile = File(testDataFileName)
         val extractor = ModuleStructureExtractorWorker(listOf(testDataFile), directivesContainer)
-        return extractor.splitTestDataByModules()
+        var result = extractor.splitTestDataByModules()
+        for (transformer in moduleStructureTransformers) {
+            result = transformer.transformModuleStructure(result)
+        }
+        return result
     }
 
     private inner class ModuleStructureExtractorWorker constructor(