[Analysis API] add API to get KtScope for resolve extensions

^KT-58262

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

@@ -65,6 +65,7 @@ class KtFe10AnalysisSession(
     override val scopeSubstitutionImpl: KtScopeSubstitution = KtFe10ScopeSubstitution(this)
     override val substitutorFactoryImpl: KtSubstitutorFactory = KtFe10SubstitutorFactory(this)
     override val symbolProviderByJavaPsiImpl: KtSymbolProviderByJavaPsi = KtFe10SymbolProviderByJavaPsi(this)
+    override val resolveExtensionProviderImpl: KtSymbolFromResolveExtensionProvider = KtFe10SymbolFromResolveExtensionProvider(this)
 
     override fun createContextDependentCopy(originalKtFile: KtFile, elementToReanalyze: KtElement): KtAnalysisSession =
         withValidityAssertion {

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

@@ -0,0 +1,25 @@
+/*
+ * 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.descriptors.components
+
+import org.jetbrains.kotlin.analysis.api.components.KtSymbolFromResolveExtensionProvider
+import org.jetbrains.kotlin.analysis.api.descriptors.KtFe10AnalysisSession
+import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnalysisSessionComponent
+import org.jetbrains.kotlin.analysis.api.impl.base.scopes.KtEmptyScope
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.scopes.KtScope
+
+internal class KtFe10SymbolFromResolveExtensionProvider(
+    override val analysisSession: KtFe10AnalysisSession,
+) : KtSymbolFromResolveExtensionProvider(), Fe10KtAnalysisSessionComponent {
+    override val token: KtLifetimeToken
+        get() = analysisSession.token
+
+    override fun getResolveExtensionScopeWithTopLevelDeclarations(): KtScope {
+        // Not supported for FE1.0
+        return KtEmptyScope(token)
+    }
+}

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSession.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
 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.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.allDirectDependencies
@@ -123,6 +124,8 @@ private constructor(
 
     override val symbolProviderByJavaPsiImpl = KtFirSymbolProviderByJavaPsi(this)
 
+    override val resolveExtensionProviderImpl: KtSymbolFromResolveExtensionProvider = KtFirSymbolFromResolveExtensionProvider(this)
+
     @Suppress("AnalysisApiMissingLifetimeCheck")
     override fun createContextDependentCopy(originalKtFile: KtFile, elementToReanalyze: KtElement): KtAnalysisSession {
         check(mode == AnalysisSessionMode.REGULAR) {
@@ -150,11 +153,13 @@ private constructor(
 
     val useSiteAnalysisScope: GlobalSearchScope = analysisScopeProviderImpl.getAnalysisScope()
 
+    val extensionTools: List<LLFirResolveExtensionTool>
     val useSiteScopeDeclarationProvider: KotlinDeclarationProvider
     val useSitePackageProvider: KotlinPackageProvider
 
+
     init {
-        val extensionTools = buildList {
+        extensionTools = buildList {
             addIfNotNull(useSiteSession.llResolveExtensionTool)
             useSiteModule.allDirectDependencies().mapNotNullTo(this) { dependency ->
                 firResolveSession.getSessionFor(dependency).llResolveExtensionTool

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtSymbolByFirBuilder.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.api.fir
 
 import com.intellij.openapi.project.Project
-import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.analysis.api.KtStarTypeProjection
 import org.jetbrains.kotlin.analysis.api.KtTypeArgumentWithVariance
@@ -25,7 +24,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecific
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withConeTypeEntry
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirEntry
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirSymbolEntry
-import org.jetbrains.kotlin.analysis.providers.createPackageProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -106,7 +105,7 @@ internal class KtSymbolByFirBuilder constructor(
 
     fun buildScriptSymbol(firSymbol: FirScriptSymbol) = KtFirScriptSymbol(firSymbol, analysisSession)
 
-    private val packageProvider = project.createPackageProvider(GlobalSearchScope.allScope(project))//todo scope
+    private val packageProvider: KotlinPackageProvider get() = analysisSession.useSitePackageProvider
 
     fun createPackageSymbolIfOneExists(packageFqName: FqName): KtFirPackageSymbol? {
         val exists = packageProvider.doesPackageExist(packageFqName, analysisSession.targetPlatform)

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolFromResolveExtensionProvider.kt

@@ -0,0 +1,90 @@
+/*
+ * 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.fir.components
+
+import org.jetbrains.kotlin.analysis.api.components.KtSymbolFromResolveExtensionProvider
+import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.analysis.api.impl.base.scopes.KtEmptyScope
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.scopes.KtScope
+import org.jetbrains.kotlin.analysis.api.scopes.KtScopeNameFilter
+import org.jetbrains.kotlin.analysis.api.symbols.*
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionTool
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionToolDeclarationProvider
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtNamedDeclaration
+
+internal class KtFirSymbolFromResolveExtensionProvider(
+    override val analysisSession: KtFirAnalysisSession,
+) : KtSymbolFromResolveExtensionProvider(), KtFirAnalysisSessionComponent {
+    override val token: KtLifetimeToken
+        get() = analysisSession.token
+
+    override fun getResolveExtensionScopeWithTopLevelDeclarations(): KtScope {
+        val tools = analysisSession.extensionTools
+        if (tools.isEmpty()) return KtEmptyScope(token)
+        return KtFirResolveExtensionScope(analysisSession, tools)
+    }
+}
+
+private class KtFirResolveExtensionScope(
+    private val analysisSession: KtFirAnalysisSession,
+    private val tools: List<LLFirResolveExtensionTool>,
+) : KtScope {
+    init {
+        require(tools.isNotEmpty())
+    }
+
+    override val token: KtLifetimeToken get() = analysisSession.token
+
+    override fun getCallableSymbols(nameFilter: KtScopeNameFilter): Sequence<KtCallableSymbol> = withValidityAssertion {
+        gelTopLevelDeclarations(nameFilter) { it.getTopLevelCallables() }
+    }
+
+    override fun getClassifierSymbols(nameFilter: KtScopeNameFilter): Sequence<KtClassifierSymbol> = withValidityAssertion {
+        gelTopLevelDeclarations(nameFilter) { it.getTopLevelClassifiers() }
+    }
+
+    private inline fun <D : KtNamedDeclaration, reified S : KtDeclaration> gelTopLevelDeclarations(
+        crossinline nameFilter: KtScopeNameFilter,
+        crossinline getDeclarationsByProvider: (LLFirResolveExtensionToolDeclarationProvider) -> Sequence<D>,
+    ): Sequence<S> = sequence {
+        for (tool in tools) {
+            for (declaration in getDeclarationsByProvider(tool.declarationProvider)) {
+                val declarationName = declaration.nameAsName ?: continue
+                if (!nameFilter(declarationName)) continue
+                with(analysisSession) {
+                    yield(declaration.getSymbolOfType())
+                }
+            }
+        }
+    }
+
+    override fun getConstructors(): Sequence<KtConstructorSymbol> = withValidityAssertion {
+        emptySequence()
+    }
+
+    override fun getPackageSymbols(nameFilter: KtScopeNameFilter): Sequence<KtPackageSymbol> = withValidityAssertion {
+        sequence {
+            for (tool in tools) {
+                for (packageName in tool.packageFilter.getAllPackages()) {
+                    if (!nameFilter(packageName.shortName())) continue
+                    analysisSession.firSymbolBuilder.createPackageSymbol(packageName)
+                }
+            }
+        }
+    }
+
+    override fun getPossibleCallableNames(): Set<Name> = withValidityAssertion {
+        tools.flatMapTo(mutableSetOf()) { it.declarationProvider.getTopLevelCallableNames() }
+    }
+
+    override fun getPossibleClassifierNames(): Set<Name> = withValidityAssertion {
+        tools.flatMapTo(mutableSetOf()) { it.declarationProvider.getTopLevelClassifierNames() }
+    }
+}

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

@@ -60,7 +60,8 @@ public abstract class KtAnalysisSession(final override val token: KtLifetimeToke
     KtAnalysisScopeProviderMixIn,
     KtSignatureSubstitutorMixIn,
     KtScopeSubstitutionMixIn,
-    KtSymbolProviderByJavaPsiMixIn {
+    KtSymbolProviderByJavaPsiMixIn,
+    KtSymbolFromResolveExtensionProviderMixIn {
 
     public abstract val useSiteModule: KtModule
 
@@ -155,6 +156,9 @@ public abstract class KtAnalysisSession(final override val token: KtLifetimeToke
     internal val scopeSubstitution: KtScopeSubstitution get() = scopeSubstitutionImpl
     protected abstract val scopeSubstitutionImpl: KtScopeSubstitution
 
+    internal val resolveExtensionProvider: KtSymbolFromResolveExtensionProvider get() = resolveExtensionProviderImpl
+    protected abstract val resolveExtensionProviderImpl: KtSymbolFromResolveExtensionProvider
+
     @KtAnalysisApiInternals
     public val substitutorFactory: KtSubstitutorFactory get() = substitutorFactoryImpl
     protected abstract val substitutorFactoryImpl: KtSubstitutorFactory

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSymbolFromResolveExtensionProvider.kt

@@ -0,0 +1,25 @@
+/*
+ * 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.analysis.api.components
+
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.scopes.KtScope
+
+public abstract class KtSymbolFromResolveExtensionProvider : KtAnalysisSessionComponent() {
+    abstract fun getResolveExtensionScopeWithTopLevelDeclarations(): KtScope
+}
+
+public interface KtSymbolFromResolveExtensionProviderMixIn : KtAnalysisSessionMixIn {
+    /**
+     * Returns [KtScope] which contains all top-level callable declarations which are generated by [org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtension]
+     *
+     * @see org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtension
+     * @see org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionProvider
+     */
+    fun getResolveExtensionScopeWithTopLevelDeclarations(): KtScope = withValidityAssertion {
+        analysisSession.resolveExtensionProvider.getResolveExtensionScopeWithTopLevelDeclarations()
+    }
+}

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

@@ -36,8 +36,9 @@ import java.util.concurrent.ConcurrentHashMap
  */
 abstract class LLFirResolveExtensionTool : FirSessionComponent {
     abstract val modificationTrackers: List<ModificationTracker>
-    abstract val declarationProvider: KotlinDeclarationProvider
+    abstract val declarationProvider: LLFirResolveExtensionToolDeclarationProvider
     abstract val packageProvider: KotlinPackageProvider
+    abstract val packageFilter: LLFirResolveExtensionToolPackageFilter
     internal abstract val symbolNameCache: LLFirSymbolProviderNameCache
 }
 
@@ -53,11 +54,11 @@ internal class LLFirNonEmptyResolveExtensionTool(
 
     private val fileProvider = LLFirResolveExtensionsFileProvider(extensions)
 
-    private val packageFilter = LLFirResolveExtensionToolPackageFilter(extensions)
+    override val packageFilter = LLFirResolveExtensionToolPackageFilter(extensions)
 
     override val modificationTrackers by lazy { extensions.map { it.getModificationTracker() } }
 
-    override val declarationProvider: KotlinDeclarationProvider =
+    override val declarationProvider: LLFirResolveExtensionToolDeclarationProvider =
         LLFirResolveExtensionToolDeclarationProvider(fileProvider, session.ktModule)
 
     override val packageProvider: KotlinPackageProvider = LLFirResolveExtensionToolPackageProvider(packageFilter)
@@ -97,21 +98,47 @@ private class LLFirResolveExtensionToolNameCache(
     }
 }
 
-private class LLFirResolveExtensionToolPackageFilter(
+class LLFirResolveExtensionToolPackageFilter(
     private val extensions: List<KtResolveExtension>
 ) {
-    val allPackages: Set<FqName> by lazy {
+    private val packageSubPackages: Map<FqName, Set<Name>> by lazy {
-        forbidAnalysis {
+        val packagesFromExtensions = forbidAnalysis {
             extensions.flatMapTo(mutableSetOf()) { it.getContainedPackages() }
         }
+        createSubPackagesMapping(packagesFromExtensions)
+    }
+
+    fun getAllPackages(): Set<FqName> {
+        return packageSubPackages.keys
+    }
+
+    fun getAllSubPackages(packageFqName: FqName): Set<Name> {
+        return packageSubPackages[packageFqName].orEmpty()
     }
 
     fun packageExists(packageFqName: FqName): Boolean {
-        return packageFqName in allPackages
+        return packageFqName in packageSubPackages
+    }
+
+    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)
+        }
+        putIfAbsent(currentPackage, mutableSetOf())
     }
 }
 
-private class LLFirResolveExtensionToolDeclarationProvider(
+class LLFirResolveExtensionToolDeclarationProvider internal constructor(
     private val extensionProvider: LLFirResolveExtensionsFileProvider,
     private val ktModule: KtModule,
 ) : KotlinDeclarationProvider() {
@@ -119,6 +146,30 @@ private class LLFirResolveExtensionToolDeclarationProvider(
     private val extensionFileToDeclarationProvider: ConcurrentHashMap<KtResolveExtensionFile, FileBasedKotlinDeclarationProvider> =
         ConcurrentHashMap()
 
+    fun getTopLevelCallables(): Sequence<KtCallableDeclaration> = sequence {
+        forEachDeclarationOfType<KtCallableDeclaration> { callable ->
+            yield(callable)
+        }
+    }
+
+    fun getTopLevelClassifiers(): Sequence<KtClassLikeDeclaration> = sequence {
+        forEachDeclarationOfType<KtClass> { classLike ->
+            yield(classLike)
+        }
+    }
+
+    fun getTopLevelCallableNames(): Sequence<Name> = sequence {
+        forEachDeclarationOfType<KtCallableDeclaration> { callable ->
+            callable.nameAsName?.let { yield(it) }
+        }
+    }
+
+    fun getTopLevelClassifierNames(): Sequence<Name> = sequence {
+        forEachDeclarationOfType<KtClassLikeDeclaration> { classLike ->
+            classLike.nameAsName?.let { yield(it) }
+        }
+    }
+
     override fun getClassLikeDeclarationByClassId(classId: ClassId): KtClassLikeDeclaration? = forbidAnalysis {
         return getDeclarationProvidersByPackage(classId.packageFqName) { it.mayHaveTopLevelClassifier(classId.getTopLevelShortClassName()) }
             .firstNotNullOfOrNull { it.getClassLikeDeclarationByClassId(classId) }
@@ -224,9 +275,22 @@ private class LLFirResolveExtensionToolDeclarationProvider(
         virtualFile.psiTargetsProvider = psiTargetsProvider
         return ktFile
     }
+
+
+    private inline fun <reified D : KtDeclaration> forEachDeclarationOfType(action: (D) -> Unit) {
+        for (file in extensionProvider.getAllFiles()) {
+            val provider = createDeclarationProviderByFile(file)
+            val ktFile = provider.kotlinFile
+            for (declaration in ktFile.declarations) {
+                if (declaration is D) {
+                    action(declaration)
+                }
+            }
+        }
+    }
 }
 
-private class LLFirResolveExtensionsFileProvider(
+internal class LLFirResolveExtensionsFileProvider(
     val extensions: List<KtResolveExtension>,
 ) {
     fun getFilesByPackage(packageFqName: FqName): Sequence<KtResolveExtensionFile> = forbidAnalysis {
@@ -236,16 +300,17 @@ private class LLFirResolveExtensionsFileProvider(
             .flatMap { it.getKtFiles() }
             .filter { it.getFilePackageName() == packageFqName }
     }
+
+    fun getAllFiles(): Sequence<KtResolveExtensionFile> {
+        return extensions
+            .asSequence()
+            .flatMap { it.getKtFiles() }
+    }
 }
 
 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)
 
@@ -261,26 +326,10 @@ private class LLFirResolveExtensionToolPackageProvider(
         packageFilter.packageExists(packageFqName)
 
     override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> {
-        val subPackageNames = packageSubPackages[packageFqName] ?: return emptySet()
+        val subPackageNames = packageFilter.getAllSubPackages(packageFqName)
         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 {