[AA] Shorten kdocs that became available because of newly added imports

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

@@ -114,17 +114,23 @@ internal class KtFirReferenceShortener(
         )
         firDeclaration.accept(collector)
 
+        val additionalImports = AdditionalImports(
+            collector.getNamesToImport(starImport = false).toSet(),
+            collector.getNamesToImport(starImport = true).toSet(),
+        )
+
         val kDocQualifiersToShorten = collectKDocQualifiersToShorten(
             file,
             selection,
+            additionalImports,
             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(),
+            additionalImports.simpleImports,
+            additionalImports.starImports,
             collector.typesToShorten.map { it.element }.distinct().map { it.createSmartPointer() },
             collector.qualifiersToShorten.map { it.element }.distinct().map { it.createSmartPointer() },
             kDocQualifiersToShorten.map { it.element }.distinct().map { it.createSmartPointer() },
@@ -134,6 +140,7 @@ internal class KtFirReferenceShortener(
     private fun collectKDocQualifiersToShorten(
         file: KtFile,
         selection: TextRange,
+        additionalImports: AdditionalImports,
         classShortenOption: (KtClassLikeSymbol) -> ShortenOption,
         callableShortenOption: (KtCallableSymbol) -> ShortenOption,
     ): List<ShortenKDocQualifier> {
@@ -157,6 +164,7 @@ internal class KtFirReferenceShortener(
 
                 val shouldShortenKDocQualifier = shouldShortenKDocQualifier(
                     element,
+                    additionalImports,
                     classShortenOption = { classShortenOption(buildSymbol(it) as KtClassLikeSymbol) },
                     callableShortenOption = { callableShortenOption(buildSymbol(it) as KtCallableSymbol) },
                 )
@@ -177,11 +185,15 @@ internal class KtFirReferenceShortener(
 
     private fun shouldShortenKDocQualifier(
         kDocName: KDocName,
+        additionalImports: AdditionalImports,
         classShortenOption: (FirClassLikeSymbol<*>) -> ShortenOption,
         callableShortenOption: (FirCallableSymbol<*>) -> ShortenOption,
     ): Boolean {
         val fqName = kDocName.getQualifiedNameAsFqName().dropFakeRootPrefixIfPresent()
 
+        // KDocs are only shortened if they are available without imports, so `additionalImports` contain all the imports to add
+        if (fqName.isInNewImports(additionalImports)) return true
+
         val resolvedSymbols = with(analysisSession) {
             val shortFqName = FqName.topLevel(fqName.shortName())
             val owner = kDocName.getContainingDoc().owner
@@ -200,6 +212,9 @@ internal class KtFirReferenceShortener(
         return false
     }
 
+    private fun FqName.isInNewImports(additionalImports: AdditionalImports): Boolean =
+        this in additionalImports.simpleImports || this.parent() in additionalImports.starImports
+
     private fun canShorten(fqNameToShorten: FqName, fqNameOfAvailableSymbol: FqName, getShortenOption: () -> ShortenOption): Boolean =
         fqNameToShorten == fqNameOfAvailableSymbol && getShortenOption() != ShortenOption.DO_NOT_SHORTEN
 
@@ -213,6 +228,8 @@ private fun FqName.dropFakeRootPrefixIfPresent(): FqName {
     } else this
 }
 
+private data class AdditionalImports(val simpleImports: Set<FqName>, val starImports: Set<FqName>)
+
 private inline fun <reified T : KtElement> KtFile.findSmallestElementOfTypeContainingSelection(selection: TextRange): T? =
     findElementAt(selection.startOffset)
         ?.parentsOfType<T>(withSelf = true)

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisSourceModuleReferenceShortenerForWholeFileTestGenerated.java

@@ -46,6 +46,12 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerForWholeFileTestG
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("fileWithKDoc.kt")
+    public void testFileWithKDoc() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithKDoc.kt");
+    }
+
     @Test
     @TestMetadata("fileWithMultipleDeclarations.kt")
     public void testFileWithMultipleDeclarations() throws Exception {

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/references/FirStandaloneNormalAnalysisSourceModuleReferenceShortenerForWholeFileTestGenerated.java

@@ -46,6 +46,12 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerForWholeFi
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("fileWithKDoc.kt")
+    public void testFileWithKDoc() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithKDoc.kt");
+    }
+
     @Test
     @TestMetadata("fileWithMultipleDeclarations.kt")
     public void testFileWithMultipleDeclarations() throws Exception {

analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithKDoc.kt

@@ -0,0 +1,14 @@
+// FILE: main.kt
+
+/**
+ * [a.b.c.dependency.Foo]
+ * [a.b.c.dependency.Foo.Nested]
+ */
+fun foo<caret_onAirContext>(): a.b.c.dependency.Foo = t
+
+// FILE: dependency.kt
+package a.b.c.dependency
+
+class Foo {
+    class Nested
+}

analysis/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/fileWithKDoc.txt

@@ -0,0 +1,10 @@
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] a.b.c.dependency.Foo
+[kdoc] a.b.c.dependency.Foo
+[kdoc] a.b.c.dependency.Foo
+with SHORTEN_AND_STAR_IMPORT:
+[type] a.b.c.dependency.Foo
+[kdoc] a.b.c.dependency.Foo
+[kdoc] a.b.c.dependency.Foo