From e4ea7334823ebedc6a1e9dfe6f780f795b88045f Mon Sep 17 00:00:00 2001
From: Roman Efremov <roman.efremov@jetbrains.com>
Date: Tue, 10 Oct 2023 15:04:11 +0200
Subject: [PATCH] [K2] Support reporting of EXPECTED_ENUM_ENTRY_WITH_BODY

Test for the case when enum entry has constructor call but doesn't have
body already exist in
`compiler/testData/diagnostics/tests/multiplatform/enum/constructorInHeaderEnum.kt`.

^KT-59978 Fixed
---
 .../fir/analysis/checkers/SourceNavigator.kt  | 26 +++++++++++++++++++
 .../FirExpectConsistencyChecker.kt            |  8 ++++++
 .../enum/enumEntryWithBody.fir.kt             |  6 ++---
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/SourceNavigator.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/SourceNavigator.kt
index 51eb6f62329..9446bf88502 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/SourceNavigator.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/SourceNavigator.kt
@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.diagnostics.getAncestors
 import org.jetbrains.kotlin.diagnostics.nameIdentifier
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
 import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -41,6 +42,20 @@ interface SourceNavigator {
 
     fun FirTypeRef.isRedundantNullable(): Boolean
 
+    /**
+     * Returns whether this [FirEnumEntry] has a body in source, or `null` if the entry does not have a source.
+     *
+     * Returns `false` if entry has a constructor call, but doesn't have a body:
+     * ```kotlin
+     * enum class E(i: Int) { FOO(42) }
+     * ```
+     *
+     * We have to go down to source level, since this cannot be checked only by FIR element. This is because in FIR all enum entries
+     * with constructor calls have a fake [FirEnumEntry.initializer] with an anonymous object, regardless of whether the entry had
+     * body originally.
+     */
+    fun FirEnumEntry.hasBody(): Boolean?
+
     companion object {
 
         private val lightTreeInstance = LightTreeSourceNavigator()
@@ -108,6 +123,12 @@ private open class LightTreeSourceNavigator : SourceNavigator {
         parent?.let { parent = source.treeStructure.getParent(it) }
         return parent
     }
+
+    override fun FirEnumEntry.hasBody(): Boolean? {
+        val source = source ?: return null
+        val childNodes = source.lighterASTNode.getChildren(source.treeStructure)
+        return childNodes.any { it.tokenType == KtNodeTypes.CLASS_BODY }
+    }
 }
 
 //by default psi tree can reuse light tree manipulations
@@ -150,4 +171,9 @@ private object PsiSourceNavigator : LightTreeSourceNavigator() {
         val typeElement = typeReference.typeElement as? KtNullableType ?: return false
         return typeElement.innerType is KtNullableType
     }
+
+    override fun FirEnumEntry.hasBody(): Boolean? {
+        val enumEntryPsi = source?.psi as? KtEnumEntry ?: return null
+        return enumEntryPsi.body != null
+    }
 }
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirExpectConsistencyChecker.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirExpectConsistencyChecker.kt
index e5b0e4ac35c..781d20c8adb 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirExpectConsistencyChecker.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirExpectConsistencyChecker.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.SourceNavigator.Companion.withNavigator
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.valOrVarKeyword
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
@@ -48,6 +49,9 @@ object FirExpectConsistencyChecker : FirBasicDeclarationChecker() {
         if (isProhibitedEnumConstructor(declaration, lastClass)) {
             reporter.reportOn(source, FirErrors.EXPECTED_ENUM_CONSTRUCTOR, context)
         }
+        if (isProhibitedEnumEntryWithBody(declaration)) {
+            reporter.reportOn(source, FirErrors.EXPECTED_ENUM_ENTRY_WITH_BODY, context)
+        }
 
         if (isProhibitedPrivateDeclaration(declaration)) {
             reporter.reportOn(source, FirErrors.EXPECTED_PRIVATE_DECLARATION, context)
@@ -94,4 +98,8 @@ object FirExpectConsistencyChecker : FirBasicDeclarationChecker() {
     private fun isProhibitedDeclarationWithBody(declaration: FirMemberDeclaration): Boolean {
         return declaration is FirFunction && declaration.hasBody
     }
+
+    private fun isProhibitedEnumEntryWithBody(declaration: FirMemberDeclaration): Boolean {
+        return declaration is FirEnumEntry && declaration.withNavigator { declaration.hasBody() == true }
+    }
 }
diff --git a/compiler/testData/diagnostics/tests/multiplatform/enum/enumEntryWithBody.fir.kt b/compiler/testData/diagnostics/tests/multiplatform/enum/enumEntryWithBody.fir.kt
index d8178b31ee5..49277f563dd 100644
--- a/compiler/testData/diagnostics/tests/multiplatform/enum/enumEntryWithBody.fir.kt
+++ b/compiler/testData/diagnostics/tests/multiplatform/enum/enumEntryWithBody.fir.kt
@@ -3,8 +3,8 @@
 
 <!NO_ACTUAL_FOR_EXPECT!>expect enum class En {
     E1,
-    E2 {
+    <!EXPECTED_ENUM_ENTRY_WITH_BODY!>E2 {
         <!EXPECTED_DECLARATION_WITH_BODY!>fun foo()<!> = ""
-    },
-    E3 { };
+    },<!>
+    <!EXPECTED_ENUM_ENTRY_WITH_BODY!>E3 { };<!>
 }<!>