From e4a9679ec2f937a4de76b07993dd818e9d88e305 Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Fri, 22 May 2020 19:17:42 +0700
Subject: [PATCH] TrailingCommaVisitor: improve tree traversal

---
 .../idea/formatter/TrailingCommaVisitor.kt    | 49 ++++---------------
 1 file changed, 10 insertions(+), 39 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/formatter/TrailingCommaVisitor.kt b/idea/src/org/jetbrains/kotlin/idea/formatter/TrailingCommaVisitor.kt
index 93f24146852..dbff2aada9d 100644
--- a/idea/src/org/jetbrains/kotlin/idea/formatter/TrailingCommaVisitor.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/formatter/TrailingCommaVisitor.kt
@@ -9,47 +9,18 @@ import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaContext
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaState
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.idea.formatter.trailingComma.canAddTrailingComma
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
 
 abstract class TrailingCommaVisitor : KtTreeVisitorVoid() {
-    override fun visitParameterList(list: KtParameterList) {
-        super.visitParameterList(list)
-        runProcessIfApplicable(list)
-    }
-
-    override fun visitValueArgumentList(list: KtValueArgumentList) {
-        super.visitValueArgumentList(list)
-        runProcessIfApplicable(list)
-    }
-
-    override fun visitArrayAccessExpression(expression: KtArrayAccessExpression) {
-        super.visitArrayAccessExpression(expression)
-        runProcessIfApplicable(expression.indicesNode)
-    }
-
-    override fun visitTypeParameterList(list: KtTypeParameterList) {
-        super.visitTypeParameterList(list)
-        runProcessIfApplicable(list)
-    }
-
-    override fun visitTypeArgumentList(typeArgumentList: KtTypeArgumentList) {
-        super.visitTypeArgumentList(typeArgumentList)
-        runProcessIfApplicable(typeArgumentList)
-    }
-
-    override fun visitCollectionLiteralExpression(expression: KtCollectionLiteralExpression) {
-        super.visitCollectionLiteralExpression(expression)
-        runProcessIfApplicable(expression)
-    }
-
-    override fun visitWhenEntry(jetWhenEntry: KtWhenEntry) {
-        super.visitWhenEntry(jetWhenEntry)
-        runProcessIfApplicable(jetWhenEntry)
-    }
-
-    override fun visitDestructuringDeclaration(destructuringDeclaration: KtDestructuringDeclaration) {
-        super.visitDestructuringDeclaration(destructuringDeclaration)
-        runProcessIfApplicable(destructuringDeclaration)
+    override fun visitKtElement(element: KtElement) {
+        super.visitKtElement(element)
+        // because KtFunctionLiteral contains KtParameterList
+        if (element !is KtFunctionLiteral && element.canAddTrailingComma()) {
+            runProcessIfApplicable(element)
+        }
     }
 
     override fun visitElement(element: PsiElement) {