From ee1b018c9a3611fdc0ccc3d80668202aa0a307be Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Fri, 17 Feb 2023 14:32:44 +0100
Subject: [PATCH] [Analysis API] introduce a helper test function to get carets
 at multimodal tests

---
 .../AbstractMultiModuleResolveCallTest.kt          |  9 +++------
 .../ExpressionMarkersSourceFilePreprocessor.kt     | 14 ++++++++++++++
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractMultiModuleResolveCallTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractMultiModuleResolveCallTest.kt
index d71deeb8b10..2511f77e02f 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/callResolver/AbstractMultiModuleResolveCallTest.kt
+++ b/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 ktFiles = testServices.ktModuleProvider.getModuleFiles(module).filterIsInstance<KtFile>()
-            testServices.expressionMarkerProvider.getElementsOfTypeAtCarets<KtExpression>(ktFiles)
-        }.single()
+        val (expression, _) =
+            testServices.expressionMarkerProvider.getElementsOfTypeAtCarets<KtExpression>(moduleStructure, testServices)
+                .single()
 
         val actual = executeOnPooledThreadInReadAction {
             analyseForTest(expression) {
diff --git a/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt b/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt
index 3d1f61419d0..1566cc50936 100644
--- a/analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/ExpressionMarkersSourceFilePreprocessor.kt
+++ b/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")