Introduce TrailingCommaContext

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

@@ -19,15 +19,16 @@ import com.intellij.psi.tree.TokenSet
 import org.jetbrains.kotlin.KtNodeTypes.*
 import org.jetbrains.kotlin.idea.core.formatter.KotlinCodeStyleSettings
 import org.jetbrains.kotlin.idea.formatter.NodeIndentStrategy.Companion.strategy
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper.trailingCommaExistsOrCanExist
 import org.jetbrains.kotlin.idea.formatter.trailingComma.addTrailingCommaIsAllowedFor
-import org.jetbrains.kotlin.idea.formatter.trailingComma.needTrailingComma
-import org.jetbrains.kotlin.idea.util.*
+import org.jetbrains.kotlin.idea.util.containsLineBreakInThis
+import org.jetbrains.kotlin.idea.util.isMultiline
+import org.jetbrains.kotlin.idea.util.requireNode
 import org.jetbrains.kotlin.kdoc.lexer.KDocTokens
 import org.jetbrains.kotlin.kdoc.parser.KDocElementTypes
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
-import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 private val QUALIFIED_OPERATION = TokenSet.create(DOT, SAFE_ACCESS)
@@ -568,36 +569,31 @@ abstract class KotlinCommonBlock(
                         additionalWrap = trailingCommaWrappingStrategyWithMultiLineCheck(LPAR, RPAR),
                     )
                     FUNCTION_TYPE -> return defaultTrailingCommaWrappingStrategy(LPAR, RPAR)
-                    FUNCTION_LITERAL -> {
-                        if (nodePsi.parent?.safeAs<KtFunctionLiteral>()?.needTrailingComma(settings) == true) {
-                            val check = thisOrPrevIsMultiLineElement(COMMA, LBRACE /* not necessary */, ARROW /* not necessary */)
-                            return { childElement ->
-                                createWrapAlwaysIf(getSiblingWithoutWhitespaceAndComments(childElement) == null || check(childElement))
-                            }
+                    FUNCTION_LITERAL -> if (trailingCommaExistsOrCanExist(nodePsi.parent, settings)) {
+                        val check = thisOrPrevIsMultiLineElement(COMMA, LBRACE /* not necessary */, ARROW /* not necessary */)
+                        return { childElement ->
+                            createWrapAlwaysIf(getSiblingWithoutWhitespaceAndComments(childElement) == null || check(childElement))
                         }
                     }
                 }
             }
 
