Trailing commas in enums

CHANGELOG.md

@@ -126,6 +126,7 @@ The callback function provided as parameter to the format function is now called
 * When a glob is specified then ensure that it matches files in the current directory and not only in subdirectories of the current directory ([#1533](https://github.com/pinterest/ktlint/issue/1533)).
 * Execute `ktlint` cli on default kotlin extensions only when an (existing) path to a directory is given. ([#917](https://github.com/pinterest/ktlint/issue/917)).
 * Invoke callback on `format` function for all errors including errors that are autocorrected ([#1491](https://github.com/pinterest/ktlint/issues/1491))
+* `trailing-comma` now handles enums as well
 
 
 ### Changed

ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/TrailingCommaRule.kt

@@ -4,8 +4,11 @@ import com.pinterest.ktlint.core.Rule
 import com.pinterest.ktlint.core.api.EditorConfigProperties
 import com.pinterest.ktlint.core.api.UsesEditorConfigProperties
 import com.pinterest.ktlint.core.ast.ElementType
+import com.pinterest.ktlint.core.ast.ElementType.ENUM_ENTRY
+import com.pinterest.ktlint.core.ast.ElementType.SEMICOLON
 import com.pinterest.ktlint.core.ast.children
 import com.pinterest.ktlint.core.ast.containsLineBreakInRange
+import com.pinterest.ktlint.core.ast.lineNumber
 import com.pinterest.ktlint.core.ast.prevCodeLeaf
 import com.pinterest.ktlint.core.ast.prevLeaf
 import kotlin.properties.Delegates
@@ -15,13 +18,17 @@ import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
 import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtEnumEntry
 import org.jetbrains.kotlin.psi.KtFunctionLiteral
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.KtValueArgumentList
 import org.jetbrains.kotlin.psi.KtWhenEntry
 import org.jetbrains.kotlin.psi.KtWhenExpression
+import org.jetbrains.kotlin.psi.psiUtil.allChildren
 import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.nextLeaf
@@ -90,9 +97,7 @@ public class TrailingCommaRule :
             ElementType.TYPE_PARAMETER_LIST -> visitTypeList(node, emit, autoCorrect)
             ElementType.VALUE_PARAMETER_LIST -> visitValueList(node, emit, autoCorrect)
             ElementType.WHEN_ENTRY -> visitWhenEntry(node, emit, autoCorrect)
-            else -> Unit
-        }
-        when (node.elementType) {
+            ElementType.CLASS -> visitClassEntry(node, emit, autoCorrect)
             ElementType.COLLECTION_LITERAL_EXPRESSION -> visitCollectionLiteralExpression(node, emit, autoCorrect)
             ElementType.INDICES -> visitIndices(node, emit, autoCorrect)
             ElementType.TYPE_ARGUMENT_LIST -> visitTypeList(node, emit, autoCorrect)
@@ -191,6 +196,57 @@ public class TrailingCommaRule :
         node.reportAndCorrectTrailingCommaNodeBefore(inspectNode, emit, autoCorrect)
     }
 
+    private fun visitClassEntry(
+        node: ASTNode,
+        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit,
+        autoCorrect: Boolean
+    ) {
+        val psi = node.psi
+        require(psi is KtClass)
+
+        if (psi.isEnum()) {
+            val classBody = node
+                .children()
+                .singleOrNull { it.psi is KtClassBody }
+                ?: return
+
+            if (classBody.lastTwoEnumEntriesAreOnSameLine()) return // Do nothing when last two entries are on same line
+            val inspectNode = classBody.determineNodeToInspectForEnum() ?: return
+            node.reportAndCorrectTrailingCommaNodeBefore(inspectNode, emit, autoCorrect)
+        }
+    }
+
+    /**
+     * Determines what [ASTNode] to inspect when dealing with Enums.
+     *
+     * If an Enum is terminated by a semicolon, that will be inspected, otherwise the element following the last
+     * enum entry (which should be the `}` terminating the enum class)
+     */
+    private fun ASTNode.determineNodeToInspectForEnum(): ASTNode? {
+        val lastEnumEntry = psi
+            .children
+            .last { it is KtEnumEntry }
+
+        val semicolon = lastEnumEntry
+            .node
+            .children()
+            .singleOrNull { it.elementType == SEMICOLON }
+
+        // Inspect the semicolon if it is semicolon terminated, otherwise next non-ignorable node in the tree
+        return semicolon
+            ?: lastEnumEntry.nextLeaf { !it.node.isIgnorable() }!!.node
+    }
+
+    private fun ASTNode.lastTwoEnumEntriesAreOnSameLine(): Boolean {
+        val enumEntries = psi
+            .children
+            .filterIsInstance<KtEnumEntry>()
+            .reversed() // We are interested in the last two ones.
+
+        if (enumEntries.count() < 2) return false
+        return enumEntries[0].node.lineNumber() == enumEntries[1].node.lineNumber()
+    }
+
     private fun ASTNode.reportAndCorrectTrailingCommaNodeBefore(
         inspectNode: ASTNode,
         emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit,
@@ -234,6 +290,7 @@ public class TrailingCommaRule :
                         true
                     )
                 }
+
                 if (autoCorrect) {
                     if (addNewLineBeforeArrowInWhenEntry) {
                         val parentIndent = (prevNode.psi.parent.prevLeaf() as? PsiWhiteSpace)?.text ?: "\n"
@@ -246,6 +303,20 @@ public class TrailingCommaRule :
                             prevNode.psi.parent.addAfter(newLine, prevNode.psi)
                         }
                     }
+
+                    if (inspectNode.treeParent.elementType == ENUM_ENTRY) {
+                        with(KtPsiFactory(prevNode.psi)) {
+                            val parentIndent = (prevNode.psi.parent.prevLeaf() as? PsiWhiteSpace)?.text ?: "\n"
+                            val newline = createWhiteSpace(parentIndent)
+                            val enumEntry = inspectNode.treeParent.psi
+                            enumEntry.apply {
+                                add(newline)
+                                removeChild(inspectNode)
+                                parent.addAfter(createSemicolon(), this)
+                            }
+                        }
+                    }
+
                     val comma = KtPsiFactory(prevNode.psi).createComma()
                     prevNode.psi.parent.addAfter(comma, prevNode.psi)
                 }
@@ -272,6 +343,13 @@ public class TrailingCommaRule :
             false
         }
 
+    private fun ASTNode?.correctTrailingCommaInSemicolonTerminatedEnum() =
+        when {
+            this == null -> false
+            psi is KtEnumEntry -> psi.allChildren.last?.node?.elementType == SEMICOLON
+            else -> false
+        }
+
     private fun ASTNode?.findPreviousTrailingCommaNodeOrNull(): ASTNode? {
         var node = this
         while (node?.isIgnorable() == true) {
@@ -324,7 +402,8 @@ public class TrailingCommaRule :
             ElementType.FUNCTION_TYPE,
             ElementType.TYPE_PARAMETER_LIST,
             ElementType.VALUE_PARAMETER_LIST,
-            ElementType.WHEN_ENTRY
+            ElementType.WHEN_ENTRY,
+            ElementType.CLASS
         )
 
         private val TYPES_ON_CALL_SITE = TokenSet.create(