TrailingCommaVisitor: adapt to TrailingCommaContext

#KT-39123 Fixed

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

@@ -7,8 +7,14 @@ package org.jetbrains.kotlin.idea.formatter.trailingComma
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.core.formatter.KotlinCodeStyleSettings
+import org.jetbrains.kotlin.psi.KtElement
 
 class TrailingCommaContext private constructor(val element: PsiElement, val state: TrailingCommaState) {
+    /**
+     * Return [KtElement] if [state] != [TrailingCommaState.NOT_APPLICABLE]
+     */
+    val ktElement: KtElement get() = element as? KtElement ?: error("State is NOT_APPLICABLE")
+
     companion object {
         fun create(element: PsiElement): TrailingCommaContext = TrailingCommaContext(
             element,

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

@@ -16,10 +16,12 @@ import com.intellij.psi.codeStyle.CodeStyleManager
 import com.intellij.psi.codeStyle.CodeStyleSettings
 import com.intellij.psi.impl.source.codeStyle.PostFormatProcessor
 import com.intellij.psi.impl.source.codeStyle.PostFormatProcessorHelper
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaContext
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper.findInvalidCommas
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper.trailingCommaOrLastElement
 import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaState
 import org.jetbrains.kotlin.idea.formatter.trailingComma.addTrailingCommaIsAllowedFor
+import org.jetbrains.kotlin.idea.formatter.trailingComma.addTrailingCommaIsAllowedForThis
 import org.jetbrains.kotlin.idea.util.leafIgnoringWhitespace
 import org.jetbrains.kotlin.idea.util.leafIgnoringWhitespaceAndComments
 import org.jetbrains.kotlin.psi.KtElement
@@ -38,8 +40,12 @@ class TrailingCommaPostFormatProcessor : PostFormatProcessor {
 private class TrailingCommaPostFormatVisitor(val settings: CodeStyleSettings) : TrailingCommaVisitor() {
     private val myPostProcessor = PostFormatProcessorHelper(settings.kotlinCommonSettings)
 
-    override fun process(commaOwner: KtElement) = processIfInRange(commaOwner) {
+    override fun process(trailingCommaContext: TrailingCommaContext) {
-        processCommaOwner(commaOwner)
+        val ktElement = trailingCommaContext.ktElement
+        if (!ktElement.addTrailingCommaIsAllowedForThis()) return
+        processIfInRange(ktElement) {
+            processCommaOwner(trailingCommaContext)
+        }
     }
 
     private fun processIfInRange(element: KtElement, block: () -> Unit = {}) {
@@ -48,25 +54,26 @@ private class TrailingCommaPostFormatVisitor(val settings: CodeStyleSettings) :
         }
     }
 
-    private fun processCommaOwner(parent: KtElement) {
+    private fun processCommaOwner(trailingCommaContext: TrailingCommaContext) {
-        if (!postFormatIsEnable(parent)) return
+        val ktElement = trailingCommaContext.ktElement
+        if (!postFormatIsEnable(ktElement)) return
 
-        val lastElementOrComma = trailingCommaOrLastElement(parent) ?: return
+        val lastElementOrComma = trailingCommaOrLastElement(ktElement) ?: return
-        updatePsi(parent) {
+        updatePsi(ktElement) {
-            val stateForElement = TrailingCommaState.stateForElement(parent)
+            val state = trailingCommaContext.state
             when {
-                stateForElement == TrailingCommaState.MISSING && settings.kotlinCustomSettings.addTrailingCommaIsAllowedFor(parent) -> {
+                state == TrailingCommaState.MISSING && settings.kotlinCustomSettings.addTrailingCommaIsAllowedFor(ktElement) -> {
                     // add a missing comma
-                    if (trailingCommaAllowedInModule(parent)) {
+                    if (trailingCommaAllowedInModule(ktElement)) {
-                        lastElementOrComma.addCommaAfter(KtPsiFactory(parent))
+                        lastElementOrComma.addCommaAfter(KtPsiFactory(ktElement))
                     }
 
-                    correctCommaPosition(parent)
+                    correctCommaPosition(ktElement)
                 }
-                stateForElement == TrailingCommaState.EXISTS -> {
+                state == TrailingCommaState.EXISTS -> {
-                    correctCommaPosition(parent)
+                    correctCommaPosition(ktElement)
                 }
-                stateForElement == TrailingCommaState.REDUNDANT -> {
+                state == TrailingCommaState.REDUNDANT -> {
                     // remove redundant comma
                     lastElementOrComma.delete()
                 }

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

@@ -7,48 +7,49 @@ package org.jetbrains.kotlin.idea.formatter
 
 import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.idea.formatter.trailingComma.addTrailingCommaIsAllowedForThis
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaContext
+import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaState
 import org.jetbrains.kotlin.psi.*
 
 abstract class TrailingCommaVisitor : KtTreeVisitorVoid() {
     override fun visitParameterList(list: KtParameterList) {
         super.visitParameterList(list)
-        runProcessIfAllowed(list)
+        runProcessIfApplicable(list)
     }
 
     override fun visitValueArgumentList(list: KtValueArgumentList) {
         super.visitValueArgumentList(list)
-        runProcessIfAllowed(list)
+        runProcessIfApplicable(list)
     }
 
     override fun visitArrayAccessExpression(expression: KtArrayAccessExpression) {
         super.visitArrayAccessExpression(expression)
-        runProcessIfAllowed(expression.indicesNode)
+        runProcessIfApplicable(expression.indicesNode)
     }
 
     override fun visitTypeParameterList(list: KtTypeParameterList) {
         super.visitTypeParameterList(list)
-        runProcessIfAllowed(list)
+        runProcessIfApplicable(list)
     }
 
     override fun visitTypeArgumentList(typeArgumentList: KtTypeArgumentList) {
         super.visitTypeArgumentList(typeArgumentList)
-        runProcessIfAllowed(typeArgumentList)
+        runProcessIfApplicable(typeArgumentList)
     }
 
     override fun visitCollectionLiteralExpression(expression: KtCollectionLiteralExpression) {
         super.visitCollectionLiteralExpression(expression)
-        runProcessIfAllowed(expression)
+        runProcessIfApplicable(expression)
     }
 
     override fun visitWhenEntry(jetWhenEntry: KtWhenEntry) {
         super.visitWhenEntry(jetWhenEntry)
-        runProcessIfAllowed(jetWhenEntry)
+        runProcessIfApplicable(jetWhenEntry)
     }
 
     override fun visitDestructuringDeclaration(destructuringDeclaration: KtDestructuringDeclaration) {
         super.visitDestructuringDeclaration(destructuringDeclaration)
-        runProcessIfAllowed(destructuringDeclaration)
+        runProcessIfApplicable(destructuringDeclaration)
     }
 
     override fun visitElement(element: PsiElement) {
@@ -57,13 +58,17 @@ abstract class TrailingCommaVisitor : KtTreeVisitorVoid() {
         if (recursively) super.visitElement(element)
     }
 
-    private fun runProcessIfAllowed(element: KtElement) {
+    private fun runProcessIfApplicable(element: KtElement) {
-        if (element.addTrailingCommaIsAllowedForThis()) {
+        val context = TrailingCommaContext.create(element)
-            process(element)
+        if (context.state != TrailingCommaState.NOT_APPLICABLE) {
+            process(context)
         }
     }
 
-    protected abstract fun process(commaOwner: KtElement)
+    /**
+     * [trailingCommaContext] doesn't contain a state [TrailingCommaState.NOT_APPLICABLE]
+     */
+    protected abstract fun process(trailingCommaContext: TrailingCommaContext)
 
     protected open val recursively: Boolean = true
 }

idea/src/org/jetbrains/kotlin/idea/inspections/TrailingCommaInspection.kt

@@ -20,8 +20,7 @@ import com.intellij.psi.codeStyle.CodeStyleManager
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.formatter.TrailingCommaVisitor
 import org.jetbrains.kotlin.idea.formatter.kotlinCustomSettings
-import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaHelper
+import org.jetbrains.kotlin.idea.formatter.trailingComma.*
-import org.jetbrains.kotlin.idea.formatter.trailingComma.TrailingCommaState
 import org.jetbrains.kotlin.idea.formatter.trailingCommaAllowedInModule
 import org.jetbrains.kotlin.idea.util.isComma
 import org.jetbrains.kotlin.idea.util.isLineBreak
@@ -39,18 +38,20 @@ class TrailingCommaInspection(
         override val recursively: Boolean = false
         private var useTrailingComma by Delegates.notNull<Boolean>()
 
-        override fun process(commaOwner: KtElement) {
+        override fun process(trailingCommaContext: TrailingCommaContext) {
-            useTrailingComma = CodeStyle.getSettings(commaOwner.project).kotlinCustomSettings.ALLOW_TRAILING_COMMA
+            val element = trailingCommaContext.ktElement
-            val state = TrailingCommaState.stateForElement(commaOwner)
+            if (!element.addTrailingCommaIsAllowedForThis()) return
-            when (state) {
+
+            useTrailingComma = CodeStyle.getSettings(element.project).kotlinCustomSettings.ALLOW_TRAILING_COMMA
+            when (trailingCommaContext.state) {
                 TrailingCommaState.MISSING, TrailingCommaState.EXISTS -> {
-                    checkCommaPosition(commaOwner)
+                    checkCommaPosition(element)
-                    checkLineBreaks(commaOwner)
+                    checkLineBreaks(element)
                 }
                 else -> Unit
             }
 
-            checkTrailingComma(commaOwner, state)
+            checkTrailingComma(trailingCommaContext)
         }
 
         private fun checkLineBreaks(commaOwner: KtElement) {
@@ -81,24 +82,29 @@ class TrailingCommaInspection(
             }
         }
 
-        private fun checkTrailingComma(commaOwner: KtElement, state: TrailingCommaState) {
+        private fun checkTrailingComma(trailingCommaContext: TrailingCommaContext) {
+            val commaOwner = trailingCommaContext.ktElement
             val trailingCommaOrLastElement = TrailingCommaHelper.trailingCommaOrLastElement(commaOwner) ?: return
-            if (state == TrailingCommaState.MISSING) {
+            when (trailingCommaContext.state) {
-                if (!trailingCommaAllowedInModule(commaOwner)) return
+                TrailingCommaState.MISSING -> {
-                reportProblem(
+                    if (!trailingCommaAllowedInModule(commaOwner)) return
-                    trailingCommaOrLastElement,
+                    reportProblem(
-                    KotlinBundle.message("inspection.trailing.comma.missing.trailing.comma"),
+                        trailingCommaOrLastElement,
-                    KotlinBundle.message("inspection.trailing.comma.add.trailing.comma"),
+                        KotlinBundle.message("inspection.trailing.comma.missing.trailing.comma"),
-                    if (addCommaWarning) ProblemHighlightType.GENERIC_ERROR_OR_WARNING else ProblemHighlightType.INFORMATION,
+                        KotlinBundle.message("inspection.trailing.comma.add.trailing.comma"),
-                )
+                        if (addCommaWarning) ProblemHighlightType.GENERIC_ERROR_OR_WARNING else ProblemHighlightType.INFORMATION,
-            } else if (state == TrailingCommaState.REDUNDANT) {
+                    )
-                reportProblem(
+                }
-                    trailingCommaOrLastElement,
+                TrailingCommaState.REDUNDANT -> {
-                    KotlinBundle.message("inspection.trailing.comma.useless.trailing.comma"),
+                    reportProblem(
-                    KotlinBundle.message("inspection.trailing.comma.remove.trailing.comma"),
+                        trailingCommaOrLastElement,
-                    ProblemHighlightType.LIKE_UNUSED_SYMBOL,
+                        KotlinBundle.message("inspection.trailing.comma.useless.trailing.comma"),
-                    checkTrailingCommaSettings = false,
+                        KotlinBundle.message("inspection.trailing.comma.remove.trailing.comma"),
-                )
+                        ProblemHighlightType.LIKE_UNUSED_SYMBOL,
+                        checkTrailingCommaSettings = false,
+                    )
+                }
+                else -> Unit
             }
         }
 

idea/testData/formatter/When.after.inv.kt

@@ -1,7 +1,7 @@
 fun some(x: Any) {
     when (x) {
         is Number -> 0
-        else -> 1
+        else      -> 1
     }
     when {
     }
@@ -37,4 +37,11 @@ fun some(x: Any) {
     }
 }
 
+val a = when {
+    true && true && true -> {
+    }
+    else                 -> {
+    }
+}
+
 // SET_FALSE: ALIGN_IN_COLUMNS_CASE_BRANCH

idea/testData/formatter/When.after.kt

@@ -37,4 +37,11 @@ fun some(x: Any) {
     }
 }
 
+val a = when {
+    true && true && true -> {
+    }
+    else -> {
+    }
+}
+
 // SET_FALSE: ALIGN_IN_COLUMNS_CASE_BRANCH

idea/testData/formatter/When.kt

@@ -43,4 +43,11 @@ else->1
     when {}
 }
 
+val a = when {
+    true && true && true -> {
+    }
+    else -> {
+    }
+}
+
 // SET_FALSE: ALIGN_IN_COLUMNS_CASE_BRANCH