Extraction Engine: When extracting to enum class body, place new declaration after the last entry

#KT-9629 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/KtClassBody.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.stubs.KotlinPlaceHolderStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes.*
+import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull
 import java.util.*
 
 class KtClassBody : KtElementImplStub<KotlinPlaceHolderStub<KtClassBody>>, KtDeclarationContainer {
@@ -57,4 +58,22 @@ class KtClassBody : KtElementImplStub<KotlinPlaceHolderStub<KtClassBody>>, KtDec
      */
     val danglingAnnotations: List<KtAnnotationEntry>
         get() = getStubOrPsiChildrenAsList(MODIFIER_LIST).flatMap { it.annotationEntries }
+
+    private fun addEnumEntriesTerminator() {
+        val klass = parent as? KtClass ?: return
+        if (!klass.isEnum()) return
+
+        val lastEntry = klass.declarations.lastIsInstanceOrNull<KtEnumEntry>()
+        if (KtPsiUtil.skipTrailingWhitespacesAndComments(lastEntry ?: firstChild)?.node?.elementType == KtTokens.SEMICOLON) return
+        super.addAfter(KtPsiFactory(this).createSemicolon(), lastEntry)
+    }
+
+    override fun addAfter(element: PsiElement, anchor: PsiElement?): PsiElement? {
+        return super.addAfter(element, anchor)
+                .apply {
+                    if (element !is KtEnumEntry) {
+                        addEnumEntriesTerminator()
+                    }
+                }
+    }
 }

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractableAnalysisUtil.kt

@@ -865,7 +865,10 @@ fun ExtractionData.getDefaultVisibility(): String {
     if (!isVisibilityApplicable()) return ""
 
     val parent = targetSibling.getStrictParentOfType<KtDeclaration>()
-    if (parent is KtClass && parent.isInterface()) return ""
+    if (parent is KtClass) {
+        if (parent.isInterface()) return ""
+        if (parent.isEnum() && commonParent.getNonStrictParentOfType<KtEnumEntry>()?.getStrictParentOfType<KtClass>() == parent) return ""
+    }
 
     return "private"
 }

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractorUtil.kt

@@ -591,18 +591,19 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
 
         return with(descriptor.extractionData) {
             val targetContainer = anchor.getParent()!!
+            // TODO: Get rid of explicit new-lines in favor of formatter rules
             val emptyLines = psiFactory.createWhiteSpace("\n\n")
             if (insertBefore) {
-                val declarationInFile = targetContainer.addBefore(declaration, anchor) as KtNamedDeclaration
-                targetContainer.addBefore(emptyLines, anchor)
-
-                declarationInFile
+                (targetContainer.addBefore(declaration, anchor) as KtNamedDeclaration).apply {
+                    targetContainer.addBefore(emptyLines, anchor)
+                }
             }
             else {
-                val declarationInFile = targetContainer.addAfter(declaration, anchor) as KtNamedDeclaration
-                targetContainer.addAfter(emptyLines, anchor)
-
-                declarationInFile
+                (targetContainer.addAfter(declaration, anchor) as KtNamedDeclaration).apply {
+                    if (!(targetContainer is KtClassBody && (targetContainer.parent as? KtClass)?.isEnum() ?: false)) {
+                        targetContainer.addAfter(emptyLines, anchor)
+                    }
+                }
             }
         }
     }
@@ -610,8 +611,13 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
     val duplicates = if (generatorOptions.inTempFile) Collections.emptyList() else descriptor.duplicates
 
     val anchor = with(descriptor.extractionData) {
+        val targetParent = targetSibling.parent
+
         val anchorCandidates = duplicates.mapTo(ArrayList<PsiElement>()) { it.range.elements.first() }
         anchorCandidates.add(targetSibling)
+        if (targetSibling is KtEnumEntry) {
+            anchorCandidates.add(targetSibling.siblings().last { it is KtEnumEntry })
+        }
 
         val marginalCandidate = if (insertBefore) {
             anchorCandidates.minBy { it.startOffset }!!
@@ -621,8 +627,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         }
 
         // Ascend to the level of targetSibling
-        val targetParent = targetSibling.getParent()
-        marginalCandidate.parentsWithSelf.first { it.getParent() == targetParent }
+        marginalCandidate.parentsWithSelf.first { it.parent == targetParent }
     }
 
     val shouldInsert = !(generatorOptions.inTempFile || generatorOptions.target == ExtractionTarget.FAKE_LAMBDALIKE_FUNCTION)

idea/testData/refactoring/extractFunction/basic/extractToEnumClassBody.kt

@@ -0,0 +1,10 @@
+// PARAM_DESCRIPTOR: value-parameter val n: kotlin.Int defined in E.A.foo
+// PARAM_TYPES: kotlin.Int
+enum class E {
+    // SIBLING:
+    A {
+        fun foo(n: Int) = <selection>n + 1</selection>
+    },
+    B,
+    C
+}

idea/testData/refactoring/extractFunction/basic/extractToEnumClassBody.kt.after

@@ -0,0 +1,12 @@
+// PARAM_DESCRIPTOR: value-parameter val n: kotlin.Int defined in E.A.foo
+// PARAM_TYPES: kotlin.Int
+enum class E {
+    // SIBLING:
+    A {
+        fun foo(n: Int) = i(n)
+    },
+    B,
+    C;
+
+    fun i(n: Int) = n + 1
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/JetExtractionTestGenerated.java

@@ -510,6 +510,12 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 doExtractFunctionTest(fileName);
             }
 
+            @TestMetadata("extractToEnumClassBody.kt")
+            public void testExtractToEnumClassBody() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/extractToEnumClassBody.kt");
+                doExtractFunctionTest(fileName);
+            }
+
             @TestMetadata("fakeOverride.kt")
             public void testFakeOverride() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/fakeOverride.kt");

j2k/testData/fileOrElement/settings/openByDefault.kt

@@ -31,6 +31,7 @@ internal open class D : I {
 
 internal enum class E {
     ;
+
     fun foo(): Int {
         return 0
     }