Create from Usage: Add commas after enum entries

#KT-10800 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -142,6 +142,8 @@ fun PsiElement.getNextSiblingIgnoringWhitespaceAndComments(): PsiElement? {
 
 inline fun <reified T : PsiElement> T.nextSiblingOfSameType() = PsiTreeUtil.getNextSiblingOfType(this, T::class.java)
 
+inline fun <reified T : PsiElement> T.prevSiblingOfSameType() = PsiTreeUtil.getPrevSiblingOfType(this, T::class.java)
+
 fun PsiElement?.isAncestor(element: PsiElement, strict: Boolean = false): Boolean {
     return PsiTreeUtil.isAncestor(this, element, strict)
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/quickfix/generateUtil.kt

@@ -187,7 +187,7 @@ fun <T : KtDeclaration> insertMembersAfter(
                     else if (bound == null && body.declarations.isNotEmpty()) {
                         afterAnchor = body.lBrace!!
                     }
-                    else if (bound != null && afterAnchor.startOffset >= bound.startOffset) {
+                    else if (bound != null && afterAnchor.startOffset > bound.startOffset) {
                         afterAnchor = bound.prevSibling!!
                     }
                 }
@@ -207,6 +207,6 @@ fun <T : KtDeclaration> insertMembersAfter(
     }
 }
 
-fun <T : KtDeclaration> insertMember(editor: Editor, classOrObject: KtClassOrObject, declaration: T): T {
-    return insertMembersAfter(editor, classOrObject, listOf(declaration)).single()
+fun <T : KtDeclaration> insertMember(editor: Editor?, classOrObject: KtClassOrObject, declaration: T, anchor: PsiElement? = null): T {
+    return insertMembersAfter(editor, classOrObject, listOf(declaration), anchor).single()
 }

idea/src/org/jetbrains/kotlin/idea/formatter/KotlinPreFormatProcessor.kt

@@ -18,29 +18,55 @@ package org.jetbrains.kotlin.idea.formatter
 
 import com.intellij.lang.ASTNode
 import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
 import com.intellij.psi.impl.source.codeStyle.PreFormatProcessor
+import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
+import org.jetbrains.kotlin.psi.psiUtil.nextSiblingOfSameType
+import org.jetbrains.kotlin.psi.psiUtil.prevSiblingOfSameType
 import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull
 
 class KotlinPreFormatProcessor : PreFormatProcessor {
     class Visitor(var range: TextRange) : KtTreeVisitorVoid() {
+        private fun PsiElement.containsToken(type: IElementType) = allChildren.any { it.node.elementType == type }
+
         override fun visitNamedDeclaration(declaration: KtNamedDeclaration) {
             if (!range.contains(declaration.textRange)) return
 
-            if (declaration is KtEnumEntry) return
             val classBody = declaration.parent as? KtClassBody ?: return
             val klass = classBody.parent as? KtClass ?: return
             if (!klass.isEnum()) return
 
-            val lastEntry = klass.declarations.lastIsInstanceOrNull<KtEnumEntry>()
-            if (lastEntry != null && lastEntry.allChildren.any { it.node.elementType == KtTokens.SEMICOLON }) return
-            if (lastEntry == null && classBody.allChildren.any { it.node.elementType == KtTokens.SEMICOLON }) return
+            var delta = 0
 
-            val semicolon = KtPsiFactory(klass).createSemicolon()
-            classBody.addAfter(semicolon, lastEntry)
-            range = TextRange(range.startOffset, range.endOffset + semicolon.textLength)
+            if (declaration is KtEnumEntry) {
+                val comma = KtPsiFactory(klass).createComma()
+
+                declaration.nextSiblingOfSameType()?.let { nextEntry ->
+                    if (declaration.containsToken(KtTokens.COMMA)) return@let
+                    classBody.addAfter(comma, declaration)
+                    delta += comma.textLength
+                }
+
+                declaration.prevSiblingOfSameType()?.let { prevEntry ->
+                    if (prevEntry.containsToken(KtTokens.COMMA)) return@let
+                    classBody.addAfter(comma, prevEntry)
+                    delta += comma.textLength
+                }
+            }
+            else {
+                val lastEntry = klass.declarations.lastIsInstanceOrNull<KtEnumEntry>()
+                if (lastEntry != null && lastEntry.containsToken(KtTokens.SEMICOLON)) return
+                if (lastEntry == null && classBody.containsToken(KtTokens.SEMICOLON)) return
+
+                val semicolon = KtPsiFactory(klass).createSemicolon()
+                classBody.addAfter(semicolon, lastEntry)
+                delta += semicolon.textLength
+            }
+
+            range = TextRange(range.startOffset, range.endOffset + delta)
         }
     }
 

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt

@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils
 import org.jetbrains.kotlin.idea.core.CollectingNameValidator
 import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.createClass.ClassKind
+import org.jetbrains.kotlin.idea.quickfix.insertMember
 import org.jetbrains.kotlin.idea.refactoring.*
 import org.jetbrains.kotlin.idea.util.DialogWithEditor
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
@@ -576,7 +577,7 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                     }
 
                     containingElement is KtClassOrObject -> {
-                        addDeclarationToClassOrObject(containingElement, declaration)
+                        insertMember(null, containingElement, declaration, containingElement.declarations.lastOrNull())
                     }
                     else -> throw AssertionError("Invalid containing element: ${containingElement.text}")
                 }

idea/testData/quickfix/createFromUsage/createClass/referenceExpression/enumEntryAddComma.kt

@@ -0,0 +1,8 @@
+// "Create enum constant 'PUBLIC'" "true"
+enum class TopicState {
+    PRIVATE
+}
+
+fun foo() {
+    TopicState.<caret>PUBLIC
+}

idea/testData/quickfix/createFromUsage/createClass/referenceExpression/enumEntryAddComma.kt.after

@@ -0,0 +1,10 @@
+// "Create enum constant 'PUBLIC'" "true"
+enum class TopicState {
+    PRIVATE,
+
+    <caret>PUBLIC
+}
+
+fun foo() {
+    TopicState.PUBLIC
+}

idea/testData/quickfix/createFromUsage/createClass/referenceExpression/enumEntryHasComma.kt

@@ -0,0 +1,8 @@
+// "Create enum constant 'PUBLIC'" "true"
+enum class TopicState {
+    PRIVATE,
+}
+
+fun foo() {
+    TopicState.<caret>PUBLIC
+}

idea/testData/quickfix/createFromUsage/createClass/referenceExpression/enumEntryHasComma.kt.after

@@ -0,0 +1,10 @@
+// "Create enum constant 'PUBLIC'" "true"
+enum class TopicState {
+    PRIVATE,
+
+    <caret>PUBLIC
+}
+
+fun foo() {
+    TopicState.PUBLIC
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -1448,6 +1448,18 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("enumEntryAddComma.kt")
+                public void testEnumEntryAddComma() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createClass/referenceExpression/enumEntryAddComma.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("enumEntryHasComma.kt")
+                public void testEnumEntryHasComma() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createClass/referenceExpression/enumEntryHasComma.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("enumEntryNoReceiver.kt")
                 public void testEnumEntryNoReceiver() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createClass/referenceExpression/enumEntryNoReceiver.kt");