[AA] Consider context receivers in FirTowerDataElement.getAvailableScope

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

@@ -236,9 +236,10 @@ private class FirShorteningContext(val analysisSession: KtFirAnalysisSession) {
     ): List<FirScope>? {
         val towerDataContext = towerContextProvider.getClosestAvailableParentContext(position) ?: return null
         val result = buildList {
-            addAll(towerDataContext.nonLocalTowerDataElements.mapNotNull {
-                // We must use `it.getAvailableScope()` instead of `it.scope` to check scopes of companion objects as well.
-                it.getAvailableScope()
+            addAll(towerDataContext.nonLocalTowerDataElements.flatMap {
+                // We must use `it.getAvailableScopes()` instead of `it.scope` to check scopes of companion objects
+                // and context receivers as well.
+                it.getAvailableScopes()
             })
             addIfNotNull(createFakeImportingScope(newImports))
             addAll(towerDataContext.localScopes)

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

@@ -160,6 +160,12 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/companionUsedOutOfClass.kt");
     }
 
+    @Test
+    @TestMetadata("contextReceiver.kt")
+    public void testContextReceiver() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/contextReceiver.kt");
+    }
+
     @Test
     @TestMetadata("enumClassCompanionAlreadyImported.kt")
     public void testEnumClassCompanionAlreadyImported() throws Exception {

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

@@ -160,6 +160,12 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/companionUsedOutOfClass.kt");
     }
 
+    @Test
+    @TestMetadata("contextReceiver.kt")
+    public void testContextReceiver() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/contextReceiver.kt");
+    }
+
     @Test
     @TestMetadata("enumClassCompanionAlreadyImported.kt")
     public void testEnumClassCompanionAlreadyImported() throws Exception {

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

@@ -0,0 +1,12 @@
+// FILE: main.kt
+package a.b.c
+
+class A {
+    fun foo() {}
+}
+fun foo() = 1
+
+context(A)
+fun test() {
+    <expr>a.b.c.foo()</expr>
+}

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

@@ -0,0 +1,5 @@
+Before shortening: a.b.c.foo()
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+with SHORTEN_AND_STAR_IMPORT: