FIR IDE: introduce KtAnalysisSessionProvider & helping functions for working with analysis session

idea/idea-fir/src/org/jetbrains/kotlin/idea/fir/highlighter/KotlinFirPsiChecker.kt

@@ -12,6 +12,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.fir.highlighter.visitors.FirAfterResolveHighlightingVisitor
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.analyze
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.highlighter.AbstractKotlinPsiChecker
 import org.jetbrains.kotlin.idea.highlighter.Diagnostic2Annotation
@@ -29,13 +30,13 @@ class KotlinFirPsiChecker : AbstractKotlinPsiChecker() {
         if (ApplicationManager.getApplication().isDispatchThread) {
             throw ProcessCanceledException()
         }
-        val analysisSession = KtFirAnalysisSession(element)
+        analyze(element) {
+            highlightDiagnostics(element, this, holder)
 
-        highlightDiagnostics(element, analysisSession, holder)
+            FirAfterResolveHighlightingVisitor
-
+                .createListOfVisitors(this, holder)
-        FirAfterResolveHighlightingVisitor
+                .forEach(element::accept)
-            .createListOfVisitors(analysisSession, holder)
+        }
-            .forEach(element::accept)
     }
 
     private fun highlightDiagnostics(element: KtElement, analysisSession: KtAnalysisSession, holder: AnnotationHolder) {

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2020 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.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api
+
+import com.intellij.openapi.components.service
+import org.jetbrains.kotlin.idea.util.application.runReadAction
+import org.jetbrains.kotlin.psi.KtElement
+
+abstract class KtAnalysisSessionProvider {
+    abstract fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession
+}
+
+fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession =
+    service<KtAnalysisSessionProvider>().getAnalysisSessionFor(contextElement)
+
+inline fun <R> analyze(contextElement: KtElement, action: KtAnalysisSession.() -> R): R =
+    getAnalysisSessionFor(contextElement).action()
+
+inline fun <R> analyzeWithReadAction(
+    contextElement: KtElement,
+    crossinline action: KtAnalysisSession.() -> R
+): R = runReadAction {
+    analyze(contextElement, action)
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSession.kt

@@ -33,7 +33,9 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 
-class KtFirAnalysisSession(
+class KtFirAnalysisSession
+@Deprecated("Please use org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProviderKt.analyze")
+constructor(
     private val element: KtElement
 ) : KtAnalysisSession(element.project) {
     internal val token get() = validityToken

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSessionProvider.kt

@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2020 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.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api.fir
+
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProvider
+import org.jetbrains.kotlin.psi.KtElement
+
+class KtFirAnalysisSessionProvider : KtAnalysisSessionProvider() {
+    @Suppress("DEPRECATION")
+    override fun getAnalysisSessionFor(contextElement: KtElement): KtAnalysisSession =
+        KtFirAnalysisSession(contextElement)
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt

@@ -11,7 +11,9 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementResolveResult
 import com.intellij.psi.ResolveResult
 import com.intellij.psi.impl.source.resolve.ResolveCache
+import org.jetbrains.kotlin.idea.frontend.api.analyze
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.getAnalysisSessionFor
 
 object KtFirReferenceResolver : ResolveCache.PolyVariantResolver<KtReference> {
     class KotlinResolveResult(element: PsiElement) : PsiElementResolveResult(element)
@@ -22,8 +24,7 @@ object KtFirReferenceResolver : ResolveCache.PolyVariantResolver<KtReference> {
         if (ApplicationManager.getApplication().isDispatchThread) {
             throw ProcessCanceledException()
         }
-        val analysisSession = KtFirAnalysisSession(ref.expression)
+        val resolveToPsiElements = analyze(ref.expression) { ref.getResolvedToPsi(this) }
-        val resolveToPsiElements = ref.getResolvedToPsi(analysisSession)
         return resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
     }
 }

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/fir/AbstractResolveCallTest.kt

@@ -11,6 +11,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.addExternalTestFiles
 import org.jetbrains.kotlin.idea.executeOnPooledThreadInReadAction
 import org.jetbrains.kotlin.idea.frontend.api.CallInfo
+import org.jetbrains.kotlin.idea.frontend.api.analyze
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionLikeSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol
@@ -39,12 +40,13 @@ abstract class AbstractResolveCallTest : @Suppress("DEPRECATION") KotlinLightCod
         }
 
         val actualText = executeOnPooledThreadInReadAction {
-            val analysisSession = KtFirAnalysisSession(file as KtFile)
+            val callInfos = analyze(file as KtFile) {
-            val callInfos = elements.map { element ->
+                elements.map { element ->
-                when (element) {
+                    when (element) {
-                    is KtCallExpression -> analysisSession.resolveCall(element)
+                        is KtCallExpression -> resolveCall(element)
-                    is KtBinaryExpression -> analysisSession.resolveCall(element)
+                        is KtBinaryExpression -> resolveCall(element)
-                    else -> error("Selected should be either KtCallExpression or KtBinaryExpression but was $element")
+                        else -> error("Selected should be either KtCallExpression or KtBinaryExpression but was $element")
+                    }
                 }
             }
 

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/symbols/AbstractStdlibSymbolsBuildingTest.kt

@@ -35,9 +35,10 @@ abstract class AbstractStdlibSymbolsBuildingTest : KotlinLightCodeInsightFixture
         val symbolData = SymbolData.create(identifier)
 
         val renderedSymbols = executeOnPooledThreadInReadAction {
-            val analysisSession = KtFirAnalysisSession(fakeKtFile)
+            analyze(fakeKtFile) {
-            val symbols = symbolData.toSymbols(analysisSession)
+                val symbols = symbolData.toSymbols(this)
-            symbols.map { DebugSymbolRenderer.render(it) }
+                symbols.map { DebugSymbolRenderer.render(it) }
+            }
         }
 
         val actual = buildString {
@@ -49,6 +50,7 @@ abstract class AbstractStdlibSymbolsBuildingTest : KotlinLightCodeInsightFixture
             append(actualSymbolsData)
         }
         KotlinTestUtils.assertEqualsToFile(testDataFile, actual)
+    }
 
     }
 

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/symbols/AbstractSymbolPointerTest.kt

@@ -10,7 +10,9 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.addExternalTestFiles
 import org.jetbrains.kotlin.idea.executeOnPooledThreadInReadAction
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.analyze
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.getAnalysisSessionFor
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
@@ -31,19 +33,21 @@ abstract class AbstractSymbolPointerTest : KotlinLightCodeInsightFixtureTestCase
         val ktFile = myFixture.configureByText(file.name, FileUtil.loadFile(file)) as KtFile
 
         val pointers = executeOnPooledThreadInReadAction {
-            val analysisSession = KtFirAnalysisSession(ktFile)
+            analyze(ktFile) {
-            val declarationSymbols = ktFile.collectDescendantsOfType<KtDeclaration>().map { declaration ->
+                val declarationSymbols = ktFile.collectDescendantsOfType<KtDeclaration>().map { declaration ->
-                analysisSession.symbolProvider.getSymbol(declaration)
+                    symbolProvider.getSymbol(declaration)
+                }
+                declarationSymbols.map { PointerWithPsi(it.createPointer(), it.psi!!) }
             }
-            declarationSymbols.map { PointerWithPsi(it.createPointer(), it.psi!!) }
         }
 
         // another read action
         executeOnPooledThreadInReadAction {
-            val analysisSession = KtFirAnalysisSession(ktFile)
+            analyze(ktFile) {
-            pointers.forEach { (pointer, psi) ->
+                pointers.forEach { (pointer, psi) ->
-                val restored = pointer.restoreSymbol(analysisSession) ?: error("Symbol $psi was not not restored correctly")
+                    val restored = pointer.restoreSymbol(this) ?: error("Symbol $psi was not not restored correctly")
-                Assert.assertEquals(restored.psi!!, psi)
+                    Assert.assertEquals(restored.psi!!, psi)
+                }
             }
         }
     }

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/symbols/AbstractSymbolsByPsiBuildingTest.kt

@@ -8,7 +8,9 @@ package org.jetbrains.kotlin.idea.frontend.api.symbols
 import com.intellij.openapi.util.io.FileUtil
 import org.jetbrains.kotlin.idea.addExternalTestFiles
 import org.jetbrains.kotlin.idea.executeOnPooledThreadInReadAction
+import org.jetbrains.kotlin.idea.frontend.api.analyze
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.getAnalysisSessionFor
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFile
@@ -23,11 +25,12 @@ abstract class AbstractSymbolsByPsiBuildingTest : KotlinLightCodeInsightFixtureT
         val ktFile = myFixture.configureByText(file.name, FileUtil.loadFile(file)) as KtFile
 
         val renderedSymbols = executeOnPooledThreadInReadAction {
-            val analysisSession = KtFirAnalysisSession(ktFile)
+            analyze(ktFile) {
-            val declarationSymbols = ktFile.collectDescendantsOfType<KtDeclaration>().map { declaration ->
+                val declarationSymbols = ktFile.collectDescendantsOfType<KtDeclaration>().map { declaration ->
-                analysisSession.symbolProvider.getSymbol(declaration)
+                    symbolProvider.getSymbol(declaration)
+                }
+                declarationSymbols.map(DebugSymbolRenderer::render)
             }
-            declarationSymbols.map(DebugSymbolRenderer::render)
         }
 
         val actual = buildString {

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/symbols/StdlibSymbolsBuildingTestGenerated.java

@@ -16,7 +16,7 @@ import java.util.regex.Pattern;
 
 /** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
 @SuppressWarnings("all")
-@TestMetadata("idea/idea-frontend-fir/testData/stdLibSymbols")
+@TestMetadata("idea/idea-frontend-fir/testData/symbolsByFqName")
 @TestDataPath("$PROJECT_ROOT")
 @RunWith(JUnit3RunnerWithInners.class)
 public class StdlibSymbolsBuildingTestGenerated extends AbstractStdlibSymbolsBuildingTest {
@@ -24,22 +24,22 @@ public class StdlibSymbolsBuildingTestGenerated extends AbstractStdlibSymbolsBui
         KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
     }
 
-    public void testAllFilesPresentInStdLibSymbols() throws Exception {
+    public void testAllFilesPresentInSymbolsByFqName() throws Exception {
-        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/stdLibSymbols"), Pattern.compile("^(.+)\\.txt$"), null, true);
+        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/symbolsByFqName"), Pattern.compile("^(.+)\\.txt$"), null, true);
     }
 
     @TestMetadata("fileWalkDirectionEnum.txt")
     public void testFileWalkDirectionEnum() throws Exception {
-        runTest("idea/idea-frontend-fir/testData/stdLibSymbols/fileWalkDirectionEnum.txt");
+        runTest("idea/idea-frontend-fir/testData/symbolsByFqName/fileWalkDirectionEnum.txt");
     }
 
     @TestMetadata("iterator.txt")
     public void testIterator() throws Exception {
-        runTest("idea/idea-frontend-fir/testData/stdLibSymbols/iterator.txt");
+        runTest("idea/idea-frontend-fir/testData/symbolsByFqName/iterator.txt");
     }
 
     @TestMetadata("listOf.txt")
     public void testListOf() throws Exception {
-        runTest("idea/idea-frontend-fir/testData/stdLibSymbols/listOf.txt");
+        runTest("idea/idea-frontend-fir/testData/symbolsByFqName/listOf.txt");
     }
 }

idea/resources-fir/META-INF/plugin.xml

@@ -81,6 +81,8 @@ The Kotlin FIR plugin provides language support in IntelliJ IDEA and Android Stu
 
   <extensions defaultExtensionNs="com.intellij">
     <!-- fir-specific -->
+    <applicationService serviceInterface="org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProvider"
+                        serviceImplementation="org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSessionProvider"/>
     <applicationService serviceInterface="org.jetbrains.kotlin.idea.references.KotlinReferenceProviderContributor"
                         serviceImplementation="org.jetbrains.kotlin.idea.references.KotlinFirReferenceContributor"/>
     <applicationService serviceImplementation="org.jetbrains.kotlin.idea.PluginStartupService"/>