[AA REx] Require a parent Disposable when creating resolve extensions.

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

@@ -5,23 +5,49 @@
 
 package org.jetbrains.kotlin.analysis.api.impl.base.projectStructure
 
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.util.Disposer
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.api.KaImplementationDetail
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinContentScopeRefiner
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
+import org.jetbrains.kotlin.analysis.api.resolve.extensions.KaResolveExtension
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KaResolveExtensionProvider
 
 @KaImplementationDetail
 class KaResolveExtensionToContentScopeRefinerBridge : KotlinContentScopeRefiner {
     override fun getEnlargementScopes(module: KaModule): List<GlobalSearchScope> =
-        buildList {
-            if (KaResolveExtensionProvider.provideExtensionsFor(module).isNotEmpty()) {
-                add(KaBaseResolveExtensionGeneratedFilesScope(listOf(module)))
+        withResolveExtensionsFor(module) { extensions ->
+            if (extensions.isNotEmpty()) {
+                listOf(KaBaseResolveExtensionGeneratedFilesScope(listOf(module)))
+            } else {
+                listOf()
             }
         }
 
     override fun getRestrictionScopes(module: KaModule): List<GlobalSearchScope> =
-        KaResolveExtensionProvider.provideExtensionsFor(module).map { resolveExtension ->
-            GlobalSearchScope.notScope(resolveExtension.getShadowedScope())
+        withResolveExtensionsFor(module) { extensions ->
+            extensions.map { GlobalSearchScope.notScope(it.getShadowedScope()) }
         }
+
+    companion object {
+        private inline fun <T> withResolveExtensionsFor(
+            module: KaModule,
+            block: (List<KaResolveExtension>) -> T,
+        ): T {
+            var disposable: Disposable? = null
+            try {
+                val extensions =
+                    KaResolveExtensionProvider.provideExtensionsFor(module) {
+                        Disposer.newDisposable("KaResolveExtensionToContentScopeRefinerBridge")
+                            .also { disposable = it }
+                    }
+                return block(extensions)
+            } finally {
+                if (disposable != null) {
+                    Disposer.dispose(disposable)
+                }
+            }
+        }
+    }
 }

analysis/analysis-api/api/analysis-api.undocumented

@@ -569,7 +569,6 @@ org.jetbrains.kotlin.analysis.api.resolution.KaSymbolBasedReference
 org.jetbrains.kotlin.analysis.api.resolution.KaSymbolBasedReference.resolveToSymbols()
 org.jetbrains.kotlin.analysis.api.resolve.extensions.KaResolveExtensionNavigationTargetsProvider
 org.jetbrains.kotlin.analysis.api.resolve.extensions.KaResolveExtensionProvider.Companion.EP_NAME
-org.jetbrains.kotlin.analysis.api.resolve.extensions.KaResolveExtensionProvider.Companion.provideExtensionsFor(KaModule)
 org.jetbrains.kotlin.analysis.api.scopes.KaScopeLike
 org.jetbrains.kotlin.analysis.api.session.KaSessionProvider.Companion.getInstance(Project)
 org.jetbrains.kotlin.analysis.api.session.KaSessionProvider.analyze(KaModule, KaSession.() -> R)

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/resolve/extensions/KaResolveExtensionProvider.kt

@@ -5,7 +5,9 @@
 
 package org.jetbrains.kotlin.analysis.api.resolve.extensions
 
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.util.Disposer
 import org.jetbrains.kotlin.analysis.api.KaExperimentalApi
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
 
@@ -37,8 +39,24 @@ public abstract class KaResolveExtensionProvider {
         public val EP_NAME: ExtensionPointName<KaResolveExtensionProvider> =
             ExtensionPointName<KaResolveExtensionProvider>("org.jetbrains.kotlin.kaResolveExtensionProvider")
 
-        public fun provideExtensionsFor(module: KaModule): List<KaResolveExtension> {
-            return EP_NAME.getExtensionList(module.project).flatMap { it.provideExtensionsFor(module) }
+        /**
+         * Creates [resolve extensions][KaResolveExtension] for the provided [KaModule].
+         * The [Disposable] provided by the factory will be used as a parent for all returned extensions.
+         *
+         * @param module The [KaModule] for which to create extensions.
+         * @param parentDisposableFactory A factory method to retrieve a parent [Disposable] for the returned
+         *   extensions. This factory will only be invoked if at least one extension is created.
+         */
+        public fun provideExtensionsFor(
+            module: KaModule,
+            parentDisposableFactory: () -> Disposable,
+        ): List<KaResolveExtension> {
+            val extensions = EP_NAME.getExtensionList(module.project).flatMap { it.provideExtensionsFor(module) }
+            if (extensions.isEmpty()) return emptyList()
+
+            val parentDisposable = parentDisposableFactory()
+            extensions.forEach { Disposer.register(parentDisposable, it) }
+            return extensions
         }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/projectStructure/sessionFactoryHelpers.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.projectStructure
 
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.Disposer
 import org.jetbrains.kotlin.analysis.api.platform.declarations.createAnnotationResolver
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KaResolutionScopeProvider
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinCompilerPluginsProvider
@@ -62,14 +61,11 @@ internal fun LLFirSession.registerIdeComponents(project: Project, languageVersio
 private fun LLFirSession.registerResolveExtensionTool() {
     val resolveExtensionTool = createResolveExtensionTool() ?: return
 
-    // `KaResolveExtension`s are disposables meant to be tied to the lifetime of the `LLFirSession`.
-    resolveExtensionTool.extensions.forEach { Disposer.register(requestDisposable(), it) }
-
     register(LLFirResolveExtensionTool::class, resolveExtensionTool)
 }
 
 private fun LLFirSession.createResolveExtensionTool(): LLFirResolveExtensionTool? {
-    val extensions = KaResolveExtensionProvider.provideExtensionsFor(ktModule)
+    val extensions = KaResolveExtensionProvider.provideExtensionsFor(ktModule) { requestDisposable() }
     if (extensions.isEmpty()) return null
     return LLFirNonEmptyResolveExtensionTool(this, extensions)
 }