Formatter: Use formatter to automatically insert semicolons after the last enum entry when needed

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

@@ -24,7 +24,6 @@ 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 {
@@ -58,22 +57,4 @@ 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/META-INF/plugin.xml

@@ -305,6 +305,7 @@
 
     <lang.foldingBuilder language="kotlin" implementationClass="org.jetbrains.kotlin.idea.KotlinFoldingBuilder"/>
     <lang.formatter language="kotlin" implementationClass="org.jetbrains.kotlin.idea.formatter.JetFormattingModelBuilder"/>
+    <preFormatProcessor implementation="org.jetbrains.kotlin.idea.formatter.KotlinPreFormatProcessor"/>
     <lang.findUsagesProvider language="kotlin" implementationClass="org.jetbrains.kotlin.idea.findUsages.JetFindUsagesProvider"/>
     <lang.elementManipulator forClass="org.jetbrains.kotlin.psi.KtStringTemplateExpression"
                              implementationClass="org.jetbrains.kotlin.psi.psiUtil.KtStringTemplateExpressionManipulator"/>

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

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.formatter
+
+import com.intellij.lang.ASTNode
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.impl.source.codeStyle.PreFormatProcessor
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull
+
+class KotlinPreFormatProcessor : PreFormatProcessor {
+    class Visitor(var range: TextRange) : KtTreeVisitorVoid() {
+        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 (KtPsiUtil.skipTrailingWhitespacesAndComments(lastEntry ?: classBody.firstChild)?.node?.elementType == KtTokens.SEMICOLON) return
+
+            val semicolon = KtPsiFactory(klass).createSemicolon()
+            classBody.addAfter(semicolon, lastEntry)
+            range = TextRange(range.startOffset, range.endOffset + semicolon.textLength)
+        }
+    }
+
+    override fun process(element: ASTNode, range: TextRange): TextRange {
+        val psi = element.psi ?: return range
+        if (!psi.isValid) return range
+        if (psi.containingFile !is KtFile) return range
+        return Visitor(range).apply { psi.accept(this) }.range
+    }
+}