From f85fc47383bb1ba0c0933c11742c147927cd8117 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Thu, 14 Jan 2021 18:23:05 +0300
Subject: [PATCH] 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.
---
 .../selfReferenceToCompanionObject.fir.txt    | 36 +++++++++++++++++++
 .../selfReferenceToCompanionObject.kt         |  9 +++++
 .../runners/FirDiagnosticTestGenerated.java   |  6 ++++
 .../fir/resolve/BodyResolveComponents.kt      |  1 +
 .../body/resolve/BodyResolveContext.kt        |  3 ++
 .../FirDeclarationsResolveTransformer.kt      | 28 ++++++++-------
 ...ectInInlineLambdaInSuperConstructorCall.kt |  1 -
 ...ectInInlineLambdaInSuperConstructorCall.kt |  1 -
 ...nionObjectSuperConstructorArguments.fir.kt | 19 ----------
 ...ompanionObjectSuperConstructorArguments.kt |  1 +
 10 files changed, 72 insertions(+), 33 deletions(-)
 create mode 100644 compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.fir.txt
 create mode 100644 compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.kt
 delete mode 100644 compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.fir.kt

diff --git a/compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.fir.txt b/compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.fir.txt
new file mode 100644
index 00000000000..52c10711859
--- /dev/null
+++ b/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|>|()
+    }
diff --git a/compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.kt b/compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/selfReferenceToCompanionObject.kt
new file mode 100644
index 00000000000..0d2dcfa78a4
--- /dev/null
+++ b/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()
diff --git a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java
index 65e5c6b13af..adb9a067766 100644
--- a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java
+++ b/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 {
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt
index 886de9e44da..38c4a2d54a9 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt
+++ b/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?,
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
index 42b6f473833..e9254f795bb 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
+++ b/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
 
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
index 16406b89295..42e5189ef16 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
+++ b/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 =
-            context.towerDataContext
-                .addNonLocalTowerDataElements(towerElementsForClass.superClassesStaticsAndCompanionReceivers)
-                .run {
-                    if (towerElementsForClass.companionReceiver != null)
-                        addReceiver(null, towerElementsForClass.companionReceiver)
-                    else
-                        this
-                }
-                .addNonLocalScopeIfNotNull(towerElementsForClass.companionStaticScope)
-                .addNonLocalScopeIfNotNull(towerElementsForClass.staticScope)
+        val base = context.towerDataContext.addNonLocalTowerDataElements(towerElementsForClass.superClassesStaticsAndCompanionReceivers)
+        val statics = base
+            .addNonLocalScopeIfNotNull(towerElementsForClass.companionStaticScope)
+            .addNonLocalScopeIfNotNull(towerElementsForClass.staticScope)
+
+        val companionReceiver = towerElementsForClass.companionReceiver
+        val staticsAndCompanion = if (companionReceiver == null) statics else base
+            .addReceiver(null, companionReceiver)
+            .addNonLocalScopeIfNotNull(towerElementsForClass.companionStaticScope)
+            .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
diff --git a/compiler/testData/codegen/box/objects/selfReferenceToCompanionObjectInInlineLambdaInSuperConstructorCall.kt b/compiler/testData/codegen/box/objects/selfReferenceToCompanionObjectInInlineLambdaInSuperConstructorCall.kt
index e877be87f1d..fc2a08dd42e 100644
--- a/compiler/testData/codegen/box/objects/selfReferenceToCompanionObjectInInlineLambdaInSuperConstructorCall.kt
+++ b/compiler/testData/codegen/box/objects/selfReferenceToCompanionObjectInInlineLambdaInSuperConstructorCall.kt
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 abstract class Base(val fn: () -> String)
 
 class Host {
diff --git a/compiler/testData/codegen/box/objects/selfReferenceToInterfaceCompanionObjectInInlineLambdaInSuperConstructorCall.kt b/compiler/testData/codegen/box/objects/selfReferenceToInterfaceCompanionObjectInInlineLambdaInSuperConstructorCall.kt
index c0f50527d87..226506094b2 100644
--- a/compiler/testData/codegen/box/objects/selfReferenceToInterfaceCompanionObjectInInlineLambdaInSuperConstructorCall.kt
+++ b/compiler/testData/codegen/box/objects/selfReferenceToInterfaceCompanionObjectInInlineLambdaInSuperConstructorCall.kt
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 abstract class Base(val fn: () -> String)
 
 interface Host {
diff --git a/compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.fir.kt b/compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.fir.kt
deleted file mode 100644
index 6de00281a0b..00000000000
--- a/compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.fir.kt
+++ /dev/null
@@ -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() {}
-    }
-}
diff --git a/compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.kt b/compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.kt
index 2eb3300447f..f03b9903174 100644
--- a/compiler/testData/diagnostics/tests/scopes/classHeader/companionObjectSuperConstructorArguments.kt
+++ b/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 {