Introduce TrailingCommaState

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

@@ -5,21 +5,15 @@
 
 package org.jetbrains.kotlin.idea.formatter
 
-import com.intellij.openapi.util.registry.Registry
-import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiErrorElement
-import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.codeStyle.CodeStyleSettings
 import com.intellij.psi.tree.TokenSet
-import com.intellij.psi.util.PsiUtil.getElementType
-import org.jetbrains.kotlin.config.LanguageFeature
-import org.jetbrains.kotlin.idea.project.languageVersionSettings
+import com.intellij.psi.util.PsiUtil
 import org.jetbrains.kotlin.idea.util.*
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getPrevSiblingIgnoringWhitespaceAndComments
-import org.jetbrains.kotlin.psi.psiUtil.nextLeaf
 import org.jetbrains.kotlin.psi.psiUtil.prevLeaf
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.utils.addToStdlib.cast
@@ -52,23 +46,26 @@ object TrailingCommaHelper {
                 commaOwner.isMultiline()
     }
 
+    fun trailingCommaExists(commaOwner: KtElement): Boolean = when (commaOwner) {
+        is KtFunctionLiteral -> commaOwner.valueParameterList?.trailingComma != null
+        is KtWhenEntry -> commaOwner.trailingComma != null
+        is KtDestructuringDeclaration -> commaOwner.trailingComma != null
+        else -> trailingCommaOrLastElement(commaOwner)?.isComma == true
+    }
+
     fun trailingCommaOrLastElement(commaOwner: KtElement): PsiElement? {
         val lastChild = commaOwner.lastSignificantChild ?: return null
-        val withSelf = when (getElementType(lastChild)) {
+        val withSelf = when (PsiUtil.getElementType(lastChild)) {
             KtTokens.COMMA -> return lastChild
             in RIGHT_BARRIERS -> false
             else -> true
         }
 
         return lastChild.getPrevSiblingIgnoringWhitespaceAndComments(withSelf)?.takeIf {
-            getElementType(it) !in LEFT_BARRIERS
+            PsiUtil.getElementType(it) !in LEFT_BARRIERS
         }?.takeIfIsNotError()
     }
 
-    fun trailingCommaAllowedInModule(source: PsiElement): Boolean =
-        Registry.`is`("kotlin.formatter.allowTrailingCommaInAnyProject", false) ||
-                source.module?.languageVersionSettings?.supportsFeature(LanguageFeature.TrailingCommas) == true
-
     fun elementBeforeFirstElement(commaOwner: KtElement): PsiElement? = when (commaOwner) {
         is KtParameterList -> {
             val parent = commaOwner.parent
@@ -100,15 +97,3 @@ object TrailingCommaHelper {
             else -> lastChild
         }
 }
-
-fun PsiElement.leafIgnoringWhitespace(forward: Boolean = true, skipEmptyElements: Boolean = true) =
-    leaf(forward) { (!skipEmptyElements || it.textLength != 0) && it !is PsiWhiteSpace }
-
-fun PsiElement.leafIgnoringWhitespaceAndComments(forward: Boolean = true, skipEmptyElements: Boolean = true) =
-    leaf(forward) { (!skipEmptyElements || it.textLength != 0) && it !is PsiWhiteSpace && it !is PsiComment }
-
-fun PsiElement.leaf(forward: Boolean = true, filter: (PsiElement) -> Boolean): PsiElement? =
-    if (forward) nextLeaf(filter)
-    else prevLeaf(filter)
-
-val PsiElement.isComma: Boolean get() = getElementType(this) == KtTokens.COMMA

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

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.formatter
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.util.canAddTrailingComma
+import org.jetbrains.kotlin.idea.util.containsLineBreakInThis
+import org.jetbrains.kotlin.idea.util.isMultiline
+import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtWhenEntry
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
+
+enum class TrailingCommaState {
+    /**
+     * The trailing comma is needed and exists
+     */
+    EXISTS,
+
+    /**
+     * The trailing comma is needed and doesn't exists
+     */
+    MISSING,
+
+    /**
+     * The trailing comma isn't needed and doesn't exists
+     */
+    NOT_EXISTS,
+
+    /**
+     * The trailing comma isn't needed, but exists
+     */
+    REDUNDANT,
+
+    /**
+     * The trailing comma isn't applicable for this element
+     */
+    NOT_APPLICABLE,
+    ;
+
+    companion object {
+        fun stateForElement(element: PsiElement): TrailingCommaState = when {
+            element !is KtElement || !element.canAddTrailingComma() -> NOT_APPLICABLE
+            isMultiline(element) ->
+                if (TrailingCommaHelper.trailingCommaExists(element))
+                    EXISTS
+                else
+                    MISSING
+            else ->
+                if (TrailingCommaHelper.trailingCommaExists(element))
+                    REDUNDANT
+                else
+                    NOT_EXISTS
+        }
+    }
+}
+
+private fun isMultiline(ktElement: KtElement): Boolean = when {
+    ktElement.parent is KtFunctionLiteral -> isMultiline(ktElement.parent as KtElement)
+
+    ktElement is KtFunctionLiteral -> ktElement.isMultiline(
+        startOffsetGetter = { valueParameterList?.startOffset },
+        endOffsetGetter = { arrow?.endOffset },
+    )
+
+    ktElement is KtWhenEntry -> ktElement.isMultiline(
+        startOffsetGetter = { startOffset },
+        endOffsetGetter = { arrow?.endOffset },
+    )
+
+    ktElement is KtDestructuringDeclaration -> ktElement.isMultiline(
+        startOffsetGetter = { lPar?.startOffset },
+        endOffsetGetter = { rPar?.endOffset },
+    )
+
+    else -> ktElement.isMultiline()
+}
+
+private fun <T : PsiElement> T.isMultiline(
+    startOffsetGetter: T.() -> Int?,
+    endOffsetGetter: T.() -> Int?,
+): Boolean {
+    val startOffset = startOffsetGetter() ?: startOffset
+    val endOffset = endOffsetGetter() ?: endOffset
+    return containsLineBreakInThis(startOffset, endOffset)
+}

idea/formatter/src/org/jetbrains/kotlin/idea/util/FormatterUtil.kt

@@ -11,20 +11,26 @@ import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.util.UserDataHolder
 import com.intellij.openapi.util.registry.Registry
 import com.intellij.openapi.util.text.StringUtil
+import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.codeStyle.CodeStyleSettings
 import com.intellij.psi.tree.IElementType
 import com.intellij.psi.tree.TokenSet
+import com.intellij.psi.util.PsiUtil
 import com.intellij.psi.util.PsiUtilCore
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.idea.core.formatter.KotlinCodeStyleSettings
 import org.jetbrains.kotlin.idea.formatter.kotlinCustomSettings
+import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
 import org.jetbrains.kotlin.psi.KtFunctionLiteral
 import org.jetbrains.kotlin.psi.KtWhenEntry
 import org.jetbrains.kotlin.psi.KtWhenExpression
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.nextLeaf
+import org.jetbrains.kotlin.psi.psiUtil.prevLeaf
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 
@@ -56,6 +62,20 @@ fun PsiElement.getLineCount(): Int {
 
 fun PsiElement.isMultiline() = getLineCount() > 1
 
+fun PsiElement?.isLineBreak() = this is PsiWhiteSpace && StringUtil.containsLineBreak(text)
+
+fun PsiElement.leafIgnoringWhitespace(forward: Boolean = true, skipEmptyElements: Boolean = true) =
+    leaf(forward) { (!skipEmptyElements || it.textLength != 0) && it !is PsiWhiteSpace }
+
+fun PsiElement.leafIgnoringWhitespaceAndComments(forward: Boolean = true, skipEmptyElements: Boolean = true) =
+    leaf(forward) { (!skipEmptyElements || it.textLength != 0) && it !is PsiWhiteSpace && it !is PsiComment }
+
+fun PsiElement.leaf(forward: Boolean = true, filter: (PsiElement) -> Boolean): PsiElement? =
+    if (forward) nextLeaf(filter)
+    else prevLeaf(filter)
+
+val PsiElement.isComma: Boolean get() = PsiUtil.getElementType(this) == KtTokens.COMMA
+
 fun KtFunctionLiteral.needTrailingComma(settings: CodeStyleSettings?, checkExistingTrailingComma: Boolean = true): Boolean =
     needTrailingComma(
         settings = settings,
@@ -79,7 +99,7 @@ fun KtDestructuringDeclaration.needTrailingComma(settings: CodeStyleSettings?, c
         globalEndOffset = { rPar?.endOffset },
     )
 
-fun <T : PsiElement> T.needTrailingComma(
+private fun <T : PsiElement> T.needTrailingComma(
     settings: CodeStyleSettings?,
     trailingComma: T.() -> PsiElement?,
     additionalCheck: () -> Boolean = { true },
@@ -136,9 +156,12 @@ private fun elementType(userDataHolder: UserDataHolder): IElementType? = when (u
     else -> null
 }
 
-fun PsiElement.canAddTrailingComma(): Boolean = canAddTrailingComma(this)
+fun ASTNode.canAddTrailingComma(): Boolean = psi?.canAddTrailingComma() == true
 
-fun ASTNode.canAddTrailingComma(): Boolean = canAddTrailingComma(this)
-
-private fun canAddTrailingComma(userDataHolder: UserDataHolder): Boolean = elementType(userDataHolder) in TYPES_WITH_TRAILING_COMMA
+fun PsiElement.canAddTrailingComma(): Boolean =
+    if (this is KtWhenEntry && (isElse || parent.cast<KtWhenExpression>().leftParenthesis == null))
+        false
+    else
+        canAddTrailingComma(this)
 
+private fun canAddTrailingComma(userDataHolder: UserDataHolder): Boolean = elementType(userDataHolder) in TYPES_WITH_TRAILING_COMMA

idea/src/org/jetbrains/kotlin/idea/codeInsight/upDownMover/KotlinExpressionMover.kt

@@ -12,7 +12,7 @@ import com.intellij.psi.*
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.idea.core.util.isMultiLine
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaHelper
-import org.jetbrains.kotlin.idea.formatter.isComma
+import org.jetbrains.kotlin.idea.util.isComma
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypesAndPredicate

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

@@ -19,8 +19,9 @@ import com.intellij.psi.impl.source.codeStyle.PostFormatProcessorHelper
 import com.intellij.psi.util.PsiUtil.getElementType
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaHelper.findInvalidCommas
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaHelper.needComma
-import org.jetbrains.kotlin.idea.formatter.TrailingCommaHelper.trailingCommaAllowedInModule
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaHelper.trailingCommaOrLastElement
+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

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -7,9 +7,14 @@ package org.jetbrains.kotlin.idea.formatter
 
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.util.TextRange
+import com.intellij.openapi.util.registry.Registry
 import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.codeStyle.CodeStyleManager
+import org.jetbrains.kotlin.config.LanguageFeature
+import org.jetbrains.kotlin.idea.project.languageVersionSettings
+import org.jetbrains.kotlin.idea.util.module
 
 fun PsiFile.commitAndUnblockDocument(): Boolean {
     val virtualFile = this.virtualFile ?: return false
@@ -24,3 +29,7 @@ fun PsiFile.adjustLineIndent(startOffset: Int, endOffset: Int) {
     if (!commitAndUnblockDocument()) return
     CodeStyleManager.getInstance(project).adjustLineIndent(this, TextRange(startOffset, endOffset))
 }
+
+fun trailingCommaAllowedInModule(source: PsiElement): Boolean =
+    Registry.`is`("kotlin.formatter.allowTrailingCommaInAnyProject", false) ||
+            source.module?.languageVersionSettings?.supportsFeature(LanguageFeature.TrailingCommas) == true

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

@@ -18,8 +18,13 @@ 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.*
+import org.jetbrains.kotlin.idea.formatter.TrailingCommaHelper
+import org.jetbrains.kotlin.idea.formatter.TrailingCommaVisitor
+import org.jetbrains.kotlin.idea.formatter.kotlinCustomSettings
+import org.jetbrains.kotlin.idea.formatter.trailingCommaAllowedInModule
+import org.jetbrains.kotlin.idea.util.isComma
 import org.jetbrains.kotlin.idea.util.isLineBreak
+import org.jetbrains.kotlin.idea.util.leafIgnoringWhitespaceAndComments
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.psiUtil.*
 import javax.swing.JComponent
@@ -75,7 +80,7 @@ class TrailingCommaInspection(
         private fun checkTrailingComma(commaOwner: KtElement, action: TrailingCommaAction) {
             val trailingCommaOrLastElement = TrailingCommaHelper.trailingCommaOrLastElement(commaOwner) ?: return
             if (action == TrailingCommaAction.ADD) {
-                if (!TrailingCommaHelper.trailingCommaAllowedInModule(commaOwner) || trailingCommaOrLastElement.isComma) return
+                if (!trailingCommaAllowedInModule(commaOwner) || trailingCommaOrLastElement.isComma) return
                 reportProblem(
                     trailingCommaOrLastElement,
                     KotlinBundle.message("inspection.trailing.comma.missing.trailing.comma"),

idea/src/org/jetbrains/kotlin/idea/util/expressionExt.kt

@@ -7,8 +7,10 @@ package org.jetbrains.kotlin.idea.util
 
 import com.intellij.lang.jvm.JvmModifier
 import com.intellij.openapi.util.TextRange
-import com.intellij.openapi.util.text.StringUtil
-import com.intellij.psi.*
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiMember
+import com.intellij.psi.PsiMethod
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.callExpression
 import org.jetbrains.kotlin.idea.references.mainReference
@@ -58,5 +60,3 @@ fun KtDotQualifiedExpression.hasNotReceiver(): Boolean {
             element is PsiMember && element.hasModifier(JvmModifier.STATIC) ||
             element is PsiMethod && element.isConstructor
 }
-
-fun PsiElement?.isLineBreak() = this is PsiWhiteSpace && StringUtil.containsLineBreak(text)