[Analysis API] introduce a helper test function to get carets at multimodal tests

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractMultiModuleResolveCallTest.kt

@@ -10,23 +10,20 @@ import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.calls.KtCallInfo
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.stringRepresentation
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
-import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.analysis.test.framework.utils.executeOnPooledThreadInReadAction
 import org.jetbrains.kotlin.psi.KtCallElement
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.services.assertions
 
 abstract class AbstractMultiModuleResolveCallTest : AbstractAnalysisApiBasedTest() {
     final override fun doTestByModuleStructure(moduleStructure: TestModuleStructure, testServices: TestServices) {
-        val (expression, _) = moduleStructure.modules.flatMap { module ->
+        val (expression, _) =
-            val ktFiles = testServices.ktModuleProvider.getModuleFiles(module).filterIsInstance<KtFile>()
+            testServices.expressionMarkerProvider.getElementsOfTypeAtCarets<KtExpression>(moduleStructure, testServices)
-            testServices.expressionMarkerProvider.getElementsOfTypeAtCarets<KtExpression>(ktFiles)
+                .single()
-        }.single()
 
         val actual = executeOnPooledThreadInReadAction {
             analyseForTest(expression) {

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.test.framework.services
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.psi.KtElement
@@ -15,6 +16,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
 import org.jetbrains.kotlin.test.model.TestFile
 import org.jetbrains.kotlin.test.services.SourceFilePreprocessor
+import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestService
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
@@ -107,6 +109,18 @@ class ExpressionMarkerProvider : TestService {
         }
     }
 
+    inline fun <reified P : KtElement> getElementsOfTypeAtCarets(
+        moduleStructure: TestModuleStructure,
+        testServices: TestServices,
+        caretTag: String? = null
+    ): Collection<Pair<P, KtFile>> {
+        return moduleStructure.modules.flatMap { module ->
+            val ktFiles = testServices.ktModuleProvider.getModuleFiles(module).filterIsInstance<KtFile>()
+            getElementsOfTypeAtCarets<P>(ktFiles, caretTag)
+        }
+
+    }
+
     fun getSelectedElement(file: KtFile): KtElement {
         val range = selected[file.name]
             ?: error("No selected expression found in file")