[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:

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/calls/FirReceivers.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildThisReceiverExpression
 import org.jetbrains.kotlin.fir.references.builder.buildImplicitThisReference
 import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.resolve.smartcastScope
 import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
@@ -29,6 +30,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.SmartcastStability
+import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 
 interface Receiver
 
@@ -88,6 +90,8 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
 
     val originalType: ConeKotlinType = type
 
+    val expandedType: ConeKotlinType = type.applyIf(type is ConeClassLikeType) { fullyExpandedType(useSiteSession) }
+
     var implicitScope: FirTypeScope? =
         type.scope(useSiteSession, scopeSession, FakeOverrideTypeCalculator.DoNothing, requiredPhase = FirResolvePhase.STATUS)
         private set

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/declarations/ImplicitReceiverUtils.kt

@@ -14,10 +14,13 @@ import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.scopes.impl.wrapNestedClassifierScopeWithSubstitutionForSuperType
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeStubType
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -219,6 +222,7 @@ class FirTowerDataContext private constructor(
     }
 }
 
+// Each FirTowerDataElement has exactly one non-null value among values of properties: scope, implicitReceiver and contextReceiverGroup.
 class FirTowerDataElement(
     val scope: FirScope?,
     val implicitReceiver: ImplicitReceiverValue<*>?,
@@ -236,11 +240,24 @@ class FirTowerDataElement(
         )
 
     /**
-     * Returns [scope] if it is not null. Otherwise, returns [implicitReceiver.implicitScope].
+     * Returns [scope] if it is not null. Otherwise, returns scopes of implicit receivers (including context receivers).
      *
      * Note that a scope for a companion object is an implicit scope.
      */
-    fun getAvailableScope() = scope ?: implicitReceiver?.implicitScope
+    fun getAvailableScopes(): List<FirScope> = when {
+        scope != null -> listOf(scope)
+        implicitReceiver != null -> listOf(implicitReceiver.getImplicitScope())
+        contextReceiverGroup != null -> contextReceiverGroup.map { it.getImplicitScope() }
+        else -> error("Tower data element is expected to have either scope or implicit receivers.")
+    }
+
+    private fun ImplicitReceiverValue<*>.getImplicitScope(): FirScope {
+        return when (expandedType) {
+            is ConeErrorType,
+            is ConeStubType -> FirTypeScope.Empty
+            else -> implicitScope ?: error("Scope for type ${type::class.simpleName} is null.")
+        }
+    }
 }
 
 fun ImplicitReceiverValue<*>.asTowerDataElement(): FirTowerDataElement =