[Analysis API] Unwrap anonymous function expressions in reference shortener

Anonymous functions are additionally wrapped into an expression

^KTIJ-26629 Fixed
^KTIJ-26597 Fixed

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

@@ -86,7 +86,7 @@ internal class KtFirReferenceShortener(
         val declarationToVisit = file.findSmallestElementOfTypeContainingSelection<KtDeclaration>(selection)
             ?: file
 
-        val firDeclaration = declarationToVisit.getOrBuildFir(firResolveSession) as? FirDeclaration ?: return ShortenCommandImpl(
+        val firDeclaration = declarationToVisit.getCorrespondingFirDeclaration() ?: return ShortenCommandImpl(
             file.createSmartPointer(),
             importsToAdd = emptySet(),
             starImportsToAdd = emptySet(),
@@ -141,6 +141,18 @@ internal class KtFirReferenceShortener(
         )
     }
 
+    private fun KtElement.getCorrespondingFirDeclaration(): FirDeclaration? {
+        require(this is KtFile || this is KtDeclaration)
+
+        val firElement = getOrBuildFir(firResolveSession)
+
+        return when (firElement) {
+            is FirDeclaration -> firElement
+            is FirAnonymousFunctionExpression -> firElement.anonymousFunction
+            else -> null
+        }
+    }
+
     private fun buildSymbol(firSymbol: FirBasedSymbol<*>): KtSymbol = analysisSession.firSymbolBuilder.buildSymbol(firSymbol)
 }
 
@@ -153,10 +165,13 @@ private fun FqName.dropFakeRootPrefixIfPresent(): FqName {
 
 private data class AdditionalImports(val simpleImports: Set<FqName>, val starImports: Set<FqName>)
 
-private inline fun <reified T : KtElement> KtFile.findSmallestElementOfTypeContainingSelection(selection: TextRange): T? =
+private inline fun <reified T : KtElement> KtFile.findSmallestElementOfTypeContainingSelection(selection: TextRange): T? {
-    findElementAt(selection.startOffset)
+    val parents = findElementAt(selection.startOffset)
-        ?.parentsOfType<T>(withSelf = true)
+        ?.parentsWithSelf
-        ?.firstOrNull { selection in it.textRange }
+        ?.toList()
+
+    return parents?.filterIsInstance<T>()?.firstOrNull { selection in it.textRange }
+}
 
 /**
  * How a symbol is imported. The order of the enum entry represents the priority of imports. If a symbol is available from multiple kinds of

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

@@ -52,6 +52,24 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/annotaiton.kt");
     }
 
+    @Test
+    @TestMetadata("anonymousFunction_annotation.kt")
+    public void testAnonymousFunction_annotation() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_annotation.kt");
+    }
+
+    @Test
+    @TestMetadata("anonymousFunction_receiverType.kt")
+    public void testAnonymousFunction_receiverType() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_receiverType.kt");
+    }
+
+    @Test
+    @TestMetadata("anonymousFunction_returnType.kt")
+    public void testAnonymousFunction_returnType() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_returnType.kt");
+    }
+
     @Test
     @TestMetadata("callInsideScriptExpression.kts")
     public void testCallInsideScriptExpression() throws Exception {

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

@@ -52,6 +52,24 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/annotaiton.kt");
     }
 
+    @Test
+    @TestMetadata("anonymousFunction_annotation.kt")
+    public void testAnonymousFunction_annotation() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_annotation.kt");
+    }
+
+    @Test
+    @TestMetadata("anonymousFunction_receiverType.kt")
+    public void testAnonymousFunction_receiverType() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_receiverType.kt");
+    }
+
+    @Test
+    @TestMetadata("anonymousFunction_returnType.kt")
+    public void testAnonymousFunction_returnType() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_returnType.kt");
+    }
+
     @Test
     @TestMetadata("callInsideScriptExpression.kts")
     public void testCallInsideScriptExpression() throws Exception {

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_annotation.kt

@@ -0,0 +1,9 @@
+// FILE: main.kt
+package test
+
+val anonymous = <expr>@dependency.Ann</expr> fun() {}
+
+// FILE: dependency.kt
+package dependency
+
+annotation class Ann

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_annotation.txt

@@ -0,0 +1,7 @@
+Before shortening: @dependency.Ann
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] dependency.Ann
+with SHORTEN_AND_STAR_IMPORT:
+[type] dependency.Ann

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_receiverType.kt

@@ -0,0 +1,9 @@
+// FILE: main.kt
+package test
+
+val anonymous = fun <expr>dependency.Type</expr>.() {}
+
+// FILE: dependency.kt
+package dependency
+
+class Type

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_receiverType.txt

@@ -0,0 +1,7 @@
+Before shortening: dependency.Type
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] dependency.Type
+with SHORTEN_AND_STAR_IMPORT:
+[type] dependency.Type

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_returnType.kt

@@ -0,0 +1,9 @@
+// FILE: main.kt
+package test
+
+val anonymous = fun(): <expr>dependency.Type</expr> = dependency.Type()
+
+// FILE: dependency.kt
+package dependency
+
+class Type

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/anonymousFunction_returnType.txt

@@ -0,0 +1,7 @@
+Before shortening: dependency.Type
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] dependency.Type
+with SHORTEN_AND_STAR_IMPORT:
+[type] dependency.Type