FIR resolve: add handling of companions as objects

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/declarations/FirJavaClass.kt

@@ -43,6 +43,9 @@ class FirJavaClass(
 
     override val declarations = mutableListOf<FirDeclaration>()
 
+    override val companionObject: FirRegularClass?
+        get() = null
+
     override fun replaceSupertypes(newSupertypes: List<FirTypeRef>): FirRegularClass {
         superTypeRefs.clear()
         superTypeRefs.addAll(newSupertypes)

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -469,8 +469,10 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
                 val delegatedSelfType = enumEntry.toDelegatedSelfType(firEnumEntry)
                 val delegatedSuperType = enumEntry.extractSuperTypeListEntriesTo(firEnumEntry, delegatedSelfType)
                 for (declaration in enumEntry.declarations) {
-                    firEnumEntry.declarations += declaration.toFirDeclaration(
-                        delegatedSuperType, delegatedSelfType, enumEntry, hasPrimaryConstructor = true
+                    firEnumEntry.addDeclaration(
+                        declaration.toFirDeclaration(
+                            delegatedSuperType, delegatedSelfType, enumEntry, hasPrimaryConstructor = true
+                        )
                     )
                 }
                 firEnumEntry
@@ -536,14 +538,16 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
                 if (primaryConstructor != null && firPrimaryConstructor != null) {
                     primaryConstructor.valueParameters.zip(firPrimaryConstructor.valueParameters).forEach { (ktParameter, firParameter) ->
                         if (ktParameter.hasValOrVar()) {
-                            firClass.declarations += ktParameter.toFirProperty(firParameter)
+                            firClass.addDeclaration(ktParameter.toFirProperty(firParameter))
                         }
                     }
                 }
 
                 for (declaration in classOrObject.declarations) {
-                    firClass.declarations += declaration.toFirDeclaration(
-                        delegatedSuperType, delegatedSelfType, classOrObject, hasPrimaryConstructor = primaryConstructor != null
+                    firClass.addDeclaration(
+                        declaration.toFirDeclaration(
+                            delegatedSuperType, delegatedSelfType, classOrObject, hasPrimaryConstructor = primaryConstructor != null
+                        )
                     )
                 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt

@@ -64,7 +64,7 @@ fun deserializeClassToSymbol(
         }
 
         // TODO: properties
-        declarations += classProto.functionList.map(classDeserializer::loadFunction)
+        addDeclarations(classProto.functionList.map(classDeserializer::loadFunction))
 
         val delegatedSelfType = FirResolvedTypeRefImpl(
             session,
@@ -77,12 +77,18 @@ fun deserializeClassToSymbol(
             isMarkedNullable = false,
             annotations = emptyList()
         )
-        declarations += classProto.constructorList.map { classDeserializer.loadConstructor(it, delegatedSelfType) }
+        addDeclarations(
+            classProto.constructorList.map {
+                classDeserializer.loadConstructor(it, delegatedSelfType)
+            }
+        )
 
-        declarations += classProto.nestedClassNameList.mapNotNull { nestedNameId ->
-            val nestedClassId = classId.createNestedClassId(Name.identifier(nameResolver.getString(nestedNameId)))
-            deserializeNestedClass(nestedClassId, context)?.fir
-        }
+        addDeclarations(
+            classProto.nestedClassNameList.mapNotNull { nestedNameId ->
+                val nestedClassId = classId.createNestedClassId(Name.identifier(nameResolver.getString(nestedNameId)))
+                deserializeNestedClass(nestedClassId, context)?.fir
+            }
+        )
     }
 }
 

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

@@ -29,7 +29,10 @@ fun ConeKotlinType.scope(useSiteSession: FirSession, scopeSession: ScopeSession)
             // For ConeClassLikeType they might be a type alias instead of a regular class
             // TODO: support that case and switch back to `firUnsafe` instead of `firSafeNullable`
             val fir = this.lookupTag.toSymbol(useSiteSession)?.firSafeNullable<FirRegularClass>() ?: return null
-            wrapSubstitutionScopeIfNeed(useSiteSession, fir.buildUseSiteScope(useSiteSession, scopeSession), scopeSession)
+            val companionScope = fir.companionObject?.buildUseSiteScope(useSiteSession, scopeSession)
+            val ownScope = wrapSubstitutionScopeIfNeed(useSiteSession, fir.buildUseSiteScope(useSiteSession, scopeSession), scopeSession)
+            if (companionScope != null) FirCompositeScope(mutableListOf(ownScope, companionScope)) else ownScope
+
         }
         is ConeTypeParameterType -> {
             // TODO: support LibraryTypeParameterSymbol or get rid of it

compiler/fir/resolve/testData/resolve/expresssions/companion.kt

@@ -0,0 +1,21 @@
+open class A {
+    companion object {
+        fun foo() {}
+    }
+
+    fun bar() {
+        foo()
+    }
+}
+
+class B {
+    companion object : A() {
+        fun baz() {}
+    }
+}
+
+fun test() {
+    A.foo()
+    B.bar()
+    B.baz()
+}

compiler/fir/resolve/testData/resolve/expresssions/companion.txt

@@ -0,0 +1,42 @@
+FILE: companion.kt
+    public open class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final companion object Companion : R|kotlin/Any| {
+            private constructor(): R|A.Companion| {
+                super<R|kotlin/Any|>()
+            }
+
+            public final fun foo(): R|kotlin/Unit| {
+            }
+
+        }
+
+        public final fun bar(): R|kotlin/Unit| {
+            R|/A.Companion.foo|()
+        }
+
+    }
+    public final class B : R|kotlin/Any| {
+        public constructor(): R|B| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final companion object Companion : R|A| {
+            private constructor(): R|B.Companion| {
+                super<R|A|>()
+            }
+
+            public final fun baz(): R|kotlin/Unit| {
+            }
+
+        }
+
+    }
+    public final fun test(): R|kotlin/Unit| {
+        R|/A|.R|/A.Companion.foo|()
+        R|/B|.<Inapplicable(WRONG_RECEIVER): [/A.bar]>#()
+        R|/B|.R|/B.Companion.baz|()
+    }

compiler/fir/resolve/testData/resolve/fromBuilder/enums.txt

@@ -70,7 +70,7 @@ FILE: enums.kt
 
         }
 
-        public final val g: R|kotlin/Double| = <Unresolved name: G>#.<Unresolved name: times>#(R|/Planet.m|).<Unresolved name: div>#(R|/Planet.r|.R|kotlin/Double.times|(R|/Planet.r|))
+        public final val g: R|kotlin/Double| = R|/Planet.Companion.G|.R|kotlin/Double.times|(R|/Planet.m|).R|kotlin/Double.div|(R|/Planet.r|.R|kotlin/Double.times|(R|/Planet.r|))
             public get(): R|kotlin/Double|
 
         public abstract fun sayHello(): R|kotlin/Unit|

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirResolveTestCaseGenerated.java

@@ -202,6 +202,11 @@ public class FirResolveTestCaseGenerated extends AbstractFirResolveTestCase {
             runTest("compiler/fir/resolve/testData/resolve/expresssions/checkArguments.kt");
         }
 
+        @TestMetadata("companion.kt")
+        public void testCompanion() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/expresssions/companion.kt");
+        }
+
         @TestMetadata("constructor.kt")
         public void testConstructor() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/expresssions/constructor.kt");

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirRegularClass.kt

@@ -24,6 +24,8 @@ interface FirRegularClass : FirClass, @VisitedSupertype FirClassLikeDeclaration,
 
     val isInline: Boolean get() = status.isInline
 
+    val companionObject: FirRegularClass?
+
     override val symbol: FirClassSymbol
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/impl/FirClassImpl.kt

@@ -46,6 +46,19 @@ open class FirClassImpl(
 
     override val declarations = mutableListOf<FirDeclaration>()
 
+    override var companionObject: FirRegularClass? = null
+
+    fun addDeclaration(declaration: FirDeclaration) {
+        declarations += declaration
+        if (companionObject == null && declaration is FirRegularClass && declaration.isCompanion) {
+            companionObject = declaration
+        }
+    }
+
+    fun addDeclarations(declarations: Collection<FirDeclaration>) {
+        declarations.forEach(this::addDeclaration)
+    }
+
     override fun replaceSupertypes(newSupertypes: List<FirTypeRef>): FirRegularClass {
         superTypeRefs.clear()
         superTypeRefs.addAll(newSupertypes)
@@ -58,6 +71,7 @@ open class FirClassImpl(
 
         // Transform declarations in last turn
         declarations.transformInplace(transformer, data)
+        companionObject = declarations.asSequence().filterIsInstance<FirRegularClass>().firstOrNull { it.isCompanion }
         return result
     }
 }