[FIR] Prohibit access to members of companion of enum from enum entry initializer

^KT-49110

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -2028,7 +2028,7 @@ sealed class KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         abstract val enumEntry: KtSymbol
     }
 
-    abstract class UninitializedEnumCompanion : KtFirDiagnostic<KtSimpleNameExpression>() {
+    abstract class UninitializedEnumCompanion : KtFirDiagnostic<KtExpression>() {
         override val diagnosticClass get() = UninitializedEnumCompanion::class
         abstract val enumClass: KtClassLikeSymbol
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -2446,7 +2446,7 @@ internal class UninitializedEnumCompanionImpl(
     override val enumClass: KtClassLikeSymbol,
     override val firDiagnostic: KtPsiDiagnostic,
     override val token: ValidityToken,
-) : KtFirDiagnostic.UninitializedEnumCompanion(), KtAbstractFirDiagnostic<KtSimpleNameExpression>
+) : KtFirDiagnostic.UninitializedEnumCompanion(), KtAbstractFirDiagnostic<KtExpression>
 
 internal class ValReassignmentImpl(
     override val variable: KtVariableLikeSymbol,

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -5751,6 +5751,18 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/tryWithAssignmentUsedInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("uninitializedCompanionOfEnum_after.kt")
+            public void testUninitializedCompanionOfEnum_after() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_after.kt");
+            }
+
+            @Test
+            @TestMetadata("uninitializedCompanionOfEnum_before.kt")
+            public void testUninitializedCompanionOfEnum_before() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_before.kt");
+            }
+
             @Test
             @TestMetadata("UninitializedEnumCompanionVals.kt")
             public void testUninitializedEnumCompanionVals() throws Exception {

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

@@ -5751,6 +5751,18 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/tryWithAssignmentUsedInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("uninitializedCompanionOfEnum_after.kt")
+            public void testUninitializedCompanionOfEnum_after() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_after.kt");
+            }
+
+            @Test
+            @TestMetadata("uninitializedCompanionOfEnum_before.kt")
+            public void testUninitializedCompanionOfEnum_before() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_before.kt");
+            }
+
             @Test
             @TestMetadata("UninitializedEnumCompanionVals.kt")
             public void testUninitializedEnumCompanionVals() throws Exception {

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

@@ -5751,6 +5751,18 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/tryWithAssignmentUsedInCatch.kt");
             }
 
