K2: update LightTree support in modularized tests

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/pipeline/compilerPipeline.kt

@@ -83,6 +83,7 @@ import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import java.io.File
+import java.util.LinkedHashSet
 
 private const val kotlinFileExtensionWithDot = ".${KotlinFileType.EXTENSION}"
 private const val javaFileExtensionWithDot = ".${JavaFileType.DEFAULT_EXTENSION}"
@@ -114,41 +115,7 @@ fun compileModulesUsingFrontendIrAndLightTree(
         val moduleConfiguration = compilerConfiguration.copy().applyModuleProperties(module, buildFile).apply {
             put(JVMConfigurationKeys.FRIEND_PATHS, module.getFriendPaths())
         }
-        val platformSources = linkedSetOf<KtSourceFile>()
-        val commonSources = linkedSetOf<KtSourceFile>()
-
-        // !!
-
-        // TODO: the scripts checking should be part of the scripting plugin functionality, as it is implemented now in ScriptingProcessSourcesBeforeCompilingExtension
-        // TODO: implement in the next round of K2 scripting support
-        val skipScriptsInLtMode = compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR) && compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE)
-        var skipScriptsInLtModeWarning = false
-
-        compilerConfiguration.kotlinSourceRoots.forAllFiles(compilerConfiguration, projectEnvironment.project) { virtualFile, isCommon ->
-            val file = KtVirtualFileSourceFile(virtualFile)
-            when {
-                file.path.endsWith(javaFileExtensionWithDot) -> {}
-                file.path.endsWith(kotlinFileExtensionWithDot) || !skipScriptsInLtMode -> {
-                    if (isCommon) commonSources.add(file)
-                    else platformSources.add(file)
-                }
-                else -> {
-                    // temporarily assume it is a script, see the TODO above
-                    skipScriptsInLtModeWarning = true
-                }
-            }
-        }
-
-        if (skipScriptsInLtModeWarning) {
-            // TODO: remove then Scripts are supported in LT (probably different K2 extension should be written for handling the case properly)
-            messageCollector.report(
-                CompilerMessageSeverity.STRONG_WARNING,
-                "Scripts are not yet supported with K2 in LightTree mode, consider using K1 or disable LightTree mode with -XuseFirLT=false"
-            )
-        }
-
-        val renderDiagnosticName = moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
-        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
+        val (platformSources, commonSources) = collectSources(compilerConfiguration, projectEnvironment, messageCollector)
 
         val compilerInput = ModuleCompilerInput(
             TargetId(module),
@@ -156,6 +123,9 @@ fun compileModulesUsingFrontendIrAndLightTree(
             JvmPlatforms.unspecifiedJvmPlatform, platformSources,
             moduleConfiguration
         )
+
+        val renderDiagnosticName = moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
+        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
         val compilerEnvironment = ModuleCompilerEnvironment(projectEnvironment, diagnosticsReporter)
 
         performanceManager?.notifyAnalysisStarted()
@@ -210,6 +180,44 @@ fun compileModulesUsingFrontendIrAndLightTree(
     )
 }
 
