[Analysis API] introduce getKtFiles util

^KT-64805

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/AnalysisApiKtModuleProviderImpl.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 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.
  */
 
@@ -9,9 +9,11 @@ import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.KtModuleProjectStructure
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.KtModuleWithFiles
 import org.jetbrains.kotlin.analysis.project.structure.*
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.model.TestModule
 import org.jetbrains.kotlin.test.services.TestService
 import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.moduleStructure
 
 abstract class AnalysisApiKtModuleProvider : TestService {
     protected abstract val testServices: TestServices
@@ -48,6 +50,8 @@ class AnalysisApiKtModuleProviderImpl(
     override fun getModuleStructure(): KtModuleProjectStructure = modulesStructure
 }
 
+fun AnalysisApiKtModuleProvider.getKtFiles(module: TestModule): List<KtFile> = getModuleFiles(module).filterIsInstance<KtFile>()
+
 fun TestServices.allKtFiles(): List<KtFile> = moduleStructure.modules.flatMap(ktModuleProvider::getKtFiles)
 
 val TestServices.ktModuleProvider: AnalysisApiKtModuleProvider by TestServices.testServiceAccessor()
@@ -66,4 +70,4 @@ fun List<KtModuleWithFiles>.associateByName(): Map<String, KtModuleWithFiles> {
             else -> error("Unsupported module type: " + ktModule.javaClass.name)
         }
     }
-}
+}

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt

@@ -10,6 +10,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.getKtFiles
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
 import org.jetbrains.kotlin.psi.KtElement
@@ -123,7 +124,7 @@ class ExpressionMarkerProvider : TestService {
     @OptIn(PrivateForInline::class)
     inline fun <reified P : KtElement> getElementsOfTypeAtCarets(
         files: Collection<KtFile>,
-        caretTag: String? = null
+        caretTag: String? = null,
     ): Collection<Pair<P, KtFile>> {
         return files.mapNotNull { file ->
             getCaretPositionOrNull(file, caretTag)?.let { offset ->
@@ -137,10 +138,10 @@ class ExpressionMarkerProvider : TestService {
     inline fun <reified P : KtElement> getElementsOfTypeAtCarets(
         moduleStructure: TestModuleStructure,
         testServices: TestServices,
-        caretTag: String? = null
+        caretTag: String? = null,
     ): Collection<Pair<P, KtFile>> {
         return moduleStructure.modules.flatMap { module ->
-            val ktFiles = testServices.ktModuleProvider.getModuleFiles(module).filterIsInstance<KtFile>()
+            val ktFiles = testServices.ktModuleProvider.getKtFiles(module)
             getElementsOfTypeAtCarets<P>(ktFiles, caretTag)
         }
 

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/SealedClassesInheritorsCaclulatorPreAnalysisHandler.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 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.
  */
 
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.FirSealedClassInheritorsProcessorFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.services.LLFirSealedClassInheritorsProcessorFactoryForTests
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.getKtFiles
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
@@ -42,7 +43,7 @@ class SealedClassesInheritorsCaclulatorPreAnalysisHandler(
         }
 
         val ktFilesByModule = moduleStructure.modules.associateWith { testModule ->
-            testServices.ktModuleProvider.getModuleFiles(testModule).filterIsInstance<KtFile>()
+            testServices.ktModuleProvider.getKtFiles(testModule)
         }
 
         for ((testModule, ktFiles) in ktFilesByModule) {
@@ -62,7 +63,7 @@ class SealedClassesInheritorsCaclulatorPreAnalysisHandler(
 
     private fun collectSealedClassInheritors(
         firFiles: List<FirFile>,
-        tmpFirResolveSession: LLFirResolveSession
+        tmpFirResolveSession: LLFirResolveSession,
     ): Map<ClassId, List<ClassId>> {
         firFiles.forEach { it.lazyResolveToPhaseRecursively(FirResolvePhase.TYPES) }
         val inheritorsCollector = FirSealedClassInheritorsProcessor.InheritorsCollector(tmpFirResolveSession.useSiteFirSession)