Refactor FIR total resolve in IDE test fixing various problems

Related to KT-24077

idea/tests/org/jetbrains/kotlin/idea/fir/FirTotalKotlinResolveInIdeTest.kt

@@ -5,12 +5,17 @@
 
 package org.jetbrains.kotlin.idea.fir
 
-import com.intellij.openapi.roots.ModuleRootModificationUtil
-import com.intellij.openapi.vfs.VfsUtil
+import com.intellij.ide.impl.ProjectUtil
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.ex.ApplicationEx
+import com.intellij.openapi.fileTypes.FileTypeManager
+import com.intellij.openapi.fileTypes.impl.FileTypeManagerImpl
 import com.intellij.psi.PsiManager
 import com.intellij.psi.search.FileTypeIndex
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.testFramework.LightPlatformTestCase
 import com.intellij.testFramework.ModuleTestCase
+import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
@@ -22,72 +27,79 @@ import org.jetbrains.kotlin.fir.resolve.impl.FirProviderImpl
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTotalResolveTransformer
 import org.jetbrains.kotlin.fir.service
 import org.jetbrains.kotlin.idea.KotlinFileType
-import org.jetbrains.kotlin.idea.caches.project.productionSourceInfo
+import org.jetbrains.kotlin.idea.caches.project.*
 import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
+import kotlin.math.ceil
 
 class FirTotalKotlinResolveInIdeTest : ModuleTestCase() {
-    private val forbiddenDirectories = listOf("testdata", "resources")
-
-    private val projectRootFile = File(".")
-
-    override fun setUpModule() {
-        super.setUpModule()
-
-        ModuleRootModificationUtil.addContentRoot(module, VfsUtil.findFileByIoFile(projectRootFile, true)!!)
-
-        ModuleRootModificationUtil.updateModel(module) {
-            projectRootFile.walkTopDown().onEnter {
-                it.name.toLowerCase() !in forbiddenDirectories
-            }.filter {
-                it.isDirectory
-            }.forEach { dir ->
-                val vdir by lazy { VfsUtil.findFileByIoFile(dir, true)!! }
-                if (dir.name in setOf("src", "test", "tests")) {
-                    it.contentEntries.single().addSourceFolder(vdir, false)
-                } else if (dir.name in setOf("build")) {
-                    it.contentEntries.single().addExcludeFolder(vdir)
-                }
-            }
-        }
+    private val projectRootFile = File("../testProject")
 
+    override fun setUpProject() {
+        (ApplicationManager.getApplication() as ApplicationEx).doNotSave()
+        myProject = ProjectUtil.openOrImport(projectRootFile.path, null, false)
+        LightPlatformTestCase.clearUncommittedDocuments(this.project)
+        this.runStartupActivities()
+        (FileTypeManager.getInstance() as FileTypeManagerImpl).drainReDetectQueue()
     }
 
-    private fun createSession(): FirSession {
-        val moduleInfo = module.productionSourceInfo()!!
-        val sessionProvider = FirProjectSessionProvider(project)
+    private lateinit var sessionProvider: FirProjectSessionProvider
+
+    private fun IdeaModuleInfo.createSession(): FirSession {
+        val moduleInfo = this
 
         return FirJavaModuleBasedSession(
             moduleInfo, sessionProvider, moduleInfo.contentScope(),
-            IdeFirDependenciesSymbolProvider(moduleInfo, project, sessionProvider)
+            IdeFirDependenciesSymbolProvider(moduleInfo as ModuleSourceInfo, project, sessionProvider)
         )
     }
 
+    private fun <T> Collection<T>.progress(label: String, step: Double = 0.1): Sequence<T> {
+        val intStep = (this.size * step).toInt()
+        val percentStep = ceil(this.size * 0.01).toInt()
+        var progress = 0
+        return asSequence().onEach {
+            if (progress % intStep == 0) {
+                println("$label: ${progress / percentStep}% ($progress/${this.size})")
+            }
+            progress++
+        }
+    }
+
+    override fun setUp() {
+        super.setUp()
+        sessionProvider = FirProjectSessionProvider(project)
+    }
+
     fun testTotalKotlin() {
-        val session = createSession()
-
-        val firFiles = mutableListOf<FirFile>()
-        val builder = RawFirBuilder(session, stubMode = true)
         val psiManager = PsiManager.getInstance(project)
-
-        val files = FileTypeIndex.getFiles(KotlinFileType.INSTANCE, GlobalSearchScope.moduleScope(module))
-
+        val files = FileTypeIndex.getFiles(KotlinFileType.INSTANCE, GlobalSearchScope.projectScope(project))
+        val firFiles = mutableListOf<FirFile>()
+        val sessionPerModule = mutableMapOf<ModuleInfo, FirSession>()
 
         println("Got vfiles: ${files.size}")
-        files.forEach {
-            val file = psiManager.findFile(it) as? KtFile ?: return@forEach
+        files.mapNotNull {
+            val file = psiManager.findFile(it) as? KtFile ?: return@mapNotNull null
+            val moduleInfo = file.getNullableModuleInfo() as? ModuleSourceInfo ?: return@mapNotNull null
+            file to moduleInfo
+        }.progress("Loading FIR").forEach { (file, moduleInfo) ->
+
+            val session = sessionPerModule.getOrPut(moduleInfo) {
+                moduleInfo.createSession()
+            }
+            val builder = RawFirBuilder(session, stubMode = true)
+
             try {
                 val firFile = builder.buildFirFile(file)
                 (session.service<FirProvider>() as FirProviderImpl).recordFile(firFile)
                 firFiles += firFile
             } catch (e: Exception) {
-                System.err.println("Error building fir for $it")
+                System.err.println("Error building fir for $file")
                 e.printStackTrace()
             }
         }
 
         println("Raw fir up, files: ${firFiles.size}")
-
         doFirResolveTestBench(firFiles, FirTotalResolveTransformer().transformers)
     }
 }