+            @Test
+            @TestMetadata("uninitializedCompanionOfEnum_after.kt")
+            public void testUninitializedCompanionOfEnum_after() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_after.kt");
+            }
+
+            @Test
+            @TestMetadata("uninitializedCompanionOfEnum_before.kt")
+            public void testUninitializedCompanionOfEnum_before() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_before.kt");
+            }
+
             @Test
             @TestMetadata("UninitializedEnumCompanionVals.kt")
             public void testUninitializedEnumCompanionVals() throws Exception {

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/FirDiagnosticsList.kt

@@ -1023,7 +1023,7 @@ object DIAGNOSTICS_LIST : DiagnosticList("FirErrors") {
         val UNINITIALIZED_ENUM_ENTRY by error<KtSimpleNameExpression> {
             parameter<FirEnumEntrySymbol>("enumEntry")
         }
-        val UNINITIALIZED_ENUM_COMPANION by error<KtSimpleNameExpression>(PositioningStrategy.REFERENCE_BY_QUALIFIED) {
+        val UNINITIALIZED_ENUM_COMPANION by error<KtExpression>(PositioningStrategy.REFERENCE_BY_QUALIFIED) {
             parameter<FirRegularClassSymbol>("enumClass")
         }
         val VAL_REASSIGNMENT by error<KtExpression> {

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt

@@ -559,7 +559,7 @@ object FirErrors {
     val UNINITIALIZED_VARIABLE by error1<KtExpression, FirPropertySymbol>()
     val UNINITIALIZED_PARAMETER by error1<KtSimpleNameExpression, FirValueParameterSymbol>()
     val UNINITIALIZED_ENUM_ENTRY by error1<KtSimpleNameExpression, FirEnumEntrySymbol>()
-    val UNINITIALIZED_ENUM_COMPANION by error1<KtSimpleNameExpression, FirRegularClassSymbol>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
+    val UNINITIALIZED_ENUM_COMPANION by error1<KtExpression, FirRegularClassSymbol>(SourceElementPositioningStrategies.REFERENCE_BY_QUALIFIED)
     val VAL_REASSIGNMENT by error1<KtExpression, FirVariableSymbol<*>>()
     val VAL_REASSIGNMENT_VIA_BACKING_FIELD by deprecationError1<KtExpression, FirBackingFieldSymbol>(RestrictionOfValReassignmentViaBackingField)
     val CAPTURED_VAL_INITIALIZATION by error1<KtExpression, FirPropertySymbol>()

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/CommonDeclarationCheckers.kt

@@ -156,4 +156,9 @@ object CommonDeclarationCheckers : DeclarationCheckers() {
         get() = setOf(
             FirAnonymousInitializerInInterfaceChecker
         )
+
+    override val enumEntryCheckers: Set<FirEnumEntryChecker>
+        get() = setOf(
+            FirEnumCompanionInEnumConstructorCallChecker,
+        )
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirEnumCompanionInEnumConstructorCallChecker.kt

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.analysis.checkers.declaration
+
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
+import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.primaryConstructorIfAny
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.fir.visitors.acceptAllElements
+import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull
+
+object FirEnumCompanionInEnumConstructorCallChecker : FirEnumEntryChecker() {
+    override fun check(declaration: FirEnumEntry, context: CheckerContext, reporter: DiagnosticReporter) {
+        val enumClass = context.containingDeclarations.lastIsInstanceOrNull<FirRegularClass>() ?: return
+        if (enumClass.classKind != ClassKind.ENUM_CLASS) return
+        val companionOfEnumSymbol = enumClass.companionObject?.symbol ?: return
+        val initializerObject = (declaration.initializer as? FirAnonymousObjectExpression)?.anonymousObject ?: return
+        val delegatingConstructorCall = initializerObject.primaryConstructorIfAny(context.session)?.resolvedDelegatedConstructorCall ?: return
+        val visitor = Visitor(context, reporter, companionOfEnumSymbol)
+        delegatingConstructorCall.argumentList.acceptChildren(visitor)
+    }
+
+    private class Visitor(
+        val context: CheckerContext,
+        val reporter: DiagnosticReporter,
+        val companionSymbol: FirRegularClassSymbol
+    ) : FirVisitorVoid() {
+        override fun visitElement(element: FirElement) {
+            element.acceptChildren(this)
+        }
+
+        override fun visitFunctionCall(functionCall: FirFunctionCall) {
+            val needVisitReceiver = checkQualifiedAccess(functionCall)
+            functionCall.argumentList.acceptChildren(this)
+            if (needVisitReceiver) {
+                functionCall.explicitReceiver?.accept(this)
+            }
+        }
+
+        override fun visitPropertyAccessExpression(propertyAccessExpression: FirPropertyAccessExpression) {
+            val needVisitReceiver = checkQualifiedAccess(propertyAccessExpression)
+            if (needVisitReceiver) {
+                propertyAccessExpression.explicitReceiver?.accept(this)
+            }
+        }
+
+
+        private fun checkQualifiedAccess(expression: FirQualifiedAccessExpression): Boolean {
+            val extensionReceiver = expression.extensionReceiver
+            if (extensionReceiver !is FirResolvedQualifier && extensionReceiver !is FirThisReceiverExpression) return true
+
+            val receiverSymbol = expression.toResolvedCallableSymbol()
+                ?.resolvedReceiverTypeRef
+                ?.toRegularClassSymbol(context.session)
+                ?: return true
+
+            if (receiverSymbol == companionSymbol) {
+                val source = expression.extensionReceiver.source ?: expression.source
+                reporter.reportOn(source, FirErrors.UNINITIALIZED_ENUM_COMPANION, companionSymbol, context)
+                return false
+            }
+
+            return true
+        }
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/visitors/FirTransformerUtil.kt

@@ -39,3 +39,11 @@ inline fun <T : FirElement, D> MutableList<T>.transformInplace(transformer: FirT
         iterator.set(result)
     }
 }
+
+fun <R, D> List<FirElement>.acceptAllElements(visitor: FirVisitor<R, D>, data: D) {
+    forEach { it.accept(visitor, data) }
+}
+
+fun List<FirElement>.acceptAllElements(visitor: FirVisitorVoid) {
+    forEach { it.accept(visitor) }
+}

compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_after.fir.kt

@@ -0,0 +1,39 @@
+// LANGUAGE: +ProhibitAccessToEnumCompanionMembersInEnumConstructorCall
+// ISSUE: KT-49110
+
+enum class SomeEnum(val x: Int) {
+    A(<!UNINITIALIZED_ENUM_COMPANION!>companionFun<!>().length),// UNINITIALIZED_ENUM_COMPANION
+    B(<!UNINITIALIZED_VARIABLE!>companionProp<!>.length), // UNINITIALIZED_VARIABLE
+
+    C(SomeEnum.<!UNINITIALIZED_ENUM_COMPANION!>companionFun<!>().length),
+    D(<!UNINITIALIZED_VARIABLE!>SomeEnum.companionProp<!>.length),
+
+    E(SomeEnum.<!UNINITIALIZED_ENUM_COMPANION!>Companion<!>.companionFun().length),
+    F(<!UNINITIALIZED_VARIABLE!>SomeEnum.Companion.companionProp<!>.length); // UNINITIALIZED_VARIABLE
+
+    companion object {
+        val companionProp = "someString"
+        fun companionFun(): String = "someString"
+    }
+}
+
+enum class OtherEnum(val x: Int) {
+    G(<!UNINITIALIZED_ENUM_COMPANION!>extensionFun<!>().length), // UNINITIALIZED_ENUM_COMPANION
+    H(<!UNINITIALIZED_ENUM_COMPANION!>extensionProp<!>.length),
+
+    I(<!UNINITIALIZED_ENUM_COMPANION!>OtherEnum<!>.extensionFun().length),
+    J(<!UNINITIALIZED_ENUM_COMPANION!>OtherEnum<!>.extensionProp.length),
+
+    K(OtherEnum.<!UNINITIALIZED_ENUM_COMPANION!>Companion<!>.extensionFun().length),
+    L(OtherEnum.<!UNINITIALIZED_ENUM_COMPANION!>Companion<!>.extensionProp.length);
+
+    companion object {
+        val companionProp = "someString"
+        fun companionFun(): String = "someString"
+    }
+}
+
+fun OtherEnum.Companion.extensionFun(): String = companionFun()
+val OtherEnum.Companion.extensionProp: String
+    get() = companionProp
+

compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_after.kt

@@ -1,5 +1,4 @@
 // LANGUAGE: +ProhibitAccessToEnumCompanionMembersInEnumConstructorCall
-// FIR_IDENTICAL
 // ISSUE: KT-49110
 
 enum class SomeEnum(val x: Int) {

compiler/testData/diagnostics/tests/controlFlowAnalysis/uninitializedCompanionOfEnum_before.fir.kt

@@ -18,14 +18,14 @@ enum class SomeEnum(val x: Int) {
 }
 
 enum class OtherEnum(val x: Int) {
-    G(extensionFun().length), // UNINITIALIZED_ENUM_COMPANION
+    G(<!UNINITIALIZED_ENUM_COMPANION!>extensionFun<!>().length), // UNINITIALIZED_ENUM_COMPANION
-    H(extensionProp.length),
+    H(<!UNINITIALIZED_ENUM_COMPANION!>extensionProp<!>.length),
 
-    I(OtherEnum.extensionFun().length),
+    I(<!UNINITIALIZED_ENUM_COMPANION!>OtherEnum<!>.extensionFun().length),
-    J(OtherEnum.extensionProp.length),
+    J(<!UNINITIALIZED_ENUM_COMPANION!>OtherEnum<!>.extensionProp.length),
 
-    K(OtherEnum.Companion.extensionFun().length),
+    K(OtherEnum.<!UNINITIALIZED_ENUM_COMPANION!>Companion<!>.extensionFun().length),
-    L(OtherEnum.Companion.extensionProp.length);
+    L(OtherEnum.<!UNINITIALIZED_ENUM_COMPANION!>Companion<!>.extensionProp.length);
 
     companion object {
         val companionProp = "someString"