[Analysis API] implement API to extend Kotlin resolution by generated declarations

^KT-57930 fixed

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

@@ -5,12 +5,20 @@
 
 package org.jetbrains.kotlin.analysis.api.impl.base.components
 
+import com.intellij.openapi.module.Module
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.components.KtAnalysisScopeProvider
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.KtModuleStructureInternals
+import org.jetbrains.kotlin.analysis.project.structure.allDirectDependencies
+import org.jetbrains.kotlin.analysis.project.structure.analysisExtensionFileContextModule
 import org.jetbrains.kotlin.analysis.providers.KotlinResolutionScopeProvider
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.contains
 
 class KtAnalysisScopeProviderImpl(
@@ -18,13 +26,40 @@ class KtAnalysisScopeProviderImpl(
     override val token: KtLifetimeToken
 ) : KtAnalysisScopeProvider() {
 
-    private val scope by lazy(LazyThreadSafetyMode.PUBLICATION) {
+    private val baseResolveScope by lazy(LazyThreadSafetyMode.PUBLICATION) {
         KotlinResolutionScopeProvider.getInstance(analysisSession.useSiteModule.project).getResolutionScope(analysisSession.useSiteModule)
     }
 
-    override fun getAnalysisScope(): GlobalSearchScope = scope
+    private val resolveScope by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        KtAnalysisScopeProviderResolveScope(baseResolveScope, analysisSession.useSiteModule)
+    }
+
+    override fun getAnalysisScope(): GlobalSearchScope = resolveScope
 
     override fun canBeAnalysed(psi: PsiElement): Boolean {
-        return scope.contains(psi)
+        return baseResolveScope.contains(psi)
+                || psi.isFromGeneratedModule()
+    }
+
+    private fun PsiElement.isFromGeneratedModule(): Boolean {
+        val ktFile = containingFile as? KtFile ?: return false
+        return ktFile.virtualFile?.isFromGeneratedModule(analysisSession.useSiteModule) == true
     }
 }
+
+private class KtAnalysisScopeProviderResolveScope(
+    private val base: GlobalSearchScope,
+    private val useSiteModule: KtModule
+) : GlobalSearchScope() {
+    override fun getProject(): Project? = base.project
+    override fun isSearchInModuleContent(aModule: Module): Boolean = base.isSearchInModuleContent(aModule)
+    override fun isSearchInLibraries(): Boolean = base.isSearchInLibraries
+    override fun contains(file: VirtualFile): Boolean = base.contains(file) || file.isFromGeneratedModule(useSiteModule)
+}
+
+@OptIn(KtModuleStructureInternals::class)
+private fun VirtualFile.isFromGeneratedModule(useSiteModule: KtModule): Boolean {
+    val analysisContextModule = analysisExtensionFileContextModule ?: return false
+    if (analysisContextModule == useSiteModule) return true
+    return analysisContextModule in useSiteModule.allDirectDependencies()
+}

analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleProvider.kt

@@ -14,6 +14,8 @@ class KtStaticModuleProvider(
     private val builtinsModule: KtBuiltinsModule,
     val projectStructure: KtModuleProjectStructure,
 ) : ProjectStructureProvider() {
+
+    @OptIn(KtModuleStructureInternals::class)
     override fun getKtModuleForKtElement(element: PsiElement): KtModule {
         val containingFileAsPsiFile = element.containingFile
         val containingFileAsVirtualFile = containingFileAsPsiFile.virtualFile
@@ -21,6 +23,8 @@ class KtStaticModuleProvider(
             return builtinsModule
         }
 
+        containingFileAsVirtualFile.analysisExtensionFileContextModule?.let { return it }
+
         return projectStructure.mainModules
             .first { module ->
                 element in module.ktModule.contentScope

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2010-2023 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.resolve.extensions
+
+import com.intellij.openapi.util.ModificationTracker
+import org.jetbrains.kotlin.name.FqName
+
+/**
+ * Provides a list of Kotlin files which provides additional, generated declarations for resolution.
+ *
+ * Provided by the [KtResolveExtensionProvider].
+ *
+ * All member implementations should:
+ * - consider caching the results for subsequent invocations.
+ * - be lightweight and not build the whole file structure inside.
+ * - not use the Kotlin resolve inside, as this function is called during session initialization, so Analysis API access is forbidden.
+ *
+ * @see KtResolveExtensionFile
+ * @see KtResolveExtensionProvider
+ */
+abstract class KtResolveExtension {
+    /**
+     * Get the list of files that should be generated for the module.
+     *
+     * The content of those files should remain valid until the tracker [getModificationTracker] is modified.
+     *
+     * Returned files should contain a valid Kotlin code.
+     *
+     * @see KtResolveExtensionFile
+     * @see KtResolveExtension
+     */
+    abstract fun getKtFiles(): List<KtResolveExtensionFile>
+
+    /**
+     * Returns a [ModificationTracker], which controls the validity lifecycle of the files provided by [getKtFiles].
+     *
+     * All files generated by [getKtFiles] should be valid if the [ModificationTracker] did not change.
+     * If files become invalid (e.g., the in-source declarations they were based on changed) the [ModificationTracker] should be incremented.
+     *
+     * @see KtResolveExtension
+     */
+    abstract fun getModificationTracker(): ModificationTracker
+
+    /**
+     * Returns the set of packages that are contained in the files provided by [getKtFiles].
+     *
+     * The returned package set should be a strict set of all file packages,
+     * so `for-all pckg: pckg in getContainedPackages() <=> exists file: file in getKtFiles() && file.getFilePackageName() == pckg`
+     *
+     * @see KtResolveExtension
+     */
+    abstract fun getContainedPackages(): Set<FqName>
+}

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

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2010-2023 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.resolve.extensions
+
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+/**
+ * Represents the Kotlin file which provides additional, generated declarations for resolution.
+ *
+ * All member implementations should:
+ * - consider caching the results for subsequent invocations.
+ * - be lightweight and not build the whole file structure inside.
+ * - not use the Kotlin resolve inside, as this function is called during session initialization, so Analysis API access is forbidden.
+ *
+ * @see KtResolveExtension
+ */
+abstract class KtResolveExtensionFile {
+    /**
+     * The name a Kotlin file which will be generated.
+     *
+     * Should have the `.kt` extension.
+     *
+     * It will be used as a Java facade name, e.g., for the file name `myFile.kt`, the `MyFileKt` facade is generated if the file contains some properties or functions.
+     *
+     * @see KtResolveExtensionFile
+     */
+    abstract fun getFileName(): String
+
+    /**
+     * [FqName] of the package specified in the file
+     *
+     * The operation might be called regularly, so the [getFilePackageName] should work fast and avoid building the whole file text.
+     *
+     * It should be equal to the package name specified in the [buildFileText].
+     *
+     * @see KtResolveExtensionFile
+     */
+    abstract fun getFilePackageName(): FqName
+
+    /**
+     * Returns the set of top-level classifier (classes, interfaces, objects, and type-aliases) names in the file.
+     *
+     * The result may have false-positive entries but cannot have false-negative entries. It should contain all the names in the package but may have some additional names that are not there.
+     *
+     * @see KtResolveExtensionFile
+     */
+    abstract fun getTopLevelClassifierNames(): Set<Name>
+
+    /**
+     * Returns the set of top-level callable (functions and properties) names in the file.
+     *
+     * The result may have false-positive entries but cannot have false-negative entries. It should contain all the names in the package but may have some additional names that are not there.
+     *
+     * @see KtResolveExtensionFile
+     */
+    abstract fun getTopLevelCallableNames(): Set<Name>
+
+    /**
+     * Creates the generated Kotlin source file text.
+     *
+     * The resulted String should be a valid Kotlin code.
+     * It should be consistent with other declarations which are present in the [KtResolveExtensionFile], more specifically:
+     * 1. [getFilePackageName] should be equal to the file's package name.
+     * 2. All classifier names should be contained in the [getTopLevelClassifierNames].
+     * 3. All callable names should be contained in the [getTopLevelCallableNames].
+     *
+     * Additional restrictions on the file text:
+     * 1. The File should not contain the `kotlin.jvm.JvmMultifileClass` and `kotlin.jvm.JvmName` annotations on the file level.
+     * 2. All declaration types should be specified explicitly.
+     *
+     * @see KtResolveExtensionFile
+     */
+    abstract fun buildFileText(): String
+
+    /**
+     * Creates a [KtResolveExtensionReferenceTargetsPsiProvider] for this [KtResolveExtensionFile].
+     *
+     * @see KtResolveExtensionReferenceTargetsPsiProvider
+     * @see KtResolveExtensionFile
+     */
+    abstract fun createTargetPsiProvider(): KtResolveExtensionReferenceTargetsPsiProvider
+
+}

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2023 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.resolve.extensions
+
+import com.intellij.openapi.extensions.ExtensionPointName
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+
+/**
+ * Allows extending Kotlin resolution by generating additional declarations.
+ *
+ * Those declarations will be analyzed the same way as they were just regular source files inside the project.
+ *
+ * All member implementations should consider caching the results for subsequent invocations.
+ */
+abstract class KtResolveExtensionProvider {
+    /**
+     * Provides a list of [KtResolveExtension]s for a given [KtModule].
+     *
+     * Should not perform any heavy analysis and the generation of the actual files. All file generation should be performed only in [KtResolveExtensionFile.buildFileText].
+     *
+     * Implementations should consider caching the results, so the subsequent invocations should be performed instantly.
+     *
+     * Implementation cannot use the Kotlin resolve inside, as this function is called during session initialization, so Analysis API access is forbidden.
+     */
+    abstract fun provideExtensionsFor(module: KtModule): List<KtResolveExtension>
+
+    companion object {
+        val EP_NAME = ExtensionPointName<KtResolveExtensionProvider>("org.jetbrains.kotlin.ktResolveExtensionProvider")
+
+        fun provideExtensionsFor(module: KtModule): List<KtResolveExtension> {
+            return EP_NAME.getExtensionList(module.project).flatMap { it.provideExtensionsFor(module) }
+        }
+    }
+}

analysis/low-level-api-fir/build.gradle.kts

@@ -24,6 +24,7 @@ dependencies {
     implementation(project(":compiler:frontend.common"))
     implementation(project(":compiler:fir:entrypoint"))
     implementation(project(":analysis:analysis-api-providers"))
+    implementation(project(":analysis:analysis-api"))
     implementation(project(":analysis:analysis-internal-utils"))
 
     // We cannot use the latest version `3.1.5` because it doesn't support Java 8.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/FileBasedKotlinDeclarationProvider.kt

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.yieldIfNotNull
 
-internal class FileBasedKotlinDeclarationProvider(private val kotlinFile: KtFile) : KotlinDeclarationProvider() {
+internal class FileBasedKotlinDeclarationProvider(val kotlinFile: KtFile) : KotlinDeclarationProvider() {
     private val topLevelDeclarations: Sequence<KtDeclaration>
         get() {
             return sequence {

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

@@ -7,11 +7,14 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure
 
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.IdeSessionComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.createSealedInheritorsProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.fir.caches.FirThreadSafeCachesFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirIdePredicateBasedProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirIdeRegisteredPluginAnnotations
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirNonEmptyResolveExtensionTool
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionTool
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSourcesSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirExceptionHandler
@@ -42,8 +45,18 @@ internal fun LLFirSession.registerIdeComponents(project: Project) {
     register(FirCachesFactory::class, FirThreadSafeCachesFactory)
     register(SealedClassInheritorsProvider::class, project.createSealedInheritorsProvider())
     register(FirExceptionHandler::class, LLFirExceptionHandler)
+    createResolveExtensionTool()?.let {
+        register(LLFirResolveExtensionTool::class, it)
+    }
 }
 
+private fun LLFirSession.createResolveExtensionTool(): LLFirResolveExtensionTool? {
+    val extensions = KtResolveExtensionProvider.provideExtensionsFor(ktModule)
+    if (extensions.isEmpty()) return null
+    return LLFirNonEmptyResolveExtensionTool(this, extensions)
+}
+
+
 internal inline fun createCompositeSymbolProvider(
     session: FirSession,
     createSubProviders: MutableList<FirSymbolProvider>.() -> Unit
@@ -52,7 +65,7 @@ internal inline fun createCompositeSymbolProvider(
 
 @SessionConfiguration
 internal fun FirSession.registerCompilerPluginExtensions(project: Project, module: KtSourceModule) {
-    val extensionProvider = project.getService(KtCompilerPluginsProvider::class.java) ?: return
+    val extensionProvider = project.getService<KtCompilerPluginsProvider>(KtCompilerPluginsProvider::class.java) ?: return
     FirSessionConfigurator(this).apply {
         val registrars = FirExtensionRegistrarAdapter.getInstances(project) +
                 extensionProvider.getRegisteredExtensions(module, FirExtensionRegistrarAdapter)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt

@@ -6,7 +6,12 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.CompositeKotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.CompositeKotlinPackageProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionTool
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.llResolveExtensionTool
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirCompositeSymbolProviderNameCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirKotlinSymbolProviderNameCache
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
@@ -32,10 +37,25 @@ import org.jetbrains.kotlin.psi.KtFile
 internal class LLFirProviderHelper(
     firSession: FirSession,
     private val firFileBuilder: LLFirFileBuilder,
-    private val declarationProvider: KotlinDeclarationProvider,
-    private val packageProvider: KotlinPackageProvider,
+    mainDeclarationProvider: KotlinDeclarationProvider,
+    mainPackageProvider: KotlinPackageProvider,
     canContainKotlinPackage: Boolean,
 ) {
+    private val extensionTool: LLFirResolveExtensionTool? = firSession.llResolveExtensionTool
+
+    private val declarationProvider = CompositeKotlinDeclarationProvider.create(
+        listOfNotNull(
+            mainDeclarationProvider,
+            extensionTool?.declarationProvider,
+        )
+    )
+
+    private val packageProvider = CompositeKotlinPackageProvider.create(
+        listOfNotNull(
+            mainPackageProvider,
+            extensionTool?.packageProvider,
+        )
+    )
     private val allowKotlinPackage = canContainKotlinPackage ||
             firSession.languageVersionSettings.getFlag(AnalysisFlags.allowKotlinPackage)
 
@@ -60,7 +80,12 @@ internal class LLFirProviderHelper(
             }
         }
 
-    val symbolNameCache = LLFirKotlinSymbolProviderNameCache(firSession, declarationProvider)
+    val symbolNameCache = LLFirCompositeSymbolProviderNameCache.create(
+        listOfNotNull(
+            LLFirKotlinSymbolProviderNameCache(firSession, declarationProvider),
+            extensionTool?.symbolNameCache,
+        )
+    )
 
     fun getFirClassifierByFqNameAndDeclaration(
         classId: ClassId,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/extensions/LLFirResolveExtensionTool.kt

@@ -0,0 +1,276 @@
+/*
+ * Copyright 2010-2023 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.low.level.api.fir.resolve.extensions
+
+import com.intellij.openapi.util.ModificationTracker
+import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtension
+import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionFile
+import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionReferenceTargetsPsiProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.FileBasedKotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirSymbolProviderNameCache
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.KtModuleStructureInternals
+import org.jetbrains.kotlin.analysis.project.structure.analysisExtensionFileContextModule
+import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.load.kotlin.PackagePartClassUtils
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.psi.*
+import java.util.concurrent.ConcurrentHashMap
+
+/**
+ * Encapsulate all the work with the [KtResolveExtension] for the LL API.
+ *
+ * Caches generated [KtResolveExtensionFile]s, creates [KotlinDeclarationProvider], [KotlinPackageProvider], [LLFirSymbolProviderNameCache] needed for the [org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider].
+ */
+internal abstract class LLFirResolveExtensionTool : FirSessionComponent {
+    abstract val modificationTrackers: List<ModificationTracker>
+    abstract val declarationProvider: KotlinDeclarationProvider
+    abstract val packageProvider: KotlinPackageProvider
+    abstract val symbolNameCache: LLFirSymbolProviderNameCache
+}
+
+internal val FirSession.llResolveExtensionTool: LLFirResolveExtensionTool? by FirSession.nullableSessionComponentAccessor()
+
+internal class LLFirNonEmptyResolveExtensionTool(
+    session: LLFirSession,
+    extensions: List<KtResolveExtension>,
+) : LLFirResolveExtensionTool() {
+    init {
+        require(extensions.isNotEmpty())
+    }
+
+    private val fileProvider = LLFirResolveExtensionsFileProvider(extensions)
+
+    private val packageFilter = LLFirResolveExtensionToolPackageFilter(extensions)
+
+    override val modificationTrackers by lazy { extensions.map { it.getModificationTracker() } }
+
+    override val declarationProvider: KotlinDeclarationProvider =
+        LLFirResolveExtensionToolDeclarationProvider(fileProvider, session.ktModule)
+
+    override val packageProvider: KotlinPackageProvider = LLFirResolveExtensionToolPackageProvider(packageFilter)
+
+    override val symbolNameCache: LLFirSymbolProviderNameCache = LLFirResolveExtensionToolNameCache(packageFilter, fileProvider)
+}
+
+private class LLFirResolveExtensionToolNameCache(
+    private val packageFilter: LLFirResolveExtensionToolPackageFilter,
+    private val fileProvider: LLFirResolveExtensionsFileProvider,
+) : LLFirSymbolProviderNameCache() {
+    override fun getTopLevelClassifierNamesInPackage(packageFqName: FqName): Set<String>? {
+        if (!packageFilter.packageExists(packageFqName)) return emptySet()
+        return fileProvider.getFilesByPackage(packageFqName)
+            .flatMap { it.getTopLevelClassifierNames() }
+            .mapTo(mutableSetOf()) { it.asString() }
+    }
+
+    override fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name> {
+        if (!packageFilter.packageExists(packageFqName)) return emptySet()
+        return fileProvider.getFilesByPackage(packageFqName)
+            .flatMapTo(mutableSetOf()) { it.getTopLevelCallableNames() }
+    }
+
+    override fun mayHaveTopLevelClassifier(classId: ClassId, mayHaveFunctionClass: Boolean): Boolean {
+        if (!packageFilter.packageExists(classId.packageFqName)) return false
+
+        return fileProvider.getFilesByPackage(classId.packageFqName)
+            .any { it.mayHaveTopLevelClassifier(classId.getTopLevelShortClassName()) }
+    }
+
+    override fun mayHaveTopLevelCallable(packageFqName: FqName, name: Name): Boolean {
+        if (!packageFilter.packageExists(packageFqName)) return false
+
+        return fileProvider.getFilesByPackage(packageFqName)
+            .any { it.mayHaveTopLevelCallable(name) }
+    }
+}
+
+private class LLFirResolveExtensionToolPackageFilter(
+    private val extensions: List<KtResolveExtension>
+) {
+    val allPackages: Set<FqName> by lazy {
+        extensions.flatMapTo(mutableSetOf()) { it.getContainedPackages() }
+    }
+
+    fun packageExists(packageFqName: FqName): Boolean {
+        return packageFqName in allPackages
+    }
+}
+
+private class LLFirResolveExtensionToolDeclarationProvider(
+    private val extensionProvider: LLFirResolveExtensionsFileProvider,
+    private val ktModule: KtModule,
+) : KotlinDeclarationProvider() {
+
+    private val extensionFileToDeclarationProvider: ConcurrentHashMap<KtResolveExtensionFile, FileBasedKotlinDeclarationProvider> =
+        ConcurrentHashMap()
+
+    override fun getClassLikeDeclarationByClassId(classId: ClassId): KtClassLikeDeclaration? {
+        return getDeclarationProvidersByPackage(classId.packageFqName) { it.mayHaveTopLevelClassifier(classId.getTopLevelShortClassName()) }
+            .firstNotNullOfOrNull { it.getClassLikeDeclarationByClassId(classId) }
+    }
+
+    override fun getAllClassesByClassId(classId: ClassId): Collection<KtClassOrObject> {
+        return getDeclarationProvidersByPackage(classId.packageFqName) { it.mayHaveTopLevelClassifier(classId.getTopLevelShortClassName()) }
+            .flatMapTo(mutableListOf()) { it.getAllClassesByClassId(classId) }
+    }
+
+    override fun getAllTypeAliasesByClassId(classId: ClassId): Collection<KtTypeAlias> {
+        return getDeclarationProvidersByPackage(classId.packageFqName) { it.mayHaveTopLevelClassifier(classId.getTopLevelShortClassName()) }
+            .flatMapTo(mutableListOf()) { it.getAllTypeAliasesByClassId(classId) }
+    }
+
+    override fun getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName: FqName): Set<Name> {
+        return getDeclarationProvidersByPackage(packageFqName) { true }
+            .flatMapTo(mutableSetOf()) { it.getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName) }
+    }
+
+    override fun getTopLevelProperties(callableId: CallableId): Collection<KtProperty> {
+        return getDeclarationProvidersByPackage(callableId.packageName) { it.mayHaveTopLevelCallable(callableId.callableName) }
+            .flatMapTo(mutableListOf()) { it.getTopLevelProperties(callableId) }
+    }
+
+    override fun getTopLevelFunctions(callableId: CallableId): Collection<KtNamedFunction> {
+        return getDeclarationProvidersByPackage(callableId.packageName) { it.mayHaveTopLevelCallable(callableId.callableName) }
+            .flatMapTo(mutableListOf()) { it.getTopLevelFunctions(callableId) }
+    }
+
+    override fun getTopLevelCallableFiles(callableId: CallableId): Collection<KtFile> {
+        return getDeclarationProvidersByPackage(callableId.packageName) { it.mayHaveTopLevelCallable(callableId.callableName) }
+            .mapTo(mutableListOf()) { it.kotlinFile }
+    }
+
+    override fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name> {
+        return extensionProvider.getFilesByPackage(packageFqName).flatMapTo(mutableSetOf()) { it.getTopLevelClassifierNames() }
+    }
+
+    override fun findFilesForFacadeByPackage(packageFqName: FqName): Collection<KtFile> {
+        return getDeclarationProvidersByPackage(packageFqName) { file ->
+            file.getTopLevelCallableNames().isNotEmpty()
+        }.mapTo(mutableListOf()) { it.kotlinFile }
+    }
+
+    override fun findFilesForFacade(facadeFqName: FqName): Collection<KtFile> {
+        if (facadeFqName.isRoot) return emptyList()
+        val packageFqName = facadeFqName.parent()
+        return getDeclarationProvidersByPackage(packageFqName) { file ->
+            facadeFqName.shortName().asString() == PackagePartClassUtils.getFilePartShortName(file.getFileName())
+        }
+            .mapTo(mutableListOf()) { it.kotlinFile }
+    }
+
+    override fun findInternalFilesForFacade(facadeFqName: FqName): Collection<KtFile> {
+        // no decompiled files here (see the `org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider.findInternalFilesForFacade` KDoc)
+        return emptyList()
+    }
+
+    private inline fun getDeclarationProvidersByPackage(
+        packageFqName: FqName,
+        crossinline filter: (KtResolveExtensionFile) -> Boolean
+    ): Sequence<FileBasedKotlinDeclarationProvider> {
+        return extensionProvider.getFilesByPackage(packageFqName)
+            .filter { filter(it) }
+            .map { createDeclarationProviderByFile(it) }
+    }
+
+    private fun createDeclarationProviderByFile(file: KtResolveExtensionFile): FileBasedKotlinDeclarationProvider {
+        return extensionFileToDeclarationProvider.getOrPut(file) {
+            val factory = KtPsiFactory(
+                ktModule.project,
+                markGenerated = true,
+                eventSystemEnabled = true // so every generated KtFile backed by some VirtualFile
+            )
+            val text = file.buildFileText()
+            val ktFile = createKtFile(factory, file.getFileName(), text)
+            FileBasedKotlinDeclarationProvider(ktFile)
+        }
+    }
+
+
+    @OptIn(KtModuleStructureInternals::class)
+    private fun createKtFile(factory: KtPsiFactory, fileName: String, fileText: String): KtFile {
+        val ktFile = factory.createFile(fileName, fileText)
+        ktFile.virtualFile.analysisExtensionFileContextModule = ktModule
+        return ktFile
+    }
+}
+
+private class LLFirResolveExtensionsFileProvider(
+    val extensions: List<KtResolveExtension>,
+) {
+    fun getFilesByPackage(packageFqName: FqName): Sequence<KtResolveExtensionFile> {
+        return extensions
+            .asSequence()
+            .filter { packageFqName in it.getContainedPackages() }
+            .flatMap { it.getKtFiles() }
+            .filter { it.getFilePackageName() == packageFqName }
+    }
+}
+
+private class LLFirResolveExtensionToolPackageProvider(
+    private val packageFilter: LLFirResolveExtensionToolPackageFilter,
+) : KotlinPackageProvider() {
+
+    private val packageSubPackages: Map<FqName, Set<Name>> by lazy {
+        createSubPackagesMapping(packageFilter.allPackages)
+    }
+
+    override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean =
+        doesKotlinOnlyPackageExist(packageFqName)
+
+    override fun getSubPackageFqNames(packageFqName: FqName, platform: TargetPlatform, nameFilter: (Name) -> Boolean): Set<Name> =
+        getKotlinOnlySubPackagesFqNames(packageFqName, nameFilter)
+
+    override fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean = false
+
+    override fun getPlatformSpecificSubPackagesFqNames(packageFqName: FqName, platform: TargetPlatform, nameFilter: (Name) -> Boolean) =
+        emptySet<Name>()
+
+    override fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean =
+        packageFilter.packageExists(packageFqName)
+
+    override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> {
+        val subPackageNames = packageSubPackages[packageFqName] ?: return emptySet()
+        if (subPackageNames.isEmpty()) return emptySet()
+        return subPackageNames.filterTo(mutableSetOf()) { nameFilter(it) }
+    }
+
+    private fun createSubPackagesMapping(packages: Set<FqName>): Map<FqName, Set<Name>> {
+        return buildMap<FqName, MutableSet<Name>> {
+            for (packageName in packages) {
+                collectAllSubPackages(packageName)
+            }
+        }
+    }
+
+    private fun MutableMap<FqName, MutableSet<Name>>.collectAllSubPackages(packageName: FqName) {
+        var currentPackage = FqName.ROOT
+        for (packagePart in packageName.pathSegments()) {
+            getOrPut(currentPackage) { mutableSetOf<Name>() }.add(packagePart)
+            currentPackage = currentPackage.child(packagePart)
+        }
+    }
+}
+
+private fun ClassId.getTopLevelShortClassName(): Name {
+    return Name.guessByFirstCharacter(relativeClassName.asString().substringBefore("."))
+}
+
+private fun KtResolveExtensionFile.mayHaveTopLevelClassifier(name: Name): Boolean {
+    return name in getTopLevelClassifierNames()
+}
+
+private fun KtResolveExtensionFile.mayHaveTopLevelCallable(name: Name): Boolean {
+    return name in getTopLevelCallableNames()
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.ModificationTracker
 import com.intellij.psi.PsiFile
 import com.intellij.psi.SmartPointerManager
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.llResolveExtensionTool
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.providers.KotlinModificationTrackerFactory
 import org.jetbrains.kotlin.analysis.providers.KtModuleStateTracker
@@ -19,6 +20,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.concurrent.atomic.AtomicBoolean
 
 @OptIn(PrivateSessionConstructor::class)
@@ -51,12 +53,13 @@ abstract class LLFirSession(
         }
 
         modificationTracker = CompositeModificationTracker.createFlattened(
-            listOfNotNull(
-                ExplicitInvalidationTracker(ktModule, isExplicitlyInvalidated),
-                ModuleStateModificationTracker(ktModule, validityTracker),
-                outOfBlockTracker,
-                dependencyTracker
-            )
+            buildList {
+                add(ExplicitInvalidationTracker(ktModule, isExplicitlyInvalidated))
+                add(ModuleStateModificationTracker(ktModule, validityTracker))
+                addIfNotNull(outOfBlockTracker)
+                add(dependencyTracker)
+                llResolveExtensionTool?.modificationTrackers?.let(::addAll)
+            }
         )
 
         initialModificationCount = modificationTracker.modificationCount

analysis/project-structure/src/org/jetbrains/kotlin/analysis/project/structure/ktModuleInternalUtils.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2023 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.project.structure
+
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.psi.UserDataProperty
+
+/**
+ * Used by the [org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionProvider] implementations
+ * to store the references on the [KtModule] for which the [VirtualFile] was generated.
+ */
+@KtModuleStructureInternals
+public var VirtualFile.analysisExtensionFileContextModule: KtModule? by UserDataProperty(Key.create("ANALYSIS_CONTEXT_MODULE"))
+
+@RequiresOptIn("Internal KtModule structure API component which should not be used outside the Analysis API implementation modules as it does not have any compatibility guarantees")
+public annotation class KtModuleStructureInternals