FIR: add a boolean tracking if a class ref resolves to the companion object

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -93,6 +93,10 @@ class FirCallResolver(
 
         val resultExpression = functionCall.transformCalleeReference(StoreNameReference, nameReference)
         val candidate = (nameReference as? FirNamedReferenceWithCandidate)?.candidate
+        val resolvedReceiver = functionCall.explicitReceiver
+        if (candidate != null && resolvedReceiver is FirResolvedQualifier) {
+            resolvedReceiver.replaceResolvedToCompanionObject(candidate.isFromCompanionObjectTypeScope)
+        }
 
         // We need desugaring
         val resultFunctionCall = if (candidate != null && candidate.callInfo != result.info) {
@@ -223,6 +227,10 @@ class FirCallResolver(
             else -> null
         }
 
+        (qualifiedAccess.explicitReceiver as? FirResolvedQualifier)?.replaceResolvedToCompanionObject(
+            reducedCandidates.isNotEmpty() && reducedCandidates.all { it.isFromCompanionObjectTypeScope }
+        )
+
         when {
             referencedSymbol is FirClassLikeSymbol<*> -> {
                 return components.buildResolvedQualifierForClass(
@@ -284,6 +292,10 @@ class FirCallResolver(
             conflictResolver.chooseMaximallySpecificCandidates(bestCandidates, discriminateGenerics = false)
         }
 
+        (callableReferenceAccess.explicitReceiver as? FirResolvedQualifier)?.replaceResolvedToCompanionObject(
+            bestCandidates.isNotEmpty() && bestCandidates.all { it.isFromCompanionObjectTypeScope }
+        )
+
         resolvedCallableReferenceAtom.hasBeenResolvedOnce = true
 
         when {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Candidate.kt

@@ -85,6 +85,7 @@ class Candidate(
     private val baseSystem: ConstraintStorage,
     val callInfo: CallInfo,
     val originScope: FirScope?,
+    val isFromCompanionObjectTypeScope: Boolean = false
 ) {
 
     var systemInitialized: Boolean = false

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CandidateFactory.kt

@@ -6,9 +6,11 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.builder.buildErrorFunction
 import org.jetbrains.kotlin.fir.declarations.builder.buildErrorProperty
+import org.jetbrains.kotlin.fir.declarations.classId
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.moduleData
@@ -17,6 +19,7 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirErrorFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirErrorPropertySymbol
+import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzerContext
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
@@ -55,9 +58,22 @@ class CandidateFactory private constructor(
                 callInfo.withReceiverAsArgument(it)
             } ?: callInfo,
             scope,
+            isFromCompanionObjectTypeScope = when (explicitReceiverKind) {
+                ExplicitReceiverKind.EXTENSION_RECEIVER -> extensionReceiverValue.isCandidateFromCompanionObjectTypeScope()
+                ExplicitReceiverKind.DISPATCH_RECEIVER -> dispatchReceiverValue.isCandidateFromCompanionObjectTypeScope()
+                // The following cases are not applicable for companion objects.
+                ExplicitReceiverKind.NO_EXPLICIT_RECEIVER, ExplicitReceiverKind.BOTH_RECEIVERS -> false
+            }
         )
     }
 
+    private fun ReceiverValue?.isCandidateFromCompanionObjectTypeScope(): Boolean {
+        val expressionReceiverValue = this as? ExpressionReceiverValue ?: return false
+        val resolvedQualifier = (expressionReceiverValue.explicitReceiver as? FirResolvedQualifier) ?: return false
+        val originClassOfCandidate = expressionReceiverValue.type.classId ?: return false
+        return (resolvedQualifier.symbol?.fir as? FirRegularClass)?.companionObject?.classId == originClassOfCandidate
+    }
+
     fun createErrorCandidate(callInfo: CallInfo, diagnostic: ConeDiagnostic): Candidate {
         val symbol: AbstractFirBasedSymbol<*> = when (callInfo.callKind) {
             is CallKind.VariableAccess -> createErrorPropertySymbol(diagnostic)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirErrorResolvedQualifier.kt

@@ -30,6 +30,7 @@ abstract class FirErrorResolvedQualifier : FirResolvedQualifier(), FirDiagnostic
     abstract override val classId: ClassId?
     abstract override val symbol: FirClassLikeSymbol<*>?
     abstract override val isNullableLHSForCallableReference: Boolean
+    abstract override val resolvedToCompanionObject: Boolean
     abstract override val typeArguments: List<FirTypeProjection>
     abstract override val diagnostic: ConeDiagnostic
 
@@ -43,6 +44,8 @@ abstract class FirErrorResolvedQualifier : FirResolvedQualifier(), FirDiagnostic
 
     abstract override fun replaceIsNullableLHSForCallableReference(newIsNullableLHSForCallableReference: Boolean)
 
+    abstract override fun replaceResolvedToCompanionObject(newResolvedToCompanionObject: Boolean)
+
     abstract override fun replaceTypeArguments(newTypeArguments: List<FirTypeProjection>)
 
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirErrorResolvedQualifier

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirResolvedQualifier.kt

@@ -28,6 +28,7 @@ abstract class FirResolvedQualifier : FirExpression() {
     abstract val classId: ClassId?
     abstract val symbol: FirClassLikeSymbol<*>?
     abstract val isNullableLHSForCallableReference: Boolean
+    abstract val resolvedToCompanionObject: Boolean
     abstract val typeArguments: List<FirTypeProjection>
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitResolvedQualifier(this, data)
@@ -40,6 +41,8 @@ abstract class FirResolvedQualifier : FirExpression() {
 
     abstract fun replaceIsNullableLHSForCallableReference(newIsNullableLHSForCallableReference: Boolean)
 
+    abstract fun replaceResolvedToCompanionObject(newResolvedToCompanionObject: Boolean)
+
     abstract fun replaceTypeArguments(newTypeArguments: List<FirTypeProjection>)
 
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirResolvedQualifier

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirAbstractResolvedQualifierBuilder.kt

@@ -31,6 +31,7 @@ interface FirAbstractResolvedQualifierBuilder {
     abstract var classId: ClassId?
     abstract var symbol: FirClassLikeSymbol<*>?
     abstract var isNullableLHSForCallableReference: Boolean
+    abstract var resolvedToCompanionObject: Boolean
     abstract val typeArguments: MutableList<FirTypeProjection>
 
     fun build(): FirResolvedQualifier

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirErrorResolvedQualifierBuilder.kt

@@ -61,6 +61,13 @@ class FirErrorResolvedQualifierBuilder : FirAbstractResolvedQualifierBuilder, Fi
         set(_) {
             throw IllegalStateException()
         }
+
+    @Deprecated("Modification of 'resolvedToCompanionObject' has no impact for FirErrorResolvedQualifierBuilder", level = DeprecationLevel.HIDDEN)
+    override var resolvedToCompanionObject: Boolean
+        get() = throw IllegalStateException()
+        set(_) {
+            throw IllegalStateException()
+        }
 }
 
 @OptIn(ExperimentalContracts::class)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirResolvedQualifierBuilder.kt

@@ -9,6 +9,7 @@ import kotlin.contracts.*
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.builder.FirAbstractResolvedQualifierBuilder
@@ -58,6 +59,13 @@ class FirResolvedQualifierBuilder : FirAbstractResolvedQualifierBuilder, FirAnno
         set(_) {
             throw IllegalStateException()
         }
+
+    @Deprecated("Modification of 'resolvedToCompanionObject' has no impact for FirResolvedQualifierBuilder", level = DeprecationLevel.HIDDEN)
+    override var resolvedToCompanionObject: Boolean
+        get() = throw IllegalStateException()
+        set(_) {
+            throw IllegalStateException()
+        }
 }
 
 @OptIn(ExperimentalContracts::class)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirErrorResolvedQualifierImpl.kt

@@ -34,6 +34,7 @@ internal class FirErrorResolvedQualifierImpl(
     override val diagnostic: ConeDiagnostic,
 ) : FirErrorResolvedQualifier() {
     override var typeRef: FirTypeRef = FirImplicitTypeRefImpl(null)
+    override val resolvedToCompanionObject: Boolean get() = false
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         typeRef.accept(visitor, data)
@@ -66,6 +67,8 @@ internal class FirErrorResolvedQualifierImpl(
         isNullableLHSForCallableReference = newIsNullableLHSForCallableReference
     }
 
+    override fun replaceResolvedToCompanionObject(newResolvedToCompanionObject: Boolean) {}
+
     override fun replaceTypeArguments(newTypeArguments: List<FirTypeProjection>) {
         typeArguments.clear()
         typeArguments.addAll(newTypeArguments)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirResolvedQualifierImpl.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.expressions.impl
 
 import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
@@ -33,6 +34,7 @@ internal class FirResolvedQualifierImpl(
     override val classId: ClassId? get() = relativeClassFqName?.let {
     ClassId(packageFqName, it, false)
 }
+    override var resolvedToCompanionObject: Boolean = (symbol?.fir as? FirRegularClass)?.companionObject != null
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         typeRef.accept(visitor, data)
@@ -65,6 +67,10 @@ internal class FirResolvedQualifierImpl(
         isNullableLHSForCallableReference = newIsNullableLHSForCallableReference
     }
 
+    override fun replaceResolvedToCompanionObject(newResolvedToCompanionObject: Boolean) {
+        resolvedToCompanionObject = newResolvedToCompanionObject
+    }
+
     override fun replaceTypeArguments(newTypeArguments: List<FirTypeProjection>) {
         typeArguments.clear()
         typeArguments.addAll(newTypeArguments)

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/ImplementationConfigurator.kt

@@ -469,6 +469,27 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
             useTypes(expressionType)
         }
 
+        impl(resolvedQualifier) {
+            // Initialize the value to true if only the companion object is present. This makes a standalone class reference expression
+            // correctly resolve to the companion object. For example
+            // ```
+            // class A {
+            //   companion object
+            // }
+            //
+            // val companionOfA = A // This standalone class reference `A` here should resolve to the companion object.
+            // ```
+            //
+            // If this `FirResolvedQualifier` is a receiver expression of some other qualified access, the value is updated in
+            // `FirCallResolver` according to the resolution result.
+            default("resolvedToCompanionObject", "(symbol?.fir as? FirRegularClass)?.companionObject != null")
+            useTypes(regularClass)
+        }
+
+        impl(errorResolvedQualifier) {
+            defaultFalse("resolvedToCompanionObject", withGetter = true)
+        }
+
         noImpl(userTypeRef)
     }
 

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -516,6 +516,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +field("classId", classIdType, nullable = true)
             +field("symbol", classLikeSymbolType, nullable = true)
             +booleanField("isNullableLHSForCallableReference", withReplace = true)
+            +booleanField("resolvedToCompanionObject", withReplace = true)
             +typeArguments.withTransform()
         }