Support retrieving class of enum entry in ultra-light classes

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightField.kt

@@ -31,7 +31,7 @@ import org.jetbrains.kotlin.resolve.jvm.annotations.VOLATILE_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOriginKind
 
 internal open class KtUltraLightField(
-    private val declaration: KtNamedDeclaration,
+    protected val declaration: KtNamedDeclaration,
     name: String,
     private val containingClass: KtUltraLightClass,
     private val support: UltraLightSupport,
@@ -126,15 +126,24 @@ internal open class KtUltraLightField(
 }
 
 internal class KtUltraLightEnumEntry(
-    declaration: KtNamedDeclaration,
+    declaration: KtEnumEntry,
     name: String,
     containingClass: KtUltraLightClass,
     support: UltraLightSupport,
     modifiers: Set<String>
 ) : KtUltraLightField(declaration, name, containingClass, support, modifiers), PsiEnumConstant {
-    override fun getInitializingClass(): PsiEnumConstantInitializer? = null
+    private val enumEntry get() = declaration as KtEnumEntry
+
+    private val _initializingClass by lazyPub {
+        if (enumEntry.body != null)
+            KtUltraLightClassForEnumEntry(enumEntry, containingClass.support, this)
+        else
+            null
+    }
+
+    override fun getInitializingClass(): PsiEnumConstantInitializer? = _initializingClass
     override fun getOrCreateInitializingClass(): PsiEnumConstantInitializer =
-        error("cannot create initializing class in light enum constant")
+        _initializingClass ?: error("cannot create initializing class in light enum constant")
 
     override fun getArgumentList(): PsiExpressionList? = null
     override fun resolveMethod(): PsiMethod? = null
@@ -145,3 +154,31 @@ internal class KtUltraLightEnumEntry(
     override fun hasInitializer() = true
     override fun computeConstantValue(visitedVars: MutableSet<PsiVariable>?) = this
 }
+
+internal class KtUltraLightClassForEnumEntry(
+    enumEntry: KtEnumEntry, support: UltraLightSupport,
+    private val enumConstant: PsiEnumConstant
+) : KtUltraLightClass(enumEntry, support), PsiEnumConstantInitializer {
+
+    private val baseClassReferenceAndType: Pair<PsiJavaCodeReferenceElement, PsiClassType> by lazyPub {
+        // It should not be null for not-too-complex classes and that is not the case because
+        // the containing class is not too complex (since we created KtUltraLightClassForEnumEntry instance)
+        val extendsList =
+            super.getExtendsList() ?: error("KtUltraLightClass::getExtendsList is null for ${enumEntry.fqName}")
+
+        Pair(
+            extendsList.referenceElements.getOrNull(0) ?: error("No referenceElements found for ${enumEntry.fqName}"),
+            extendsList.referencedTypes.getOrNull(0) ?: error("No referencedTypes found for ${enumEntry.fqName}")
+        )
+    }
+
+    override fun getBaseClassType() = baseClassReferenceAndType.second
+
+    override fun getBaseClassReference() = baseClassReferenceAndType.first
+
+    override fun getArgumentList(): PsiExpressionList? = null
+
+    override fun getEnumConstant() = enumConstant
+
+    override fun isInQualifiedNew() = false
+}

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightPsi.kt

@@ -44,7 +44,7 @@ import org.jetbrains.kotlin.resolve.jvm.annotations.SYNCHRONIZED_ANNOTATION_FQ_N
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.isAnyOrNullableAny
 
-class KtUltraLightClass(classOrObject: KtClassOrObject, private val support: UltraLightSupport) :
+open class KtUltraLightClass(classOrObject: KtClassOrObject, internal val support: UltraLightSupport) :
     KtLightClassImpl(classOrObject) {
     companion object {
 
@@ -316,8 +316,11 @@ class KtUltraLightClass(classOrObject: KtClassOrObject, private val support: Ult
 
     private fun defaultConstructor(): KtUltraLightMethod {
         val visibility =
-            if (classOrObject is KtObjectDeclaration || classOrObject.hasModifier(SEALED_KEYWORD) || isEnum) PsiModifier.PRIVATE
+            when {
-            else PsiModifier.PUBLIC
+                classOrObject is KtObjectDeclaration || classOrObject.hasModifier(SEALED_KEYWORD) || isEnum -> PsiModifier.PRIVATE
+                classOrObject is KtEnumEntry -> PsiModifier.PACKAGE_LOCAL
+                else -> PsiModifier.PUBLIC
+            }
         return noArgConstructor(visibility, classOrObject)
     }
 

idea/idea-core/src/org/jetbrains/kotlin/idea/perf/UltraLightChecker.kt

@@ -110,6 +110,16 @@ object UltraLightChecker {
                 } +
                 ";"
 
+    private fun PsiEnumConstant.renderEnumConstant(): String {
+        val initializingClass = initializingClass ?: return name
+
+        return buildString {
+            appendln("$name {")
+            append(initializingClass.renderMembers())
+            append("}")
+        }
+    }
+
     fun PsiClass.renderClass(): String {
         val classWord = when {
             isAnnotationType -> "@interface"
@@ -118,18 +128,51 @@ object UltraLightChecker {
             else -> "class"
         }
 
-        return renderModifiers() +
+        return buildString {
-                classWord + " " +
+            append(renderModifiers())
-                name + " /* " + qualifiedName + "*/" +
+            append("$classWord ")
-                renderTypeParams() +
+            append("$name /* $qualifiedName*/")
-                extendsList.renderRefList("extends") +
+            append(renderTypeParams())
-                implementsList.renderRefList("implements") +
+            append(extendsList.renderRefList("extends"))
-                " {\n" +
+            append(implementsList.renderRefList("implements"))
-                (if (isEnum) fields.filterIsInstance<PsiEnumConstant>().joinToString(",\n") { it.name } + ";\n\n" else "") +
+            appendln(" {")
-                fields.filterNot { it is PsiEnumConstant }.map { it.renderVar().prependIndent("  ") + ";\n\n" }.sorted().joinToString("") +
+
-                methods.map { it.renderMethod().prependIndent("  ") + "\n\n" }.sorted().joinToString("") +
+            if (isEnum) {
-                innerClasses.map { "class ${it.name} ...\n\n".prependIndent("  ") }.sorted().joinToString("") +
+                append(
-                "}"
+                    fields
+                        .filterIsInstance<PsiEnumConstant>()
+                        .joinToString(",\n") { it.renderEnumConstant() }.prependDefaultIndent()
+                )
+                append(";\n\n")
+            }
+
+            append(renderMembers())
+            append("}")
+        }
     }
 
+    private fun PsiClass.renderMembers(): String {
+        return buildString {
+            appendSorted(
+                fields
+                    .filterNot { it is PsiEnumConstant }
+                    .map { it.renderVar().prependDefaultIndent() + ";\n\n" }
+            )
+
+            appendSorted(
+                methods
+                    .map { it.renderMethod().prependDefaultIndent() + "\n\n" }
+            )
+
+            appendSorted(
+                innerClasses.map { "class ${it.name} ...\n\n".prependDefaultIndent() }
+            )
+        }
+    }
+
+    private fun StringBuilder.appendSorted(list: List<String>) {
+        append(list.sorted().joinToString(""))
+    }
+
+    private fun String.prependDefaultIndent() = prependIndent("  ")
 }

idea/testData/kotlinAndJavaChecker/javaWithKotlin/OverrideInEnumEntry.java

@@ -0,0 +1,5 @@
+package test;
+
+public interface <lineMarker>OverrideInEnumEntry</lineMarker> {
+    void <lineMarker>foo</lineMarker>();
+}

idea/testData/kotlinAndJavaChecker/javaWithKotlin/OverrideInEnumEntry.kt

@@ -0,0 +1,7 @@
+package test
+
+enum class E : OverrideInEnumEntry {
+    X {
+        override fun foo() {}
+    }
+}

idea/testData/kotlinAndJavaChecker/javaWithKotlin/OverrideInEnumEntry.txt

@@ -0,0 +1 @@
+// LINE_MARKERS

idea/tests/org/jetbrains/kotlin/checkers/AbstractJavaAgainstKotlinCheckerTest.java

@@ -29,6 +29,7 @@ import com.siyeh.ig.bugs.StaticCallOnSubclassInspection;
 import com.siyeh.ig.bugs.StaticFieldReferenceOnSubclassInspection;
 import kotlin.collections.CollectionsKt;
 import kotlin.jvm.functions.Function1;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.idea.KotlinDaemonAnalyzerTestCase;
 import org.jetbrains.kotlin.idea.test.ConfigLibraryUtil;
@@ -64,7 +65,11 @@ public abstract class AbstractJavaAgainstKotlinCheckerTest extends KotlinDaemonA
 
     @Nullable
     protected String getConfigFileText() {
-        File configureFile = new File(PluginTestCaseBase.getTestDataPathBase() + "/kotlinAndJavaChecker/javaAgainstKotlin/" + getTestName(false) + ".txt");
+        String className = getClass().getName();
+
+        boolean isJavaWithKotlin = className.contains("JavaWithKotlin");
+
+        File configureFile = new File(configPath(isJavaWithKotlin ? "javaWithKotlin" : "javaAgainstKotlin"));
         if (!configureFile.exists()) return null;
 
         try {
@@ -75,6 +80,11 @@ public abstract class AbstractJavaAgainstKotlinCheckerTest extends KotlinDaemonA
         }
     }
 
+    @NotNull
+    private String configPath(String testKind) {
+        return PluginTestCaseBase.getTestDataPathBase() + "/kotlinAndJavaChecker/" + testKind + "/" + getTestName(false) + ".txt";
+    }
+
     @Override
     protected LocalInspectionTool[] configureLocalInspectionTools() {
         String configFileText = getConfigFileText();
@@ -90,6 +100,14 @@ public abstract class AbstractJavaAgainstKotlinCheckerTest extends KotlinDaemonA
         }), LocalInspectionTool.class);
     }
 
+    @Override
+    protected boolean doTestLineMarkers() {
+        String configFileText = getConfigFileText();
+        if (configFileText == null) return super.doTestLineMarkers();
+
+        return InTextDirectivesUtils.isDirectiveDefined(configFileText, "LINE_MARKERS");
+    }
+
     @Override
     protected Module createMainModule() throws IOException {
         Module module = super.createMainModule();

idea/tests/org/jetbrains/kotlin/checkers/JavaAgainstKotlinSourceCheckerTestGenerated.java

@@ -213,5 +213,10 @@ public class JavaAgainstKotlinSourceCheckerTestGenerated extends AbstractJavaAga
         public void testNoNotNullOnParameterInOverride() throws Exception {
             runTest("idea/testData/kotlinAndJavaChecker/javaWithKotlin/NoNotNullOnParameterInOverride.kt");
         }
+
+        @TestMetadata("OverrideInEnumEntry.kt")
+        public void testOverrideInEnumEntry() throws Exception {
+            runTest("idea/testData/kotlinAndJavaChecker/javaWithKotlin/OverrideInEnumEntry.kt");
+        }
     }
 }

idea/tests/org/jetbrains/kotlin/checkers/JavaAgainstKotlinSourceCheckerWithUltraLightTestGenerated.java

@@ -213,5 +213,10 @@ public class JavaAgainstKotlinSourceCheckerWithUltraLightTestGenerated extends A
         public void testNoNotNullOnParameterInOverride() throws Exception {
             runTest("idea/testData/kotlinAndJavaChecker/javaWithKotlin/NoNotNullOnParameterInOverride.kt");
         }
+
+        @TestMetadata("OverrideInEnumEntry.kt")
+        public void testOverrideInEnumEntry() throws Exception {
+            runTest("idea/testData/kotlinAndJavaChecker/javaWithKotlin/OverrideInEnumEntry.kt");
+        }
     }
 }