[analysis api] get rid of additional analyseWithCustomValidityToken

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10CliAnalysisSessionProvider.kt

@@ -37,7 +37,7 @@ class KtFe10AnalysisSessionProvider : KtAnalysisSessionProvider() {
         throw UnsupportedOperationException("getAnalysisSessionBySymbol() should not be used on KtFe10AnalysisSession")
     }
 
-    override fun getAnalysisSessionByModule(ktModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession {
+    override fun getAnalysisSessionByUseSiteKtModule(ktModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession {
         throw UnsupportedOperationException("getAnalysisSessionByModule() should not be used on KtFe10AnalysisSession")
     }
 

analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/CachingKtAnalysisSessionProvider.kt

@@ -24,7 +24,7 @@ import java.util.concurrent.ConcurrentHashMap
 import kotlin.reflect.KClass
 
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
-abstract class CachingKtAnalysisSessionProvider<State : Any>(private val project: Project) : KtAnalysisSessionProvider() {
+abstract class CachingKtAnalysisSessionProvider<State : Any>(project: Project) : KtAnalysisSessionProvider(project) {
     private val cache = KtAnalysisSessionCache<Pair<State, KClass<out KtLifetimeToken>>>(project)
 
     protected abstract fun getFirResolveSession(contextElement: KtElement): State
@@ -46,7 +46,7 @@ abstract class CachingKtAnalysisSessionProvider<State : Any>(private val project
         }
     }
 
-    final override fun getAnalysisSessionByModule(ktModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession {
+    final override fun getAnalysisSessionByUseSiteKtModule(ktModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession {
         val firResolveSession = getFirResolveSession(ktModule)
         return cache.getAnalysisSession(firResolveSession to factory.identifier) {
             val validityToken = factory.create(project)

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/AnalysisApiBaseTestServiceRegistrar.kt

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.analysis.api.impl.base.test.configurators
 import com.intellij.mock.MockApplication
 import com.intellij.mock.MockProject
 import org.jetbrains.kotlin.analysis.api.impl.base.references.HLApiReferenceProviderService
+import org.jetbrains.kotlin.analysis.api.session.KtDefaultLifetimeTokenProvider
+import org.jetbrains.kotlin.analysis.api.session.KtReadActionConfinementDefaultLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModuleScopeProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModuleScopeProviderImpl
 import org.jetbrains.kotlin.analysis.providers.KotlinAnnotationsResolverFactory
@@ -33,6 +35,7 @@ object AnalysisApiBaseTestServiceRegistrar: AnalysisApiTestServiceRegistrar()  {
 
         project.apply {
             registerService(KotlinModificationTrackerFactory::class.java, KotlinStaticModificationTrackerFactory::class.java)
+            registerService(KtDefaultLifetimeTokenProvider::class.java, KtReadActionConfinementDefaultLifetimeTokenProvider::class.java)
 
             registerService(KtModuleScopeProvider::class.java, KtModuleScopeProviderImpl())
             registerService(KotlinModificationTrackerFactory::class.java, KotlinStaticModificationTrackerFactory::class.java)

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/KtAnalysisSessionProvider.kt

@@ -13,10 +13,9 @@ import com.intellij.openapi.progress.Task
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Computable
 import org.jetbrains.annotations.TestOnly
+import org.jetbrains.kotlin.analysis.api.session.KtDefaultLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
-import org.jetbrains.kotlin.analysis.api.tokens.KtAlwaysAccessibleLifetimeTokenFactory
 import org.jetbrains.kotlin.analysis.api.tokens.KtLifetimeTokenFactory
-import org.jetbrains.kotlin.analysis.api.tokens.ReadActionConfinementValidityTokenFactoryFactory
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -32,7 +31,8 @@ public annotation class KtAnalysisSessionProviderInternals
  * Should not be used directly, consider using [analyse]/[analyseWithReadAction]/[analyseInModalWindow] instead
  */
 @InvalidWayOfUsingAnalysisSession
-public abstract class KtAnalysisSessionProvider : Disposable {
+public abstract class KtAnalysisSessionProvider(public val project: Project) : Disposable {
+
     @Suppress("LeakingThis")
     @OptIn(KtInternalApiMarker::class)
     public val noWriteActionInAnalyseCallChecker: NoWriteActionInAnalyseCallChecker = NoWriteActionInAnalyseCallChecker(this)
@@ -44,7 +44,7 @@ public abstract class KtAnalysisSessionProvider : Disposable {
     public abstract fun getAnalysisSessionBySymbol(contextSymbol: KtSymbol): KtAnalysisSession
 
     @InvalidWayOfUsingAnalysisSession
-    public abstract fun getAnalysisSessionByModule(ktModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession
+    public abstract fun getAnalysisSessionByUseSiteKtModule(ktModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession
 
     @InvalidWayOfUsingAnalysisSession
     public inline fun <R> analyzeWithSymbolAsContext(
@@ -59,20 +59,48 @@ public abstract class KtAnalysisSessionProvider : Disposable {
     public inline fun <R> analyseInDependedAnalysisSession(
         originalFile: KtFile,
         elementToReanalyze: KtElement,
+        nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory?,
         action: KtAnalysisSession.() -> R
     ): R {
-        val dependedAnalysisSession = getAnalysisSession(originalFile, ReadActionConfinementValidityTokenFactoryFactory)
+        val factory =
+            nonDefaultLifetimeTokenFactory ?: KtDefaultLifetimeTokenProvider.getService(project).getDefaultLifetimeTokenFactory()
+
+        val originalAnalysisSession = getAnalysisSession(originalFile, factory)
+        val dependedAnalysisSession = originalAnalysisSession
             .createContextDependentCopy(originalFile, elementToReanalyze)
-        return analyse(dependedAnalysisSession, ReadActionConfinementValidityTokenFactoryFactory, action)
+        return analyse(dependedAnalysisSession, factory, action)
     }
 
     @InvalidWayOfUsingAnalysisSession
-    public inline fun <R> analyse(contextElement: KtElement, tokenFactory: KtLifetimeTokenFactory, action: KtAnalysisSession.() -> R): R =
+    public inline fun <R> analyse(
-        analyse(getAnalysisSession(contextElement, tokenFactory), tokenFactory, action)
+        contextElement: KtElement,
+        nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory?,
+        action: KtAnalysisSession.() -> R
+    ): R {
+        val factory =
+            nonDefaultLifetimeTokenFactory ?: KtDefaultLifetimeTokenProvider.getService(project).getDefaultLifetimeTokenFactory()
+        return analyse(getAnalysisSession(contextElement, factory), factory, action)
+    }
+
+    @InvalidWayOfUsingAnalysisSession
+    public inline fun <R> analyze(
+        useSiteKtModule: KtModule,
+        nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory?,
+        action: KtAnalysisSession.() -> R
+    ): R {
+        val factory =
+            nonDefaultLifetimeTokenFactory ?: KtDefaultLifetimeTokenProvider.getService(project).getDefaultLifetimeTokenFactory()
+
+        return analyse(getAnalysisSessionByUseSiteKtModule(useSiteKtModule, factory), factory, action)
+    }
 
     @OptIn(KtAnalysisSessionProviderInternals::class, KtInternalApiMarker::class)
     @InvalidWayOfUsingAnalysisSession
-    public inline fun <R> analyse(analysisSession: KtAnalysisSession, factory: KtLifetimeTokenFactory, action: KtAnalysisSession.() -> R): R {
+    public inline fun <R> analyse(
+        analysisSession: KtAnalysisSession,
+        factory: KtLifetimeTokenFactory,
+        action: KtAnalysisSession.() -> R
+    ): R {
         noWriteActionInAnalyseCallChecker.beforeEnteringAnalysisContext()
         factory.beforeEnteringAnalysisContext()
         return try {
@@ -107,37 +135,35 @@ public abstract class KtAnalysisSessionProvider : Disposable {
  * @see analyseWithReadAction
  */
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
-public inline fun <R> analyse(contextElement: KtElement, action: KtAnalysisSession.() -> R): R =
+public inline fun <R> analyse(
-    KtAnalysisSessionProvider.getInstance(contextElement.project)
-        .analyse(contextElement, ReadActionConfinementValidityTokenFactoryFactory, action)
-
-@OptIn(InvalidWayOfUsingAnalysisSession::class)
-public inline fun <R> analyseWithCustomToken(
     contextElement: KtElement,
-    tokenFactory: KtLifetimeTokenFactory,
+    nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory? = null,
     action: KtAnalysisSession.() -> R
 ): R =
     KtAnalysisSessionProvider.getInstance(contextElement.project)
-        .analyse(contextElement, tokenFactory, action)
+        .analyse(contextElement, nonDefaultLifetimeTokenFactory, action)
+
 
-/**
- * UAST-specific version of [analyse] that executes the given [action] in [KtAnalysisSession] context
- */
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
-public inline fun <R> analyseForUast(
+public inline fun <R> analyze(
-    contextElement: KtElement,
+    ktModule: KtModule,
-    action: KtAnalysisSession.() -> R
+    nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory? = null,
-): R =
+    crossinline action: KtAnalysisSession.() -> R
-    analyseWithCustomToken(contextElement, KtAlwaysAccessibleLifetimeTokenFactory, action)
+): R {
+    checkNotNull(ktModule.project)
+    val sessionProvider = KtAnalysisSessionProvider.getInstance(ktModule.project!!)
+    return sessionProvider.analyze(ktModule, nonDefaultLifetimeTokenFactory, action)
+}
 
 @OptIn(InvalidWayOfUsingAnalysisSession::class)
 public inline fun <R> analyseInDependedAnalysisSession(
     originalFile: KtFile,
     elementToReanalyze: KtElement,
+    nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory? = null,
     action: KtAnalysisSession.() -> R
 ): R =
     KtAnalysisSessionProvider.getInstance(originalFile.project)
-        .analyseInDependedAnalysisSession(originalFile, elementToReanalyze, action)
+        .analyseInDependedAnalysisSession(originalFile, elementToReanalyze, nonDefaultLifetimeTokenFactory, action)
 
 /**
  * Execute given [action] in [KtAnalysisSession] context like [analyse] does but execute it in read action
@@ -153,9 +179,10 @@ public inline fun <R> analyseInDependedAnalysisSession(
  */
 public inline fun <R> analyseWithReadAction(
     contextElement: KtElement,
+    nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory? = null,
     crossinline action: KtAnalysisSession.() -> R
 ): R = ApplicationManager.getApplication().runReadAction(Computable {
-    analyse(contextElement, action)
+    analyse(contextElement, nonDefaultLifetimeTokenFactory, action)
 })
 
 /**
@@ -168,23 +195,14 @@ public inline fun <R> analyseWithReadAction(
 public inline fun <R> analyseInModalWindow(
     contextElement: KtElement,
     windowTitle: String,
+    nonDefaultLifetimeTokenFactory: KtLifetimeTokenFactory? = null,
     crossinline action: KtAnalysisSession.() -> R
 ): R {
     ApplicationManager.getApplication().assertIsDispatchThread()
     val task = object : Task.WithResult<R, Exception>(contextElement.project, windowTitle, /*canBeCancelled*/ true) {
-        override fun compute(indicator: ProgressIndicator): R = analyseWithReadAction(contextElement) { action() }
+        override fun compute(indicator: ProgressIndicator): R =
+            analyseWithReadAction(contextElement, nonDefaultLifetimeTokenFactory) { action() }
     }
     task.queue()
     return task.result
 }
-
-@OptIn(InvalidWayOfUsingAnalysisSession::class)
-public inline fun <R> analyzeWithKtModule(
-    ktModule: KtModule,
-    tokenFactory: KtLifetimeTokenFactory,
-    crossinline action: KtAnalysisSession.() -> R
-): R {
-    checkNotNull(ktModule.project)
-    val sessionProvider = KtAnalysisSessionProvider.getInstance(ktModule.project!!)
-    return sessionProvider.getAnalysisSessionByModule(ktModule, tokenFactory).action()
-}

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/session/KtDefaultLifetimeTokenProvider.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2010-2022 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.analysis.api.session
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analysis.api.tokens.KtLifetimeTokenFactory
+import org.jetbrains.kotlin.analysis.api.tokens.ReadActionConfinementValidityTokenFactory
+
+public abstract class KtDefaultLifetimeTokenProvider {
+    public abstract fun getDefaultLifetimeTokenFactory(): KtLifetimeTokenFactory
+
+    public companion object {
+        public fun getService(project: Project): KtDefaultLifetimeTokenProvider =
+            project.getService(KtDefaultLifetimeTokenProvider::class.java)
+    }
+}
+
+public class KtReadActionConfinementDefaultLifetimeTokenProvider: KtDefaultLifetimeTokenProvider() {
+    override fun getDefaultLifetimeTokenFactory(): KtLifetimeTokenFactory {
+        return ReadActionConfinementValidityTokenFactory
+    }
+}

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/tokens/KtReadActionConfinementLifetimeToken.kt

@@ -31,7 +31,7 @@ public class KtReadActionConfinementLifetimeToken(project: Project) : KtLifetime
         if (application.isDispatchThread && !allowOnEdt.get()) return false
         if (ForbidKtResolve.resovleIsForbidenInActionWithName.get() != null) return false
         if (!application.isReadAccessAllowed) return false
-        if (!ReadActionConfinementValidityTokenFactoryFactory.isInsideAnalysisContext()) return false
+        if (!ReadActionConfinementValidityTokenFactory.isInsideAnalysisContext()) return false
         return true
     }
 
@@ -43,7 +43,7 @@ public class KtReadActionConfinementLifetimeToken(project: Project) : KtLifetime
         ForbidKtResolve.resovleIsForbidenInActionWithName.get()?.let { actionName ->
             return "Resolve is forbidden in $actionName"
         }
-        if (!ReadActionConfinementValidityTokenFactoryFactory.isInsideAnalysisContext()) return "Called outside analyse method"
+        if (!ReadActionConfinementValidityTokenFactory.isInsideAnalysisContext()) return "Called outside analyse method"
         error("Getting inaccessibility reason for validity token when it is accessible")
     }
 
@@ -53,10 +53,10 @@ public class KtReadActionConfinementLifetimeToken(project: Project) : KtLifetime
         public val allowOnEdt: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
     }
 
-    public override val factory: KtLifetimeTokenFactory = ReadActionConfinementValidityTokenFactoryFactory
+    public override val factory: KtLifetimeTokenFactory = ReadActionConfinementValidityTokenFactory
 }
 
-public object ReadActionConfinementValidityTokenFactoryFactory : KtLifetimeTokenFactory() {
+public object ReadActionConfinementValidityTokenFactory : KtLifetimeTokenFactory() {
     override val identifier: KClass<out KtLifetimeToken> = KtReadActionConfinementLifetimeToken::class
 
     override fun create(project: Project): KtLifetimeToken = KtReadActionConfinementLifetimeToken(project)

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/AbstractAnalysisApiBasedTest.kt

@@ -171,9 +171,13 @@ abstract class AbstractAnalysisApiBasedTest : TestWithDisposable() {
         return if (configurator.analyseInDependentSession) {
             val originalContainingFile = contextElement.containingKtFile
             val fileCopy = originalContainingFile.copy() as KtFile
-            analyseInDependedAnalysisSession(originalContainingFile, PsiTreeUtil.findSameElementInCopy(contextElement, fileCopy), action)
+            analyseInDependedAnalysisSession(
+                originalContainingFile,
+                PsiTreeUtil.findSameElementInCopy(contextElement, fileCopy),
+                action = action
+            )
         } else {
-            analyse(contextElement, action)
+            analyse(contextElement, action = action)
         }
     }
 

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/analyseForLightClasses.kt

@@ -6,9 +6,9 @@
 package org.jetbrains.kotlin.light.classes.symbol.classes
 
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
-import org.jetbrains.kotlin.analysis.api.analyseWithCustomToken
+import org.jetbrains.kotlin.analysis.api.analyse
 import org.jetbrains.kotlin.analysis.api.tokens.KtAlwaysAccessibleLifetimeTokenFactory
 import org.jetbrains.kotlin.psi.KtElement
 
 internal inline fun <R> analyseForLightClasses(context: KtElement, action: KtAnalysisSession.() -> R): R =
-    analyseWithCustomToken(context, KtAlwaysAccessibleLifetimeTokenFactory, action)
+    analyse(context, KtAlwaysAccessibleLifetimeTokenFactory, action)