FIR: introduce separate companion object resolve context

Before this commit, during the resolve of companion objects we used
the same context than for any nested class. However, during companion
object resolve we should not have companion object receiver itself
accessible in any case (in particular, it should not be accessible
in constructor). So in this commit we introduced separate context
for this purpose.

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.fir.txt

@@ -0,0 +1,36 @@
+FILE: selfReferenceToCompanionObject.kt
+    public abstract class Base : R|kotlin/Any| {
+        public constructor(fn: R|() -> kotlin/String|): R|Base| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final val fn: R|() -> kotlin/String| = R|<local>/fn|
+            public get(): R|() -> kotlin/String|
+
+    }
+    public final class Host : R|kotlin/Any| {
+        public constructor(): R|Host| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final companion object Companion : R|Base| {
+            private constructor(): R|Host.Companion| {
+                super<R|Base|>(R|kotlin/run|<R|() -> kotlin/String|>(<L> = run@fun <anonymous>(): R|() -> kotlin/String| <kind=EXACTLY_ONCE>  {
+                    ^ run@fun <anonymous>(): R|kotlin/String| {
+                        ^ Q|Host|.R|/Host.Companion.ok|()
+                    }
+
+                }
+                ))
+            }
+
+            public final fun ok(): R|kotlin/String| {
+                ^ok String(OK)
+            }
+
+        }
+
+    }
+    public final fun box(): R|kotlin/String| {
+        ^box Q|Host.Companion|.R|/Base.fn|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/String|>|()
+    }

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.kt

@@ -0,0 +1,9 @@
+abstract class Base(val fn: () -> String)
+
+class Host {
+    companion object : Base(run { { Host.ok() } }) {
+        fun ok() = "OK"
+    }
+}
+
+fun box() = Host.Companion.fn()

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java

@@ -4889,6 +4889,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/receiverResolutionInLambda.kt");
             }
 
+            @Test
+            @TestMetadata("selfReferenceToCompanionObject.kt")
+            public void testSelfReferenceToCompanionObject() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.kt");
+            }
+
             @Test
             @TestMetadata("SpecialCallsWithLambdas.kt")
             public void testSpecialCallsWithLambdas() throws Exception {

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

@@ -143,6 +143,7 @@ typealias TowerDataContextForAnonymousFunctions = Map<FirAnonymousFunctionSymbol
 
 class FirTowerDataContextsForClassParts(
     val forNestedClasses: FirTowerDataContext,
+    val forCompanionObject: FirTowerDataContext,
     val forConstructorHeaders: FirTowerDataContext,
     val primaryConstructorPureParametersScope: FirLocalScope?,
     val primaryConstructorAllParametersScope: FirLocalScope?,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -82,6 +82,9 @@ class BodyResolveContext(
     fun getTowerDataContextForStaticNestedClassesUnsafe(): FirTowerDataContext =
         firTowerDataContextsForClassParts().forNestedClasses
 
+    fun getTowerDataContextForCompanionUnsafe(): FirTowerDataContext =
+        firTowerDataContextsForClassParts().forCompanionObject
+
     fun getTowerDataContextForConstructorResolution(): FirTowerDataContext =
         firTowerDataContextsForClassParts().forConstructorHeaders
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -348,7 +348,11 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         return context.withTowerDataCleanup {
             if (!regularClass.isInner && context.containerIfAny is FirRegularClass) {
                 context.replaceTowerDataContext(
-                    context.getTowerDataContextForStaticNestedClassesUnsafe()
+                    if (regularClass.isCompanion) {
+                        context.getTowerDataContextForCompanionUnsafe()
+                    } else {
+                        context.getTowerDataContextForStaticNestedClassesUnsafe()
+                    }
                 )
             }
 
@@ -816,17 +820,16 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     ): T = context.withTowerDataCleanup {
         val towerElementsForClass = components.collectTowerDataElementsForClass(owner, type)
 
-        val staticsAndCompanion =
+        val base = context.towerDataContext.addNonLocalTowerDataElements(towerElementsForClass.superClassesStaticsAndCompanionReceivers)
-            context.towerDataContext
+        val statics = base
-                .addNonLocalTowerDataElements(towerElementsForClass.superClassesStaticsAndCompanionReceivers)
+            .addNonLocalScopeIfNotNull(towerElementsForClass.companionStaticScope)
-                .run {
+            .addNonLocalScopeIfNotNull(towerElementsForClass.staticScope)
-                    if (towerElementsForClass.companionReceiver != null)
+
-                        addReceiver(null, towerElementsForClass.companionReceiver)
+        val companionReceiver = towerElementsForClass.companionReceiver
-                    else
+        val staticsAndCompanion = if (companionReceiver == null) statics else base
-                        this
+            .addReceiver(null, companionReceiver)
-                }
+            .addNonLocalScopeIfNotNull(towerElementsForClass.companionStaticScope)
-                .addNonLocalScopeIfNotNull(towerElementsForClass.companionStaticScope)
+            .addNonLocalScopeIfNotNull(towerElementsForClass.staticScope)
-                .addNonLocalScopeIfNotNull(towerElementsForClass.staticScope)
 
         val typeParameterScope = (owner as? FirRegularClass)?.let(this::createTypeParameterScope)
 
@@ -856,6 +859,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
         val newContexts = FirTowerDataContextsForClassParts(
             newTowerDataContextForStaticNestedClasses,
+            statics,
             scopeForConstructorHeader,
             primaryConstructorPureParametersScope,
             primaryConstructorAllParametersScope

compiler/testData/codegen/box/objects/selfReferenceToCompanionObjectInInlineLambdaInSuperConstructorCall.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 abstract class Base(val fn: () -> String)
 
 class Host {

compiler/testData/codegen/box/objects/selfReferenceToInterfaceCompanionObjectInInlineLambdaInSuperConstructorCall.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 abstract class Base(val fn: () -> String)
 
 interface Host {

compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.fir.kt

@@ -1,19 +0,0 @@
-open class S(val a: Any, val b: Any, val c: Any) {}
-
-interface A {
-    companion object : S(prop1, prop2, func()) {
-        val prop1 = 1
-        val prop2: Int
-            get() = 1
-        fun func() {}
-    }
-}
-
-class B {
-    companion object : S(prop1, prop2, func()) {
-        val prop1 = 1
-        val prop2: Int
-            get() = 1
-        fun func() {}
-    }
-}

compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 open class S(val a: Any, val b: Any, val c: Any) {}
 
 interface A {