Get rid of TrailingCommaHelper.needComma

idea/formatter/src/org/jetbrains/kotlin/idea/formatter/trailingComma/TrailingCommaHelper.kt

@@ -11,10 +11,15 @@ import com.intellij.psi.codeStyle.CodeStyleSettings
 import com.intellij.psi.tree.TokenSet
 import com.intellij.psi.util.PsiUtil
 import org.jetbrains.kotlin.idea.formatter.kotlinCustomSettings
-import org.jetbrains.kotlin.idea.util.*
+import org.jetbrains.kotlin.idea.util.isComma
+import org.jetbrains.kotlin.idea.util.isLineBreak
+import org.jetbrains.kotlin.idea.util.leafIgnoringWhitespace
+import org.jetbrains.kotlin.idea.util.leafIgnoringWhitespaceAndComments
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.*
+import org.jetbrains.kotlin.psi.psiUtil.getPrevSiblingIgnoringWhitespaceAndComments
+import org.jetbrains.kotlin.psi.psiUtil.prevLeaf
+import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 
 object TrailingCommaHelper {
@@ -25,26 +30,6 @@ object TrailingCommaHelper {
             it.prevLeaf(true)?.isLineBreak() == true || it.leafIgnoringWhitespace(false) != it.leafIgnoringWhitespaceAndComments(false)
         }?.toList().orEmpty()
 
-    fun needComma(
-        commaOwner: KtElement,
-        settings: CodeStyleSettings?,
-        checkExistingTrailingComma: Boolean = true,
-    ): Boolean = when {
-        commaOwner is KtWhenEntry ->
-            commaOwner.needTrailingComma(settings, checkExistingTrailingComma)
-
-        commaOwner.parent is KtFunctionLiteral ->
-            commaOwner.parent.cast<KtFunctionLiteral>().needTrailingComma(settings, checkExistingTrailingComma)
-
-        commaOwner is KtDestructuringDeclaration ->
-            commaOwner.needTrailingComma(settings, checkExistingTrailingComma)
-
-        else -> (checkExistingTrailingComma &&
-                trailingCommaOrLastElement(commaOwner)?.isComma == true ||
-                settings?.kotlinCustomSettings?.addTrailingCommaIsAllowedFor(commaOwner) != false) &&
-                commaOwner.isMultiline()
-    }
-
     fun trailingCommaExistsOrCanExist(psiElement: PsiElement, settings: CodeStyleSettings): Boolean =
         TrailingCommaContext.create(psiElement).commaExistsOrMayExist(settings.kotlinCustomSettings)
 
@@ -99,45 +84,3 @@ object TrailingCommaHelper {
             else -> lastChild
         }
 }
-
-private fun KtFunctionLiteral.needTrailingComma(settings: CodeStyleSettings?, checkExistingTrailingComma: Boolean = true): Boolean =
-    needTrailingComma(
-        settings = settings,
-        trailingComma = { if (checkExistingTrailingComma) valueParameterList?.trailingComma else null },
-        globalStartOffset = { valueParameterList?.startOffset },
-        globalEndOffset = { arrow?.endOffset },
-    )
-
-private fun KtWhenEntry.needTrailingComma(settings: CodeStyleSettings?, checkExistingTrailingComma: Boolean = true): Boolean =
-    needTrailingComma(
-        settings = settings,
-        trailingComma = { if (checkExistingTrailingComma) trailingComma else null },
-        additionalCheck = { !isElse && parent.cast<KtWhenExpression>().leftParenthesis != null },
-        globalEndOffset = { arrow?.endOffset },
-    )
-
-private fun KtDestructuringDeclaration.needTrailingComma(
-    settings: CodeStyleSettings?,
-    checkExistingTrailingComma: Boolean = true
-): Boolean =
-    needTrailingComma(
-        settings = settings,
-        trailingComma = { if (checkExistingTrailingComma) trailingComma else null },
-        globalStartOffset = { lPar?.startOffset },
-        globalEndOffset = { rPar?.endOffset },
-    )
-
-private fun <T : PsiElement> T.needTrailingComma(
-    settings: CodeStyleSettings?,
-    trailingComma: T.() -> PsiElement?,
-    additionalCheck: () -> Boolean = { true },
-    globalStartOffset: T.() -> Int? = PsiElement::startOffset,
-    globalEndOffset: T.() -> Int? = PsiElement::endOffset,
-): Boolean {
-    if (trailingComma() == null && settings?.kotlinCustomSettings?.addTrailingCommaIsAllowedFor(this) == false) return false
-    if (!additionalCheck()) return false
-
-    val startOffset = globalStartOffset() ?: return false
-    val endOffset = globalEndOffset() ?: return false
-    return containsLineBreakInThis(startOffset, endOffset)
-}

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

@@ -16,13 +16,12 @@ import com.intellij.psi.codeStyle.CodeStyleManager
 import com.intellij.psi.codeStyle.CodeStyleSettings
 import com.intellij.psi.impl.source.codeStyle.PostFormatProcessor
 import com.intellij.psi.impl.source.codeStyle.PostFormatProcessorHelper
-import com.intellij.psi.util.PsiUtil.getElementType
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper.findInvalidCommas
-import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper.needComma
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper.trailingCommaOrLastElement
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaState
+import org.jetbrains.kotlin.idea.formatter.trailingComma.addTrailingCommaIsAllowedFor
 import org.jetbrains.kotlin.idea.util.leafIgnoringWhitespace
 import org.jetbrains.kotlin.idea.util.leafIgnoringWhitespaceAndComments