-            elementType === FUNCTION_LITERAL -> {
-                if (nodePsi.cast<KtFunctionLiteral>().needTrailingComma(settings))
-                    return trailingCommaWrappingStrategy(leftAnchor = LBRACE, rightAnchor = ARROW)
+            elementType === FUNCTION_LITERAL -> if (trailingCommaExistsOrCanExist(nodePsi, settings)) {
+                return trailingCommaWrappingStrategy(leftAnchor = LBRACE, rightAnchor = ARROW)
             }
 
-            elementType === WHEN_ENTRY -> {
-                // with argument
-                if (nodePsi.cast<KtWhenEntry>().needTrailingComma(settings)) {
-                    val check = thisOrPrevIsMultiLineElement(COMMA, LBRACE /* not necessary */, ARROW /* not necessary */)
-                    return trailingCommaWrappingStrategy(rightAnchor = ARROW) {
-                        getSiblingWithoutWhitespaceAndComments(it, true) != null && check(it)
-                    }
+            // with argument
+            elementType === WHEN_ENTRY -> if (trailingCommaExistsOrCanExist(nodePsi, settings)) {
+                val check = thisOrPrevIsMultiLineElement(COMMA, LBRACE /* not necessary */, ARROW /* not necessary */)
+                return trailingCommaWrappingStrategy(rightAnchor = ARROW) {
+                    getSiblingWithoutWhitespaceAndComments(it, true) != null && check(it)
                 }
             }
 
             elementType === DESTRUCTURING_DECLARATION -> {
                 nodePsi as KtDestructuringDeclaration
                 if (nodePsi.valOrVarKeyword == null) return defaultTrailingCommaWrappingStrategy(LPAR, RPAR)
-                else if (nodePsi.needTrailingComma(settings)) {
+                else if (trailingCommaExistsOrCanExist(nodePsi, settings)) {
                     val check = thisOrPrevIsMultiLineElement(COMMA, LPAR, RPAR)
                     return trailingCommaWrappingStrategy(leftAnchor = LPAR, rightAnchor = RPAR, filter = { it.elementType !== EQ }) {
                         getSiblingWithoutWhitespaceAndComments(it, true) != null && check(it)

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

@@ -0,0 +1,24 @@
+/*
+ * 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.trailingComma
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.core.formatter.KotlinCodeStyleSettings
+
+class TrailingCommaContext private constructor(val element: PsiElement, val state: TrailingCommaState) {
+    companion object {
+        fun create(element: PsiElement): TrailingCommaContext = TrailingCommaContext(
+            element,
+            TrailingCommaState.stateForElement(element),
+        )
+    }
+}
+
+fun TrailingCommaContext.commaExistsOrMayExist(settings: KotlinCodeStyleSettings): Boolean = when (state) {
+    TrailingCommaState.EXISTS -> true
+    TrailingCommaState.MISSING -> settings.addTrailingCommaIsAllowedFor(element)
+    else -> false
+}

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

@@ -14,9 +14,7 @@ import org.jetbrains.kotlin.idea.formatter.kotlinCustomSettings
 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.prevLeaf
-import org.jetbrains.kotlin.psi.psiUtil.siblings
+import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 
 object TrailingCommaHelper {
@@ -47,6 +45,9 @@ object TrailingCommaHelper {
                 commaOwner.isMultiline()
     }
 
+    fun trailingCommaExistsOrCanExist(psiElement: PsiElement, settings: CodeStyleSettings): Boolean =
+        TrailingCommaContext.create(psiElement).commaExistsOrMayExist(settings.kotlinCustomSettings)
+
     fun trailingCommaExists(commaOwner: KtElement): Boolean = when (commaOwner) {
         is KtFunctionLiteral -> commaOwner.valueParameterList?.trailingComma != null
         is KtWhenEntry -> commaOwner.trailingComma != null
@@ -98,3 +99,45 @@ 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/formatter/src/org/jetbrains/kotlin/idea/formatter/trailingComma/TrailingCommaState.kt

@@ -59,6 +59,12 @@ enum class TrailingCommaState {
     }
 }
 
+val TrailingCommaState.existsOrMissing: Boolean
+    get() = when (this) {
+        TrailingCommaState.EXISTS, TrailingCommaState.MISSING -> true
+        else -> false
+    }
+
 private fun isMultiline(ktElement: KtElement): Boolean = when {
     ktElement.parent is KtFunctionLiteral -> isMultiline(ktElement.parent as KtElement)
 

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

@@ -9,60 +9,15 @@ import com.intellij.lang.ASTNode
 import com.intellij.openapi.util.UserDataHolder
 import com.intellij.openapi.util.registry.Registry
 import com.intellij.psi.PsiElement
-import com.intellij.psi.codeStyle.CodeStyleSettings
 import com.intellij.psi.tree.IElementType
 import com.intellij.psi.tree.TokenSet
 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.idea.util.containsLineBreakInThis
-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.startOffset
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 
-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 },
-    )
-
-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 },
-)
-
-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)
-}
-
 fun trailingCommaIsAllowedOnCallSite(): Boolean = Registry.`is`("kotlin.formatter.allowTrailingCommaOnCallSite")
 
 private val TYPES_WITH_TRAILING_COMMA_ON_DECLARATION_SITE = TokenSet.create(

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

@@ -130,11 +130,11 @@ class KotlinExpressionMover : AbstractKotlinUpDownMover() {
             val withTrailingComma = lastElementOnFirstLine.parent
                 ?.safeAs<KtElement>()
                 ?.let {
-                    TrailingCommaHelper.needComma(it, CodeStyle.getSettings(it.project), true)
+                    TrailingCommaHelper.trailingCommaExistsOrCanExist(it, CodeStyle.getSettings(it.project))
                 } == true
 
             fixCommaIfNeeded(lastElementOnFirstLine, down && isLastOfItsKind(lastElementOnSecondLine, true), withTrailingComma)
-            fixCommaIfNeeded(lastElementOnSecondLine, !down && isLastOfItsKind(lastElementOnFirstLine, true),withTrailingComma)
+            fixCommaIfNeeded(lastElementOnSecondLine, !down && isLastOfItsKind(lastElementOnFirstLine, true), withTrailingComma)
             editor.project?.let { PsiDocumentManager.getInstance(it).doPostponedOperationsAndUnblockDocument(editor.document) }
         }
     }

idea/src/org/jetbrains/kotlin/idea/joinLines/JoinWithTrailingCommaHandler.kt

@@ -12,7 +12,9 @@ import com.intellij.psi.PsiFile
 import com.intellij.psi.codeStyle.CodeStyleManager
 import org.jetbrains.kotlin.idea.core.util.containsLineBreakInRange
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaState
 import org.jetbrains.kotlin.idea.formatter.trailingComma.addTrailingCommaIsAllowedForThis
+import org.jetbrains.kotlin.idea.formatter.trailingComma.existsOrMissing
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
@@ -27,7 +29,7 @@ class JoinWithTrailingCommaHandler : JoinLinesHandlerDelegate {
             ?.findLast { it.addTrailingCommaIsAllowedForThis() } as? KtElement
             ?: return CANNOT_JOIN
 
-        if (TrailingCommaHelper.needComma(commaOwner, null, true)) return CANNOT_JOIN
+        if (TrailingCommaState.stateForElement(commaOwner).existsOrMissing) return CANNOT_JOIN
         val result = CodeStyleManager.getInstance(file.project).reformat(commaOwner) as KtElement
         return TrailingCommaHelper.elementAfterLastElement(result)?.startOffset ?: end - 1
     }