+fun collectSources(
+    compilerConfiguration: CompilerConfiguration,
+    projectEnvironment: VfsBasedProjectEnvironment,
+    messageCollector: MessageCollector
+): Pair<LinkedHashSet<KtSourceFile>, LinkedHashSet<KtSourceFile>> {
+    val platformSources = linkedSetOf<KtSourceFile>()
+    val commonSources = linkedSetOf<KtSourceFile>()
+
+    // TODO: the scripts checking should be part of the scripting plugin functionality, as it is implemented now in ScriptingProcessSourcesBeforeCompilingExtension
+    // TODO: implement in the next round of K2 scripting support (https://youtrack.jetbrains.com/issue/KT-55728)
+    val skipScriptsInLtMode = compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR) && compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE)
+    var skipScriptsInLtModeWarning = false
+
+    compilerConfiguration.kotlinSourceRoots.forAllFiles(compilerConfiguration, projectEnvironment.project) { virtualFile, isCommon ->
+        val file = KtVirtualFileSourceFile(virtualFile)
+        when {
+            file.path.endsWith(javaFileExtensionWithDot) -> {}
+            file.path.endsWith(kotlinFileExtensionWithDot) || !skipScriptsInLtMode -> {
+                if (isCommon) commonSources.add(file)
+                else platformSources.add(file)
+            }
+            else -> {
+                // temporarily assume it is a script, see the TODO above
+                skipScriptsInLtModeWarning = true
+            }
+        }
+    }
+
+    if (skipScriptsInLtModeWarning) {
+        // TODO: remove then Scripts are supported in LT (probably different K2 extension should be written for handling the case properly)
+        messageCollector.report(
+            CompilerMessageSeverity.STRONG_WARNING,
+            "Scripts are not yet supported with K2 in LightTree mode, consider using K1 or disable LightTree mode with -Xuse-fir-lt=false"
+        )
+    }
+    return Pair(platformSources, commonSources)
+}
+
 fun convertAnalyzedFirToIr(
     input: ModuleCompilerInput,
     analysisResults: ModuleCompilerAnalyzedOutput,

compiler/fir/modularized-tests/tests/org/jetbrains/kotlin/fir/FirResolveModularizedTotalKotlinTest.kt

@@ -12,10 +12,13 @@ import com.intellij.psi.search.ProjectScope
 import com.sun.jna.Library
 import com.sun.jna.Native
 import com.sun.management.HotSpotDiagnosticMXBean
+import org.jetbrains.kotlin.KtPsiSourceFile
+import org.jetbrains.kotlin.KtSourceFile
 import org.jetbrains.kotlin.ObsoleteTestInfrastructure
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.cli.common.toBooleanLenient
 import org.jetbrains.kotlin.cli.jvm.compiler.*
+import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.collectSources
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
@@ -122,20 +125,33 @@ class FirResolveModularizedTotalKotlinTest : AbstractModularizedTest() {
 
     @OptIn(ObsoleteTestInfrastructure::class)
     private fun runAnalysis(moduleData: ModuleData, environment: KotlinCoreEnvironment) {
-        val project = environment.project
-        val ktFiles = environment.getSourceFiles()
 
-        val scope = GlobalSearchScope.filesScope(project, ktFiles.map { it.virtualFile })
-            .uniteWith(TopDownAnalyzerFacadeForJVM.AllJavaSourcesInProjectScope(project))
+        val projectEnvironment = environment.toAbstractProjectEnvironment() as VfsBasedProjectEnvironment
+        val project = environment.project
+
+        val (sourceFiles: Collection<KtSourceFile>, scope) =
+            if (USE_LIGHT_TREE) {
+                val (platformSources, _) = collectSources(environment.configuration, projectEnvironment, environment.messageCollector)
+                platformSources to projectEnvironment.getSearchScopeBySourceFiles(platformSources)
+            } else {
+                val ktFiles = environment.getSourceFiles()
+                ktFiles.map { KtPsiSourceFile(it) } to
+                        GlobalSearchScope.filesScope(project, ktFiles.map { it.virtualFile })
+                            .uniteWith(TopDownAnalyzerFacadeForJVM.AllJavaSourcesInProjectScope(project))
+                            .toAbstractProjectFileSearchScope()
+            }
         val librariesScope = ProjectScope.getLibrariesScope(project)
-        val session = FirTestSessionFactoryHelper.createSessionForTests(
-            environment.toAbstractProjectEnvironment(),
-            scope.toAbstractProjectFileSearchScope(),
-            librariesScope.toAbstractProjectFileSearchScope(),
-            moduleData.qualifiedName,
-            moduleData.friendDirs.map { it.toPath() },
-            environment.configuration.languageVersionSettings
-        )
+
+        val session =
+            FirTestSessionFactoryHelper.createSessionForTests(
+                projectEnvironment,
+                scope,
+                librariesScope.toAbstractProjectFileSearchScope(),
+                moduleData.qualifiedName,
+                moduleData.friendDirs.map { it.toPath() },
+                environment.configuration.languageVersionSettings
+            )
+
         val scopeSession = ScopeSession()
         val processors = createAllCompilerResolveProcessors(session, scopeSession).let {
             if (RUN_CHECKERS) {
@@ -149,20 +165,10 @@ class FirResolveModularizedTotalKotlinTest : AbstractModularizedTest() {
 
         val firFiles = if (USE_LIGHT_TREE) {
             val lightTree2Fir = LightTree2Fir(session, firProvider.kotlinScopeProvider, diagnosticsReporter = null)
-
-            val allSourceFiles = moduleData.sources.flatMap {
-                if (it.isDirectory) {
-                    it.walkTopDown().toList()
-                } else {
-                    listOf(it)
-                }
-            }.filter {
-                it.extension == "kt"
-            }
-            bench.buildFiles(lightTree2Fir, allSourceFiles)
+            bench.buildFiles(lightTree2Fir, sourceFiles)
         } else {
             val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
-            bench.buildFiles(builder, ktFiles)
+            bench.buildFiles(builder, sourceFiles.map { it as KtPsiSourceFile })
         }
 
 

compiler/tests-common/tests/org/jetbrains/kotlin/fir/FirResolveBench.kt

@@ -8,6 +8,8 @@ import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.KtIoFileSourceFile
+import org.jetbrains.kotlin.KtPsiSourceFile
+import org.jetbrains.kotlin.KtSourceFile
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.diagnostics.ConeStubDiagnostic
@@ -113,10 +115,11 @@ class FirResolveBench(val withProgress: Boolean, val listener: BenchListener? =
 
     fun buildFiles(
         builder: RawFirBuilder,
-        ktFiles: List<KtFile>
+        files: Collection<KtPsiSourceFile>
     ): List<FirFile> {
         listener?.before()
-        return ktFiles.map { file ->
+        return files.map { sourceFile ->
+            val file = sourceFile.psiFile as KtFile
             val before = vmStateSnapshot()
             val firFile: FirFile
             val time = measureNanoTime {
@@ -136,19 +139,18 @@ class FirResolveBench(val withProgress: Boolean, val listener: BenchListener? =
 
     fun buildFiles(
         builder: LightTree2Fir,
-        files: List<File>
+        files: Collection<KtSourceFile>
     ): List<FirFile> {
         listener?.before()
         return files.map { file ->
             val before = vmStateSnapshot()
             val firFile: FirFile
             val time = measureNanoTime {
-                val sourceFile = KtIoFileSourceFile(file)
-                val (code, linesMapping) = with(file.inputStream().reader(Charsets.UTF_8)) {
+                val (code, linesMapping) = with(file.getContentsAsStream().reader(Charsets.UTF_8)) {
                     this.readSourceFileWithMapping()
                 }
                 totalLines += linesMapping.linesCount
-                firFile = builder.buildFirFile(code, sourceFile, linesMapping)
+                firFile = builder.buildFirFile(code, file, linesMapping)
                 (builder.session.firProvider as FirProviderImpl).recordFile(firFile)
             }
             val after = vmStateSnapshot()