Formatter: support trailing comma in when entry

#KT-34744

compiler/psi/src/org/jetbrains/kotlin/psi/KtWhenEntry.java

@@ -53,6 +53,11 @@ public class KtWhenEntry extends KtElementImpl {
     }
 
     public PsiElement getTrailingComma() {
-        return KtPsiUtilKt.getTrailingCommaByClosingElement(findChildByType(KtTokens.ARROW));
+        return KtPsiUtilKt.getTrailingCommaByClosingElement(getArrow());
+    }
+
+    @Nullable
+    public PsiElement getArrow() {
+        return findChildByType(KtTokens.ARROW);
     }
 }

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

@@ -304,7 +304,7 @@ abstract class KotlinCommonBlock(
         // do not indent child after heading comments inside declaration
         if (childParent != null && childParent.psi is KtDeclaration) {
             val prev = getPrevWithoutWhitespace(child)
-            if (prev != null && COMMENTS.contains(prev.elementType) && getPrevWithoutWhitespaceAndComments(prev) == null) {
+            if (prev != null && COMMENTS.contains(prev.elementType) && getSiblingWithoutWhitespaceAndComments(prev) == null) {
                 return Indent.getNoneIndent()
             }
         }
@@ -613,7 +613,7 @@ abstract class KotlinCommonBlock(
                         if (nodePsi.parent?.safeAs<KtFunctionLiteral>()?.needTrailingComma(settings) == true) {
                             val check = thisOrPrevIsMultiLineElement(COMMA, LBRACE /* not necessary */, ARROW /* not necessary */)
                             return { childElement ->
-                                createWrapAlwaysIf(getPrevWithoutWhitespaceAndComments(childElement) == null || check(childElement))
+                                createWrapAlwaysIf(getSiblingWithoutWhitespaceAndComments(childElement) == null || check(childElement))
                             }
                         }
                     }
@@ -621,11 +621,23 @@ abstract class KotlinCommonBlock(
             }
 
             elementType === FUNCTION_LITERAL -> {
-                val withTrailingComma = nodePsi.cast<KtFunctionLiteral>().needTrailingComma(settings)
+                if (nodePsi.cast<KtFunctionLiteral>().needTrailingComma(settings))
-                return { childElement ->
+                    return { childElement ->
-                    createWrapAlwaysIf(
+                        createWrapAlwaysIf(childElement.elementType === ARROW || getSiblingWithoutWhitespaceAndComments(childElement)?.elementType === LBRACE)
-                        withTrailingComma && (childElement.elementType === ARROW || getPrevWithoutWhitespaceAndComments(childElement)?.elementType === LBRACE)
+                    }
-                    )
+            }
+
+            elementType === WHEN_ENTRY -> {
+                // with argument
+                if (nodePsi.cast<KtWhenEntry>().needTrailingComma(settings)) {
+                    val check = thisOrPrevIsMultiLineElement(COMMA, LBRACE /* not necessary */, ARROW /* not necessary */)
+                    return { childElement ->
+                        createWrapAlwaysIf(
+                            childElement.elementType === ARROW ||
+                                    getSiblingWithoutWhitespaceAndComments(childElement, true) != null &&
+                                    check(childElement)
+                        )
+                    }
                 }
             }
 
@@ -686,7 +698,7 @@ abstract class KotlinCommonBlock(
                     getWrapAfterAnnotation(childElement, commonSettings.METHOD_ANNOTATION_WRAP)?.let {
                         return@wrap it
                     }
-                    if (getPrevWithoutWhitespaceAndComments(childElement)?.elementType == EQ) {
+                    if (getSiblingWithoutWhitespaceAndComments(childElement)?.elementType == EQ) {
                         return@wrap Wrap.createWrap(settings.kotlinCustomSettings.WRAP_EXPRESSION_BODY_FUNCTIONS, true)
                     }
                     null
@@ -698,7 +710,7 @@ abstract class KotlinCommonBlock(
                     getWrapAfterAnnotation(childElement, wrapSetting)?.let {
                         return@wrap it
                     }
-                    if (getPrevWithoutWhitespaceAndComments(childElement)?.elementType == EQ) {
+                    if (getSiblingWithoutWhitespaceAndComments(childElement)?.elementType == EQ) {
                         return@wrap Wrap.createWrap(settings.kotlinCommonSettings.ASSIGNMENT_WRAP, true)
                     }
                     null
@@ -707,7 +719,7 @@ abstract class KotlinCommonBlock(
             nodePsi is KtBinaryExpression -> {
                 if (nodePsi.operationToken == EQ) {
                     return { childElement ->
-                        if (getPrevWithoutWhitespaceAndComments(childElement)?.elementType == OPERATION_REFERENCE) {
+                        if (getSiblingWithoutWhitespaceAndComments(childElement)?.elementType == OPERATION_REFERENCE) {
                             Wrap.createWrap(settings.kotlinCommonSettings.ASSIGNMENT_WRAP, true)
                         } else {
                             null
@@ -737,7 +749,7 @@ abstract class KotlinCommonBlock(
 
     private val ASTNode.withTrailingComma: Boolean
         get() = when {
-            lastChildNode?.let { getPrevWithoutWhitespaceAndComments(it) }?.elementType === COMMA -> true
+            lastChildNode?.let { getSiblingWithoutWhitespaceAndComments(it) }?.elementType === COMMA -> true
             settings.kotlinCustomSettings.ALLOW_TRAILING_COMMA -> psi?.let(PsiElement::isMultiline) == true
             else -> false
         }
@@ -796,7 +808,7 @@ abstract class KotlinCommonBlock(
         val childElementType = childElement.elementType
         createWrapAlwaysIf(
             childElement.treeParent.withTrailingComma && (childElementType === rightAnchor ||
-                    getPrevWithoutWhitespaceAndComments(childElement)?.elementType === leftAnchor ||
+                    getSiblingWithoutWhitespaceAndComments(childElement)?.elementType === leftAnchor ||
                     additionalCheck(childElement)
                     )
         )
@@ -1110,8 +1122,8 @@ private fun getPrevWithoutWhitespace(pNode: ASTNode): ASTNode? {
     return pNode.siblings(forward = false).firstOrNull { it.elementType != TokenType.WHITE_SPACE }
 }
 
-private fun getPrevWithoutWhitespaceAndComments(pNode: ASTNode): ASTNode? {
+private fun getSiblingWithoutWhitespaceAndComments(pNode: ASTNode, forward: Boolean = false): ASTNode? {
-    return pNode.siblings(forward = false).firstOrNull {
+    return pNode.siblings(forward = forward).firstOrNull {
         it.elementType != TokenType.WHITE_SPACE && it.elementType !in COMMENTS
     }
 }

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

@@ -14,8 +14,11 @@ import com.intellij.psi.codeStyle.CodeStyleSettings
 import org.jetbrains.kotlin.idea.core.formatter.KotlinCodeStyleSettings
 import org.jetbrains.kotlin.idea.formatter.kotlinCustomSettings
 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.safeAs
 
 /*
  * ASTBlock.node is nullable, this extension was introduced to minimize changes
@@ -53,6 +56,15 @@ fun KtFunctionLiteral.needTrailingComma(settings: CodeStyleSettings): Boolean =
             return containsLineBreakInThis(startOffset, endOffset)
         })
 
+fun KtWhenEntry.needTrailingComma(settings: CodeStyleSettings): Boolean = trailingComma != null ||
+        !isElse &&
+        settings.kotlinCustomSettings.ALLOW_TRAILING_COMMA &&
+        parent.safeAs<KtWhenExpression>()?.leftParenthesis != null &&
+        run(fun(): Boolean {
+            val endOffset = arrow?.endOffset ?: return false
+            return containsLineBreakInThis(startOffset, endOffset)
+        })
+
 fun PsiElement.containsLineBreakInThis(globalStartOffset: Int, globalEndOffset: Int): Boolean {
     val textRange = TextRange.create(globalStartOffset, globalEndOffset).shiftLeft(startOffset)
     return StringUtil.containsLineBreak(textRange.subSequence(text))

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
 import org.jetbrains.kotlin.psi.psiUtil.getPrevSiblingIgnoringWhitespace
 import org.jetbrains.kotlin.psi.psiUtil.getPrevSiblingIgnoringWhitespaceAndComments
 import org.jetbrains.kotlin.psi.psiUtil.siblings
+import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class TrailingCommaPostFormatProcessor : PostFormatProcessor {
@@ -59,10 +60,18 @@ private class TrailingCommaVisitor(val settings: CodeStyleSettings) : KtTreeVisi
         super.visitCollectionLiteralExpression(expression)
     }
 
-    private fun processCommaOwnerIfInRange(element: KtElement, preHook: () -> Unit = {}) {
+    override fun visitWhenEntry(jetWhenEntry: KtWhenEntry) = processCommaOwnerIfInRange(jetWhenEntry) {
+        super.visitWhenEntry(jetWhenEntry)
+    }
+
+    private fun processCommaOwnerIfInRange(element: KtElement, preHook: () -> Unit = {}) = processIfInRange(element) {
+        preHook()
+        processCommaOwner(element)
+    }
+
+    private fun processIfInRange(element: KtElement, block: () -> Unit = {}) {
         if (myPostProcessor.isElementPartlyInRange(element)) {
-            preHook()
+            block()
-            processCommaOwner(element)
         }
     }
 
@@ -82,8 +91,9 @@ private class TrailingCommaVisitor(val settings: CodeStyleSettings) : KtTreeVisi
     }
 
     private val KtElement.needComma: Boolean
-        get() = when (val parent = parent) {
+        get() = when {
-            is KtFunctionLiteral -> parent.needTrailingComma(settings)
+            this is KtWhenEntry -> needTrailingComma(settings)
+            parent is KtFunctionLiteral -> parent.cast<KtFunctionLiteral>().needTrailingComma(settings)
             else -> settings.kotlinCustomSettings.ALLOW_TRAILING_COMMA && isMultiline()
         }
 
@@ -138,12 +148,18 @@ private class TrailingCommaVisitor(val settings: CodeStyleSettings) : KtTreeVisi
 
 private val PsiElement.lastCommaOwnerOrComma: PsiElement?
     get() {
-        val lastChild = lastChild ?: return null
+        val lastChild = lastSignificantChild ?: return null
         val withSelf = when (lastChild.safeAs<ASTNode>()?.elementType) {
             KtTokens.COMMA -> return lastChild
-            KtTokens.RBRACKET, KtTokens.RPAR, KtTokens.RBRACE, KtTokens.GT -> false
+            KtTokens.RBRACKET, KtTokens.RPAR, KtTokens.RBRACE, KtTokens.GT, KtTokens.ARROW -> false
             else -> true
         }
 
         return lastChild.getPrevSiblingIgnoringWhitespaceAndComments(withSelf)
     }
+
+private val PsiElement.lastSignificantChild: PsiElement?
+    get() = when (this) {
+        is KtWhenEntry -> arrow
+        else -> lastChild
+    }

idea/testData/formatter/WhenEntryExpr.after.kt

@@ -39,16 +39,16 @@ fun some(x: Any) {
     }
     when (x) {
         is
-        Int
+        Int,
         -> {
             0
         }
-        3
+        3,
         -> {
             2
         }
         in
-        0..3
+        0..3,
         -> {
             2
         }

idea/testData/formatter/trailingComma/whenEntry/WhenEntry.after.inv.kt

@@ -0,0 +1,67 @@
+fun foo(x: Any) = when (x) {
+    Comparable::class, Iterable::class, String::class, // trailing comma
+    -> println(1)
+    else -> println(3)
+}
+
+fun foo(x: Any) {
+    when (x) {
+        Comparable::class, Iterable::class, String::class, /*// trailing comma*/
+        -> println(1)
+        else -> println(3)
+    }
+
+    when (x) {
+        Comparable::class, Iterable::class, String::class /*// trailing comma*/ -> println(1)
+        else -> println(3)
+    }
+
+    when (x) {
+        1 -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1,
+        -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1
+        -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1, 2,
+        3 /**/ -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (val c = x) {
+        1, 2,
+        3 /**/ -> {
+
+        }
+        else -> println(3)
+    }
+
+    when {
+        x in coll
+        -> {
+
+        }
+        else -> println(3)
+    }
+}
+
+// SET_TRUE: ALLOW_TRAILING_COMMA

idea/testData/formatter/trailingComma/whenEntry/WhenEntry.after.kt

@@ -0,0 +1,69 @@
+fun foo(x: Any) = when (x) {
+    Comparable::class, Iterable::class, String::class, // trailing comma
+    -> println(1)
+    else -> println(3)
+}
+
+fun foo(x: Any) {
+    when (x) {
+        Comparable::class, Iterable::class, String::class, /*// trailing comma*/
+        -> println(1)
+        else -> println(3)
+    }
+
+    when (x) {
+        Comparable::class, Iterable::class, String::class /*// trailing comma*/ -> println(1)
+        else -> println(3)
+    }
+
+    when (x) {
+        1 -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1,
+        -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1,
+        -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1, 2,
+        3, /**/
+        -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (val c = x) {
+        1, 2,
+        3, /**/
+        -> {
+
+        }
+        else -> println(3)
+    }
+
+    when {
+        x in coll,
+        -> {
+
+        }
+        else -> println(3)
+    }
+}
+
+// SET_TRUE: ALLOW_TRAILING_COMMA

idea/testData/formatter/trailingComma/whenEntry/WhenEntry.kt

@@ -0,0 +1,65 @@
+fun foo(x: Any) = when (x) {
+    Comparable::class, Iterable::class, String::class, // trailing comma
+    -> println(1)
+    else -> println(3)
+}
+
+fun foo(x: Any) {
+    when (x) {
+        Comparable::class, Iterable::class, String::class, /*// trailing comma*/ -> println(1)
+        else -> println(3)
+    }
+
+    when (x) {
+        Comparable::class, Iterable::class, String::class /*// trailing comma*/ -> println(1)
+        else -> println(3)
+    }
+
+    when (x) {
+        1 -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1, -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1
+        -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (x) {
+        1, 2,
+        3 /**/ -> {
+
+        }
+        else -> println(3)
+    }
+
+    when (val c = x) {
+        1, 2,
+        3 /**/ -> {
+
+        }
+        else -> println(3)
+    }
+
+    when {
+        x in coll
+        -> {
+
+        }
+        else -> println(3)
+    }
+}
+
+// SET_TRUE: ALLOW_TRAILING_COMMA

idea/tests/org/jetbrains/kotlin/formatter/FormatterTestGenerated.java

@@ -1319,6 +1319,24 @@ public class FormatterTestGenerated extends AbstractFormatterTest {
                     runTest("idea/testData/formatter/trailingComma/valueParameters/ParameterListWrapAsNeeded.after.kt");
                 }
             }
+
+            @TestMetadata("idea/testData/formatter/trailingComma/whenEntry")
+            @TestDataPath("$PROJECT_ROOT")
+            @RunWith(JUnit3RunnerWithInners.class)
+            public static class WhenEntry extends AbstractFormatterTest {
+                private void runTest(String testDataFilePath) throws Exception {
+                    KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                }
+
+                public void testAllFilesPresentInWhenEntry() throws Exception {
+                    KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/formatter/trailingComma/whenEntry"), Pattern.compile("^([^\\.]+)\\.after\\.kt.*$"), null, true);
+                }
+
+                @TestMetadata("WhenEntry.after.kt")
+                public void testWhenEntry() throws Exception {
+                    runTest("idea/testData/formatter/trailingComma/whenEntry/WhenEntry.after.kt");
+                }
+            }
         }
     }
 
@@ -1853,6 +1871,24 @@ public class FormatterTestGenerated extends AbstractFormatterTest {
                     runTest("idea/testData/formatter/trailingComma/valueParameters/ParameterListWrapAsNeeded.after.inv.kt");
                 }
             }
+
+            @TestMetadata("idea/testData/formatter/trailingComma/whenEntry")
+            @TestDataPath("$PROJECT_ROOT")
+            @RunWith(JUnit3RunnerWithInners.class)
+            public static class WhenEntry extends AbstractFormatterTest {
+                private void runTest(String testDataFilePath) throws Exception {
+                    KotlinTestUtils.runTest(this::doTestInverted, this, testDataFilePath);
+                }
+
+                public void testAllFilesPresentInWhenEntry() throws Exception {
+                    KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/formatter/trailingComma/whenEntry"), Pattern.compile("^([^\\.]+)\\.after\\.inv\\.kt.*$"), null, true);
+                }
+
+                @TestMetadata("WhenEntry.after.inv.kt")
+                public void testWhenEntry() throws Exception {
+                    runTest("idea/testData/formatter/trailingComma/whenEntry/WhenEntry.after.inv.kt");
+                }
+            }
         }
     }
 }