TrailingCommaPostFormatProcessor: introduce TrailingCommaVisitor

#KT-34744

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

@@ -1,10 +1,11 @@
 /*
- * Copyright 2010-2019 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.
  */
 
 package org.jetbrains.kotlin.idea.formatter
 
+import com.intellij.application.options.CodeStyle
 import com.intellij.lang.ASTNode
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.util.TextRange
@@ -22,6 +23,7 @@ import com.intellij.psi.tree.TokenSet
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaPostFormatProcessor.Companion.findInvalidCommas
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaPostFormatProcessor.Companion.needComma
+import org.jetbrains.kotlin.idea.formatter.TrailingCommaPostFormatProcessor.Companion.trailingCommaAllowedInModule
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaPostFormatProcessor.Companion.trailingCommaOrLastElement
 import org.jetbrains.kotlin.idea.project.languageVersionSettings
 import org.jetbrains.kotlin.idea.util.isMultiline
@@ -35,10 +37,10 @@ import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class TrailingCommaPostFormatProcessor : PostFormatProcessor {
     override fun processElement(source: PsiElement, settings: CodeStyleSettings): PsiElement =
-        TrailingCommaVisitor(settings).process(source)
+        TrailingCommaPostFormatVisitor(settings).process(source)
 
     override fun processText(source: PsiFile, rangeToReformat: TextRange, settings: CodeStyleSettings): TextRange =
-        TrailingCommaVisitor(settings).processText(source, rangeToReformat)
+        TrailingCommaPostFormatVisitor(settings).processText(source, rangeToReformat)
 
     companion object {
         fun findInvalidCommas(commaOwner: KtElement): List<PsiElement> = commaOwner.firstChild?.siblings(withItself = false)?.mapNotNull {
@@ -48,7 +50,11 @@ class TrailingCommaPostFormatProcessor : PostFormatProcessor {
                 null
         }?.toList().orEmpty()
 
-        fun needComma(commaOwner: KtElement, settings: CodeStyleSettings, checkExistingTrailingComma: Boolean = true): Boolean = when {
+        fun needComma(
+            commaOwner: KtElement,
+            settings: CodeStyleSettings = CodeStyle.getSettings(commaOwner.project),
+            checkExistingTrailingComma: Boolean = true
+        ): Boolean = when {
             commaOwner is KtWhenEntry ->
                 commaOwner.needTrailingComma(settings, checkExistingTrailingComma)
 
@@ -75,54 +81,18 @@ class TrailingCommaPostFormatProcessor : PostFormatProcessor {
                 it.safeAs<ASTNode>()?.elementType !in LEFT_BARRIERS
             }
         }
+
+        fun trailingCommaAllowedInModule(source: PsiElement): Boolean =
+            Registry.`is`("kotlin.formatter.allowTrailingCommaInAnyProject", false) ||
+                    source.module?.languageVersionSettings?.supportsFeature(LanguageFeature.TrailingCommas) == true
     }
 }
 
-private fun trailingCommaAllowedInModule(source: PsiElement): Boolean =
+private class TrailingCommaPostFormatVisitor(val settings: CodeStyleSettings) : TrailingCommaVisitor() {
-    Registry.`is`("kotlin.formatter.allowTrailingCommaInAnyProject", false) ||
-            source.module?.languageVersionSettings?.supportsFeature(LanguageFeature.TrailingCommas) == true
-
-private fun postFormatIsEnable(source: PsiElement): Boolean = !PostprocessReformattingAspect.getInstance(source.project).isDisabled
-
-private class TrailingCommaVisitor(val settings: CodeStyleSettings) : KtTreeVisitorVoid() {
     private val myPostProcessor = PostFormatProcessorHelper(settings.kotlinCommonSettings)
 
-    override fun visitParameterList(list: KtParameterList) = processCommaOwnerIfInRange(list) {
+    override fun process(commaOwner: KtElement) = processIfInRange(commaOwner) {
-        super.visitParameterList(list)
+        processCommaOwner(commaOwner)
-    }
-
-    override fun visitValueArgumentList(list: KtValueArgumentList) = processCommaOwnerIfInRange(list) {
-        super.visitValueArgumentList(list)
-    }
-
-    override fun visitArrayAccessExpression(expression: KtArrayAccessExpression) = processCommaOwnerIfInRange(expression.indicesNode) {
-        super.visitArrayAccessExpression(expression)
-    }
-
-    override fun visitTypeParameterList(list: KtTypeParameterList) = processCommaOwnerIfInRange(list) {
-        super.visitTypeParameterList(list)
-    }
-
-    override fun visitTypeArgumentList(typeArgumentList: KtTypeArgumentList) = processCommaOwnerIfInRange(typeArgumentList) {
-        super.visitTypeArgumentList(typeArgumentList)
-    }
-
-    override fun visitCollectionLiteralExpression(expression: KtCollectionLiteralExpression) = processCommaOwnerIfInRange(expression) {
-        super.visitCollectionLiteralExpression(expression)
-    }
-
-    override fun visitWhenEntry(jetWhenEntry: KtWhenEntry) = processCommaOwnerIfInRange(jetWhenEntry) {
-        super.visitWhenEntry(jetWhenEntry)
-    }
-
-    override fun visitDestructuringDeclaration(destructuringDeclaration: KtDestructuringDeclaration) =
-        processCommaOwnerIfInRange(destructuringDeclaration) {
-            super.visitDestructuringDeclaration(destructuringDeclaration)
-        }
-
-    private fun processCommaOwnerIfInRange(element: KtElement, preHook: () -> Unit = {}) = processIfInRange(element) {
-        preHook()
-        processCommaOwner(element)
     }
 
     private fun processIfInRange(element: KtElement, block: () -> Unit = {}) {
@@ -194,6 +164,9 @@ private class TrailingCommaVisitor(val settings: CodeStyleSettings) : KtTreeVisi
     }
 }
 
+
+private fun postFormatIsEnable(source: PsiElement): Boolean = !PostprocessReformattingAspect.getInstance(source.project).isDisabled
+
 private fun PsiElement.addCommaAfter(factory: KtPsiFactory) {
     val comma = factory.createComma()
     parent.addAfter(comma, this)
@@ -220,8 +193,11 @@ private val PsiElement.lastSignificantChild: PsiElement?
         else -> lastChild
     }
 
-fun PsiElement.leafIgnoringWhitespace(forward: Boolean = true) = leaf(forward) { it !is PsiWhiteSpace }
+fun PsiElement.leafIgnoringWhitespace(forward: Boolean = true, skipEmptyElements: Boolean = true) =
-fun PsiElement.leafIgnoringWhitespaceAndComments(forward: Boolean = true) = leaf(forward) { it !is PsiWhiteSpace && it !is PsiComment }
+    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)

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

@@ -0,0 +1,63 @@
+/*
+ * 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.openapi.progress.ProgressIndicatorProvider
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.*
+
+abstract class TrailingCommaVisitor : KtTreeVisitorVoid() {
+    override fun visitParameterList(list: KtParameterList) {
+        super.visitParameterList(list)
+        process(list)
+    }
+
+    override fun visitValueArgumentList(list: KtValueArgumentList) {
+        super.visitValueArgumentList(list)
+        process(list)
+    }
+
+    override fun visitArrayAccessExpression(expression: KtArrayAccessExpression) {
+        super.visitArrayAccessExpression(expression)
+        process(expression.indicesNode)
+    }
+
+    override fun visitTypeParameterList(list: KtTypeParameterList) {
+        super.visitTypeParameterList(list)
+        process(list)
+    }
+
+    override fun visitTypeArgumentList(typeArgumentList: KtTypeArgumentList) {
+        super.visitTypeArgumentList(typeArgumentList)
+        process(typeArgumentList)
+    }
+
+    override fun visitCollectionLiteralExpression(expression: KtCollectionLiteralExpression) {
+        super.visitCollectionLiteralExpression(expression)
+        process(expression)
+    }
+
+    override fun visitWhenEntry(jetWhenEntry: KtWhenEntry) {
+        super.visitWhenEntry(jetWhenEntry)
+        process(jetWhenEntry)
+    }
+
+    override fun visitDestructuringDeclaration(destructuringDeclaration: KtDestructuringDeclaration) {
+        super.visitDestructuringDeclaration(destructuringDeclaration)
+        process(destructuringDeclaration)
+    }
+
+    override fun visitElement(element: PsiElement) {
+        if (recursively)
+            super.visitElement(element)
+        else
+            ProgressIndicatorProvider.checkCanceled()
+    }
+
+    protected abstract fun process(commaOwner: KtElement)
+
+    protected open val recursively: Boolean = true
+}