-import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
@@ -52,25 +51,26 @@ private class TrailingCommaPostFormatVisitor(val settings: CodeStyleSettings) :
     private fun processCommaOwner(parent: KtElement) {
         if (!postFormatIsEnable(parent)) return
 
-        val lastElement = trailingCommaOrLastElement(parent) ?: return
+        val lastElementOrComma = trailingCommaOrLastElement(parent) ?: return
-        val elementType = getElementType(lastElement)
         updatePsi(parent) {
+            val stateForElement = TrailingCommaState.stateForElement(parent)
             when {
-                needComma(parent, settings, false) -> {
+                stateForElement == TrailingCommaState.MISSING && settings.kotlinCustomSettings.addTrailingCommaIsAllowedFor(parent) -> {
                     // add a missing comma
-                    if (elementType != KtTokens.COMMA && trailingCommaAllowedInModule(parent)) {
+                    if (trailingCommaAllowedInModule(parent)) {
-                        lastElement.addCommaAfter(KtPsiFactory(parent))
+                        lastElementOrComma.addCommaAfter(KtPsiFactory(parent))
                     }
 
                     correctCommaPosition(parent)
                 }
-                needComma(parent, settings) -> {
+                stateForElement == TrailingCommaState.EXISTS -> {
                     correctCommaPosition(parent)
                 }
-                elementType == KtTokens.COMMA -> {
+                stateForElement == TrailingCommaState.REDUNDANT -> {
                     // remove redundant comma
-                    lastElement.delete()
+                    lastElementOrComma.delete()
                 }
+                else -> Unit
             }
         }
     }

idea/src/org/jetbrains/kotlin/idea/inspections/TrailingCommaInspection.kt

@@ -18,9 +18,10 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.codeStyle.CodeStyleManager
 import org.jetbrains.kotlin.idea.KotlinBundle
-import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaVisitor
 import org.jetbrains.kotlin.idea.formatter.kotlinCustomSettings
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaState
 import org.jetbrains.kotlin.idea.formatter.trailingCommaAllowedInModule
 import org.jetbrains.kotlin.idea.util.isComma
 import org.jetbrains.kotlin.idea.util.isLineBreak
@@ -40,13 +41,16 @@ class TrailingCommaInspection(
 
         override fun process(commaOwner: KtElement) {
             useTrailingComma = CodeStyle.getSettings(commaOwner.project).kotlinCustomSettings.ALLOW_TRAILING_COMMA
-            val action = TrailingCommaAction.create(commaOwner)
+            val state = TrailingCommaState.stateForElement(commaOwner)
-            if (action != TrailingCommaAction.REMOVE) {
+            when (state) {
-                checkCommaPosition(commaOwner)
+                TrailingCommaState.MISSING, TrailingCommaState.EXISTS -> {
-                checkLineBreaks(commaOwner)
+                    checkCommaPosition(commaOwner)
+                    checkLineBreaks(commaOwner)
+                }
+                else -> Unit
             }
 
-            checkTrailingComma(commaOwner, action)
+            checkTrailingComma(commaOwner, state)
         }
 
         private fun checkLineBreaks(commaOwner: KtElement) {
@@ -77,18 +81,17 @@ class TrailingCommaInspection(
             }
         }
 
-        private fun checkTrailingComma(commaOwner: KtElement, action: TrailingCommaAction) {
+        private fun checkTrailingComma(commaOwner: KtElement, state: TrailingCommaState) {
             val trailingCommaOrLastElement = TrailingCommaHelper.trailingCommaOrLastElement(commaOwner) ?: return
-            if (action == TrailingCommaAction.ADD) {
+            if (state == TrailingCommaState.MISSING) {
-                if (!trailingCommaAllowedInModule(commaOwner) || trailingCommaOrLastElement.isComma) return
+                if (!trailingCommaAllowedInModule(commaOwner)) return
                 reportProblem(
                     trailingCommaOrLastElement,
                     KotlinBundle.message("inspection.trailing.comma.missing.trailing.comma"),
                     KotlinBundle.message("inspection.trailing.comma.add.trailing.comma"),
                     if (addCommaWarning) ProblemHighlightType.GENERIC_ERROR_OR_WARNING else ProblemHighlightType.INFORMATION,
                 )
-            } else if (action == TrailingCommaAction.REMOVE) {
+            } else if (state == TrailingCommaState.REDUNDANT) {
-                if (!trailingCommaOrLastElement.isComma) return
                 reportProblem(
                     trailingCommaOrLastElement,
                     KotlinBundle.message("inspection.trailing.comma.useless.trailing.comma"),
@@ -181,15 +184,3 @@ class TrailingCommaInspection(
         return panel
     }
 }
-
-private enum class TrailingCommaAction {
-    ADD, REFORMAT, REMOVE;
-
-    companion object {
-        fun create(commaOwner: KtElement): TrailingCommaAction = when {
-            TrailingCommaHelper.needComma(commaOwner, null, checkExistingTrailingComma = false) -> ADD
-            TrailingCommaHelper.needComma(commaOwner, null, checkExistingTrailingComma = true) -> REFORMAT
-            else -> REMOVE
-        }
-    }
-}