From fbddc6921f9968584b6224952bd09b974a16ea21 Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Thu, 21 May 2020 15:29:19 +0700
Subject: [PATCH] KotlinCodeStyleSettings: add hidden settings for trailing
 comma on call site

#KT-39079
---
 .../formatter/KotlinCodeStyleSettings.java    |  1 +
 .../kotlin/idea/util/FormatterUtil.kt         | 30 ++++++++++++++-----
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/idea/formatter/src/org/jetbrains/kotlin/idea/core/formatter/KotlinCodeStyleSettings.java b/idea/formatter/src/org/jetbrains/kotlin/idea/core/formatter/KotlinCodeStyleSettings.java
index 4fe43b83da3..b7bcf77bbe1 100644
--- a/idea/formatter/src/org/jetbrains/kotlin/idea/core/formatter/KotlinCodeStyleSettings.java
+++ b/idea/formatter/src/org/jetbrains/kotlin/idea/core/formatter/KotlinCodeStyleSettings.java
@@ -53,6 +53,7 @@ public class KotlinCodeStyleSettings extends CustomCodeStyleSettings {
     public int WRAP_ELVIS_EXPRESSIONS = 1;
     public boolean IF_RPAREN_ON_NEW_LINE = false;
     public boolean ALLOW_TRAILING_COMMA = false;
+    public boolean ALLOW_TRAILING_COMMA_ON_CALL_SITE = false;
     public int BLANK_LINES_BEFORE_DECLARATION_WITH_COMMENT_OR_ANNOTATION_ON_SEPARATE_LINE = 1;
 
     @ReflectionUtil.SkipInEquals
diff --git a/idea/formatter/src/org/jetbrains/kotlin/idea/util/FormatterUtil.kt b/idea/formatter/src/org/jetbrains/kotlin/idea/util/FormatterUtil.kt
index ab6407f8811..7c9aa2ae002 100644
--- a/idea/formatter/src/org/jetbrains/kotlin/idea/util/FormatterUtil.kt
+++ b/idea/formatter/src/org/jetbrains/kotlin/idea/util/FormatterUtil.kt
@@ -14,6 +14,7 @@ import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.PsiDocumentManager
 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
@@ -100,7 +101,7 @@ fun PsiElement.containsLineBreakInThis(globalStartOffset: Int, globalEndOffset:
 
 fun trailingCommaIsAllowedOnCallSite(): Boolean = Registry.`is`("kotlin.formatter.allowTrailingCommaOnCallSite")
 
-private val TYPES_WITH_TRAILING_COMMA = TokenSet.create(
+private val TYPES_WITH_TRAILING_COMMA_ON_DECLARATION_SITE = TokenSet.create(
     KtNodeTypes.TYPE_PARAMETER_LIST,
     KtNodeTypes.DESTRUCTURING_DECLARATION,
     KtNodeTypes.WHEN_ENTRY,
@@ -115,16 +116,29 @@ private val TYPES_WITH_TRAILING_COMMA_ON_CALL_SITE = TokenSet.create(
     KtNodeTypes.VALUE_ARGUMENT_LIST,
 )
 
-fun UserDataHolder.addTrailingCommaIsAllowedForThis(): Boolean {
-    val type = when (this) {
-        is ASTNode -> PsiUtilCore.getElementType(this)
-        is PsiElement -> PsiUtilCore.getElementType(this)
-        else -> return false
-    }
+private val TYPES_WITH_TRAILING_COMMA = TokenSet.orSet(
+    TYPES_WITH_TRAILING_COMMA_ON_DECLARATION_SITE,
+    TYPES_WITH_TRAILING_COMMA_ON_CALL_SITE,
+)
 
-    return type in TYPES_WITH_TRAILING_COMMA || trailingCommaIsAllowedOnCallSite() && type in TYPES_WITH_TRAILING_COMMA_ON_CALL_SITE
+fun UserDataHolder.addTrailingCommaIsAllowedForThis(): Boolean {
+    val type = elementType(this) ?: return false
+    return type in TYPES_WITH_TRAILING_COMMA_ON_DECLARATION_SITE ||
+            trailingCommaIsAllowedOnCallSite() && type in TYPES_WITH_TRAILING_COMMA_ON_CALL_SITE
 }
 
 fun KotlinCodeStyleSettings.addTrailingCommaIsAllowedFor(element: UserDataHolder): Boolean =
     ALLOW_TRAILING_COMMA && element.addTrailingCommaIsAllowedForThis()
 
+private fun elementType(userDataHolder: UserDataHolder): IElementType? = when (userDataHolder) {
+    is ASTNode -> PsiUtilCore.getElementType(userDataHolder)
+    is PsiElement -> PsiUtilCore.getElementType(userDataHolder)
+    else -> null
+}
+
+fun PsiElement.canAddTrailingComma(): Boolean = canAddTrailingComma(this)
+
+fun ASTNode.canAddTrailingComma(): Boolean = canAddTrailingComma(this)
+
+private fun canAddTrailingComma(userDataHolder: UserDataHolder): Boolean = elementType(userDataHolder) in TYPES_WITH_TRAILING_COMMA
+