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

^KT-58262
2023-04-26 00:19:01 +02:00
parent 06b616957c
commit 96b7abc6cf
8 changed files with 234 additions and 36 deletions
@@ -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 {
@@ -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)
+    }
+}
@@ -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
@@ -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)
@@ -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() }
+    }
+}
@@ -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
@@ -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()
+    }
+}
@@ -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 {
-        forbidAnalysis {
+    private val packageSubPackages: Map<FqName, Set<Name>> by lazy {
+        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 {