FIR IDE/UAST: introduce UAST-specific version of analyse

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

@@ -10,6 +10,7 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.Task
+import org.jetbrains.kotlin.idea.frontend.api.tokens.AlwaysAccessibleValidityTokenFactory
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ReadActionConfinementValidityTokenFactory
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityTokenFactory
 import org.jetbrains.kotlin.idea.util.application.runReadAction
@@ -92,6 +93,16 @@ inline fun <R> analyseWithCustomToken(
 ): R =
     contextElement.project.service<KtAnalysisSessionProvider>().analyse(contextElement, tokenFactory, action)
 
+/**
+ * UAST-specific version of [analyse] that executes the given [action] in [KtAnalysisSession] context
+ */
+@OptIn(InvalidWayOfUsingAnalysisSession::class)
+inline fun <R> analyseForUast(
+    contextElement: KtElement,
+    action: KtAnalysisSession.() -> R
+): R =
+    analyseWithCustomToken(contextElement, AlwaysAccessibleValidityTokenFactory, action)
+
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
 inline fun <R> analyseInDependedAnalysisSession(
     originalFile: KtFile,
@@ -137,4 +148,4 @@ inline fun <R> analyseInModalWindow(
     }
     task.queue()
     return task.result
-}
+}

plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/FirKotlinUastResolveProviderService.kt

@@ -7,8 +7,7 @@ package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiType
-import org.jetbrains.kotlin.idea.frontend.api.analyseWithCustomToken
-import org.jetbrains.kotlin.idea.frontend.api.tokens.AlwaysAccessibleValidityTokenFactory
+import org.jetbrains.kotlin.idea.frontend.api.analyseForUast
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.psi.KtExpression
@@ -28,7 +27,7 @@ interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderSer
     override fun resolveToDeclaration(ktExpression: KtExpression): PsiElement? {
         when (ktExpression) {
             is KtReferenceExpression -> {
-                analyseWithCustomToken(ktExpression, AlwaysAccessibleValidityTokenFactory) {
+                analyseForUast(ktExpression) {
                     return ktExpression.mainReference.resolve()
                 }
             }
@@ -38,14 +37,14 @@ interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderSer
     }
 
     override fun resolveToType(ktTypeReference: KtTypeReference, source: UElement): PsiType? {
-        analyseWithCustomToken(ktTypeReference, AlwaysAccessibleValidityTokenFactory) {
+        analyseForUast(ktTypeReference) {
             return ktTypeReference.getPsiType(TypeMappingMode.DEFAULT_UAST)
         }
     }
 
     override fun getExpressionType(uExpression: UExpression): PsiType? {
         val ktExpression = uExpression.sourcePsi as? KtExpression ?: return null
-        analyseWithCustomToken(ktExpression, AlwaysAccessibleValidityTokenFactory) {
+        analyseForUast(ktExpression) {
             return ktExpression.getPsiType(TypeMappingMode.DEFAULT_UAST)
         }
     }