[Analysis API] Shortener: enable shortening in KDoc

^KTIJ-21103
2023-05-16 09:51:53 +03:00
parent e7c213e06e
commit 67c3849538
15 changed files with 230 additions and 20 deletions
@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnaly
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
+import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtFile
@@ -39,10 +40,11 @@ internal class KtFe10ReferenceShortener(

        return object : ShortenCommand {
            override val targetFile: SmartPsiElementPointer<KtFile> get() = ktFilePointer
-            override val importsToAdd: List<FqName> get() = emptyList()
-            override val starImportsToAdd: List<FqName> get() = emptyList()
+            override val importsToAdd: Set<FqName> get() = emptySet()
+            override val starImportsToAdd: Set<FqName> get() = emptySet()
            override val typesToShorten: List<SmartPsiElementPointer<KtUserType>> get() = emptyList()
            override val qualifiersToShorten: List<SmartPsiElementPointer<KtDotQualifiedExpression>> get() = emptyList()
+            override val kDocQualifiersToShorten: List<SmartPsiElementPointer<KDocName>> get() = emptyList()

            override val isEmpty: Boolean get() = true
        }
@@ -15,11 +15,14 @@ import org.jetbrains.kotlin.analysis.api.components.ShortenOption
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.components.ElementsToShortenCollector.PartialOrderOfScope.Companion.toPartialOrder
 import org.jetbrains.kotlin.analysis.api.fir.isImplicitDispatchReceiver
+import org.jetbrains.kotlin.analysis.api.fir.references.KDocReferenceResolver
 import org.jetbrains.kotlin.analysis.api.fir.utils.FirBodyReanalyzingVisitorVoid
 import org.jetbrains.kotlin.analysis.api.fir.utils.computeImportableName
+import org.jetbrains.kotlin.analysis.api.fir.utils.firSymbol
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 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.api.getOrBuildFir
@@ -56,11 +59,13 @@ import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull

 internal class KtFirReferenceShortener(
    override val analysisSession: KtFirAnalysisSession,
@@ -77,15 +82,16 @@ internal class KtFirReferenceShortener(
    ): ShortenCommand {
        require(!file.isCompiled) { "No sense to collect references for shortening in compiled file $file" }

-        val declarationToVisit = file.findSmallestDeclarationContainingSelection(selection)
+        val declarationToVisit = file.findSmallestElementOfTypeContainingSelection<KtDeclaration>(selection)
            ?: file

        val firDeclaration = declarationToVisit.getOrBuildFir(firResolveSession) as? FirDeclaration ?: return ShortenCommandImpl(
            file.createSmartPointer(),
-            importsToAdd = emptyList(),
-            starImportsToAdd = emptyList(),
+            importsToAdd = emptySet(),
+            starImportsToAdd = emptySet(),
            typesToShorten = emptyList(),
            qualifiersToShorten = emptyList(),
+            kDocQualifiersToShorten = emptyList(),
        )

        val towerContext = when (declarationToVisit) {
@@ -102,25 +108,114 @@ internal class KtFirReferenceShortener(
            context,
            towerContext,
            selection,
-            classShortenOption = { classShortenOption(analysisSession.firSymbolBuilder.buildSymbol(it) as KtClassLikeSymbol) },
-            callableShortenOption = { callableShortenOption(analysisSession.firSymbolBuilder.buildSymbol(it) as KtCallableSymbol) },
+            classShortenOption = { classShortenOption(buildSymbol(it) as KtClassLikeSymbol) },
+            callableShortenOption = { callableShortenOption(buildSymbol(it) as KtCallableSymbol) },
            firResolveSession,
        )
        firDeclaration.accept(collector)

+        val kDocQualifiersToShorten = collectKDocQualifiersToShorten(
+            file,
+            selection,
+            classShortenOption = { minOf(classShortenOption(it), ShortenOption.SHORTEN_IF_ALREADY_IMPORTED) },
+            callableShortenOption = { minOf(callableShortenOption(it), ShortenOption.SHORTEN_IF_ALREADY_IMPORTED) },
+        )
+
        return ShortenCommandImpl(
            file.createSmartPointer(),
            collector.getNamesToImport(starImport = false).toList(),
            collector.getNamesToImport(starImport = true).toList(),
            collector.typesToShorten.map { it.element }.distinct().map { it.createSmartPointer() },
-            collector.qualifiersToShorten.map { it.element }.distinct().map { it.createSmartPointer() }
+            collector.qualifiersToShorten.map { it.element }.distinct().map { it.createSmartPointer() },
+            kDocQualifiersToShorten.map { it.element }.distinct().map { it.createSmartPointer() },
        )
    }
+
+    private fun collectKDocQualifiersToShorten(
+        file: KtFile,
+        selection: TextRange,
+        classShortenOption: (KtClassLikeSymbol) -> ShortenOption,
+        callableShortenOption: (KtCallableSymbol) -> ShortenOption,
+    ): List<ShortenKDocQualifier> {
+        val kDocQualifiersToShorten = mutableListOf<ShortenKDocQualifier>()
+        val elementToVisit = file.findSmallestElementOfTypeContainingSelection<KtElement>(selection)
+
+        fun addKDocToShorten(kDocName: KDocName) {
+            kDocQualifiersToShorten.add(ShortenKDocQualifier(kDocName))
+        }
+
+        elementToVisit?.accept(object : KtVisitorVoid() {
+            override fun visitElement(element: PsiElement) {
+                if (!element.textRange.intersects(selection)) return
+
+                if (!selection.contains(element.textRange) || element !is KDocName) {
+                    element.acceptChildren(this)
+                    return
+                }
+
+                if (element.getQualifier() == null) return
+
+                val shouldShortenKDocQualifier = shouldShortenKDocQualifier(
+                    element,
+                    classShortenOption = { classShortenOption(buildSymbol(it) as KtClassLikeSymbol) },
+                    callableShortenOption = { callableShortenOption(buildSymbol(it) as KtCallableSymbol) },
+                )
+                if (shouldShortenKDocQualifier) {
+                    addKDocToShorten(element)
+                } else {
+                    element.acceptChildren(this)
+
+                    if (element.getQualifier()?.getNameText() == ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE) {
+                        addKDocToShorten(element)
+                    }
+                }
+            }
+        })
+
+        return kDocQualifiersToShorten
+    }
+
+    private fun shouldShortenKDocQualifier(
+        kDocName: KDocName,
+        classShortenOption: (FirClassLikeSymbol<*>) -> ShortenOption,
+        callableShortenOption: (FirCallableSymbol<*>) -> ShortenOption,
+    ): Boolean {
+        val fqName = kDocName.getQualifiedNameAsFqName().dropFakeRootPrefixIfPresent()
+
+        val resolvedSymbols = with(analysisSession) {
+            val shortFqName = FqName.topLevel(fqName.shortName())
+            val owner = kDocName.getContainingDoc().owner
+
+            KDocReferenceResolver.resolveKdocFqName(shortFqName, shortFqName, owner ?: kDocName.containingKtFile)
+        }
+
+        resolvedSymbols.firstIsInstanceOrNull<KtCallableSymbol>()?.firSymbol?.let { availableCallable ->
+            return canShorten(fqName, availableCallable.callableId.asSingleFqName()) { callableShortenOption(availableCallable) }
+        }
+
+        resolvedSymbols.firstIsInstanceOrNull<KtClassLikeSymbol>()?.firSymbol?.let { availableClassifier ->
+            return canShorten(fqName, availableClassifier.classId.asSingleFqName()) { classShortenOption(availableClassifier) }
+        }
+
+        return false
+    }
+
+    private fun canShorten(fqNameToShorten: FqName, fqNameOfAvailableSymbol: FqName, getShortenOption: () -> ShortenOption): Boolean =
+        fqNameToShorten == fqNameOfAvailableSymbol && getShortenOption() != ShortenOption.DO_NOT_SHORTEN
+
+    private fun buildSymbol(firSymbol: FirBasedSymbol<*>): KtSymbol = analysisSession.firSymbolBuilder.buildSymbol(firSymbol)
 }

-private fun KtFile.findSmallestDeclarationContainingSelection(selection: TextRange): KtDeclaration? =
+private fun FqName.dropFakeRootPrefixIfPresent(): FqName {
+    val pathSegments = pathSegments()
+    return if (pathSegments.firstOrNull()?.asString() == ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE) {
+        FqName.fromSegments(pathSegments.drop(1).map { it.asString() })
+    } else this
+}
+
+private inline fun <reified T : KtElement> KtFile.findSmallestElementOfTypeContainingSelection(selection: TextRange): T? =
    findElementAt(selection.startOffset)
-        ?.parentsOfType<KtDeclaration>(withSelf = true)
+        ?.parentsOfType<T>(withSelf = true)
        ?.firstOrNull { selection in it.textRange }

 /**
@@ -301,7 +396,7 @@ private sealed class ElementToShorten {
 private class ShortenType(
    val element: KtUserType,
    override val nameToImport: FqName? = null,
-    override val importAllInParent: Boolean = false
+    override val importAllInParent: Boolean = false,
 ) : ElementToShorten()

 private class ShortenQualifier(
@@ -310,6 +405,12 @@ private class ShortenQualifier(
    override val importAllInParent: Boolean = false
 ) : ElementToShorten()

+private class ShortenKDocQualifier(
+    val element: KDocName,
+    override val nameToImport: FqName? = null,
+    override val importAllInParent: Boolean = false
+) : ElementToShorten()
+
 private class ElementsToShortenCollector(
    private val shorteningContext: FirShorteningContext,
    private val towerContextProvider: FirTowerContextProvider,
@@ -1098,6 +1199,7 @@ private class ElementsToShortenCollector(
        when (elementInfoToShorten) {
            is ShortenType -> addElementToShorten(elementInfoToShorten.element, nameToImport, isImportWithStar)
            is ShortenQualifier -> addElementToShorten(elementInfoToShorten.element, nameToImport, isImportWithStar)
+            is ShortenKDocQualifier -> addElementToShorten(elementInfoToShorten.element, nameToImport, isImportWithStar)
        }
    }

@@ -1124,10 +1226,11 @@ private class ElementsToShortenCollector(

 private class ShortenCommandImpl(
    override val targetFile: SmartPsiElementPointer<KtFile>,
-    override val importsToAdd: List<FqName>,
-    override val starImportsToAdd: List<FqName>,
+    override val importsToAdd: Set<FqName>,
+    override val starImportsToAdd: Set<FqName>,
    override val typesToShorten: List<SmartPsiElementPointer<KtUserType>>,
    override val qualifiersToShorten: List<SmartPsiElementPointer<KtDotQualifiedExpression>>,
+    override val kDocQualifiersToShorten: List<SmartPsiElementPointer<KDocName>>,
 ) : ShortenCommand

 private fun KtUserType.hasFakeRootPrefix(): Boolean =
@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.scopes.impl.importedFromObjectOrStaticData
 import org.jetbrains.kotlin.fir.scopes.impl.originalForWrappedIntegerOperator
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol

 internal interface KtFirSymbol<out S : FirBasedSymbol<*>> : KtSymbol, KtLifetimeOwner {
@@ -106,7 +105,7 @@ internal tailrec fun FirDeclaration.ktSymbolOrigin(): KtSymbolOrigin = when (ori
 }

 internal fun KtClassLikeSymbol.getSymbolKind(): KtSymbolKind {
-    val firSymbol = firSymbol as FirClassLikeSymbol<*>
+    val firSymbol = firSymbol
    if (firSymbol.isLocal) {
        // TODO: hack should be dropped after KT-54390
        when {
@@ -30,6 +30,7 @@ internal val KtEnumEntrySymbol.firSymbol: FirEnumEntrySymbol get() = (this as Kt
 internal val KtConstructorSymbol.firSymbol: FirConstructorSymbol get() = (this as KtFirSymbol<*>).firSymbol as FirConstructorSymbol
 internal val KtPropertyAccessorSymbol.firSymbol: FirPropertyAccessorSymbol get() = (this as KtFirSymbol<*>).firSymbol as FirPropertyAccessorSymbol
 internal val KtClassInitializerSymbol.firSymbol: FirAnonymousInitializerSymbol get() = (this as KtFirSymbol<*>).firSymbol as FirAnonymousInitializerSymbol
+internal val KtClassLikeSymbol.firSymbol: FirClassLikeSymbol<*> get() = (this as KtFirSymbol<*>).firSymbol as FirClassLikeSymbol<*>


 fun FirBasedSymbol<*>.getContainingKtModule(firResolveSession: LLFirResolveSession): KtModule {
@@ -238,6 +238,24 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/extensionFromObject.kt");
    }

+    @Test
+    @TestMetadata("kdoc.kt")
+    public void testKdoc() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/kdoc.kt");
+    }
+
+    @Test
+    @TestMetadata("kdocQualifierSelected.kt")
+    public void testKdocQualifierSelected() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/kdocQualifierSelected.kt");
+    }
+
+    @Test
+    @TestMetadata("kdocUnresolved.kt")
+    public void testKdocUnresolved() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/kdocUnresolved.kt");
+    }
+
    @Test
    @TestMetadata("multipleImport.kt")
    public void testMultipleImport() throws Exception {
@@ -24,5 +24,10 @@ internal object ShorteningResultsRenderer {
                appendLine("[qualifier] $it")
            }
        }
+        shortening.kDocQualifiersToShorten.forEach { kdoc ->
+            kdoc.element?.text?.let {
+                appendLine("[kdoc] $it")
+            }
+        }
    }
 }
@@ -238,6 +238,24 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/extensionFromObject.kt");
    }

+    @Test
+    @TestMetadata("kdoc.kt")
+    public void testKdoc() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/kdoc.kt");
+    }
+
+    @Test
+    @TestMetadata("kdocQualifierSelected.kt")
+    public void testKdocQualifierSelected() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/kdocQualifierSelected.kt");
+    }
+
+    @Test
+    @TestMetadata("kdocUnresolved.kt")
+    public void testKdocUnresolved() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/kdocUnresolved.kt");
+    }
+
    @Test
    @TestMetadata("multipleImport.kt")
    public void testMultipleImport() throws Exception {
@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtElement
@@ -86,7 +87,7 @@ public interface KtReferenceShortenerMixIn : KtAnalysisSessionMixIn {

    /**
     * Collects possible references to shorten. By default, it shortens a fully-qualified members to the outermost class and does not
-     * shorten enum entries.
+     * shorten enum entries.  In case of KDoc shortens reference only if it is already imported.
     *
     * N.B. This API is not implemented for the FE10 implementation!
     * For a K1- and K2-compatible API, use [org.jetbrains.kotlin.idea.base.codeInsight.ShortenReferencesFacility].
@@ -134,11 +135,12 @@ public interface KtReferenceShortenerMixIn : KtAnalysisSessionMixIn {

 public interface ShortenCommand {
    public val targetFile: SmartPsiElementPointer<KtFile>
-    public val importsToAdd: List<FqName>
-    public val starImportsToAdd: List<FqName>
+    public val importsToAdd: Set<FqName>
+    public val starImportsToAdd: Set<FqName>
    public val typesToShorten: List<SmartPsiElementPointer<KtUserType>>
    public val qualifiersToShorten: List<SmartPsiElementPointer<KtDotQualifiedExpression>>
+    public val kDocQualifiersToShorten: List<SmartPsiElementPointer<KDocName>>

    public val isEmpty: Boolean
-        get() = typesToShorten.isEmpty() && qualifiersToShorten.isEmpty()
+        get() = typesToShorten.isEmpty() && qualifiersToShorten.isEmpty() && kDocQualifiersToShorten.isEmpty()
 }
@@ -0,0 +1,12 @@
+// FILE: main.kt
+import a.b.c.dependency.Foo
+
+/**
+ * [<expr>a.b.c.dependency.Foo</expr>]
+ */
+fun test() {}
+
+// FILE: dependency.kt
+package a.b.c.dependency
+
+class Foo
@@ -0,0 +1,8 @@
+Before shortening: a.b.c.dependency.Foo
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[kdoc] a.b.c.dependency.Foo
+with SHORTEN_AND_IMPORT:
+[kdoc] a.b.c.dependency.Foo
+with SHORTEN_AND_STAR_IMPORT:
+[kdoc] a.b.c.dependency.Foo
@@ -0,0 +1,14 @@
+// FILE: main.kt
+import dependency.Foo
+
+/**
+ * [<expr>dependency.Foo</expr>.foo]
+ */
+fun test() {}
+
+// FILE: dependency.kt
+package dependency
+
+object Foo {
+    fun foo() {}
+}
@@ -0,0 +1,8 @@
+Before shortening: dependency.Foo
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[kdoc] dependency.Foo
+with SHORTEN_AND_IMPORT:
+[kdoc] dependency.Foo
+with SHORTEN_AND_STAR_IMPORT:
+[kdoc] dependency.Foo
@@ -0,0 +1,12 @@
+// FILE: main.kt
+import dependency.Foo
+
+/**
+ * [<expr>dependency.Foo.unresolved</expr>]
+ */
+fun test() {}
+
+// FILE: dependency.kt
+package dependency
+
+class Foo
@@ -0,0 +1,8 @@
+Before shortening: dependency.Foo.unresolved
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[kdoc] dependency.Foo
+with SHORTEN_AND_IMPORT:
+[kdoc] dependency.Foo
+with SHORTEN_AND_STAR_IMPORT:
+[kdoc] dependency.Foo
@@ -134,7 +134,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
            FileTowerProvider(place, onAirGetTowerContextForFile(firResolveSession, place))
        } else {
            val validPlace = PsiTreeUtil.findFirstParent(place, false) {
-                RawFirReplacement.isApplicableForReplacement(it as KtElement)
+                it is KtElement && RawFirReplacement.isApplicableForReplacement(it)
            } as KtElement

            FirTowerDataContextAllElementsCollector().also {