TrailingCommaInspection: should report Missing trailing comma on symbol after argument

#KT-34744

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.inspections
 
+import com.intellij.application.options.CodeStyle
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.codeInspection.ui.MultipleCheckboxOptionsPanel
@@ -73,23 +74,26 @@ class TrailingCommaInspection(
             highlightType: ProblemHighlightType = ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
         ) {
             val commaOwner = commaOrElement.parent
+            // case for KtFunctionLiteral, where PsiWhiteSpace after KtTypeParameterList isn't included in this list
+            val problemOwner = commaOwner.parent
             holder.registerProblem(
-                commaOwner,
+                problemOwner,
                 message,
                 highlightType,
-                commaOrElement.textRangeOfLastSymbol.shiftLeft(commaOwner.startOffset),
+                commaOrElement.textRangeOfCommaOrSymbolAfter.shiftLeft(problemOwner.startOffset),
-                ReformatQuickFix(fixMessage),
+                ReformatQuickFix(fixMessage, commaOwner),
             )
         }
 
-        private val PsiElement.textRangeOfLastSymbol: TextRange
+        private val PsiElement.textRangeOfCommaOrSymbolAfter: TextRange
             get() {
                 val textRange = textRange
                 if (textRange.length <= 1) return textRange
 
-                return nextLeaf()?.leafIgnoringWhitespaceAndComments(false)?.endOffset?.takeIf { it > 0 }?.let {
+                val resultRange = nextLeaf()?.leafIgnoringWhitespaceAndComments(false)?.endOffset?.takeIf { it > 0 }?.let {
-                    TextRange.create(it - 1, it).intersection(textRange)
+                    TextRange.create(it - 1, it).intersection(parent.textRange)
                 } ?: TextRange.create(textRange.endOffset - 1, textRange.endOffset)
+                return resultRange.shiftRight(1)
             }
     }
 
@@ -104,10 +108,13 @@ private enum class TrailingCommaAction {
     ADD, REFORMAT, REMOVE;
 
     companion object {
-        fun create(commaOwner: KtElement): TrailingCommaAction = when {
+        fun create(commaOwner: KtElement): TrailingCommaAction {
-            needComma(commaOwner, checkExistingTrailingComma = false) -> ADD
+            val settings = CodeStyle.getSettings(commaOwner.project)
-            needComma(commaOwner, checkExistingTrailingComma = true) -> REFORMAT
+            return when {
-            else -> REMOVE
+                needComma(commaOwner, settings, checkExistingTrailingComma = false) -> ADD
+                needComma(commaOwner, settings, checkExistingTrailingComma = true) -> REFORMAT
+                else -> REMOVE
+            }
         }
     }
 }

idea/src/org/jetbrains/kotlin/idea/quickfix/ReformatQuickFix.kt

@@ -8,15 +8,19 @@ package org.jetbrains.kotlin.idea.quickfix
 import com.intellij.codeInspection.LocalQuickFix
 import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import com.intellij.psi.codeStyle.CodeStyleManager
+import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
+
+class ReformatQuickFix(private val description: String, element: PsiElement? = null) : LocalQuickFix {
+    private val elementPoint = element?.createSmartPointer()
 
-class ReformatQuickFix(val description: String) : LocalQuickFix {
     override fun getName(): String = description
 
     override fun getFamilyName(): String = name
 
     override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
-        descriptor.psiElement?.let {
+        (elementPoint?.element ?: descriptor.psiElement)?.let {
             CodeStyleManager.getInstance(project).reformat(it)
         }
     }

idea/testData/inspectionsLocal/trailingComma/addComma2.kt

@@ -0,0 +1,5 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Add trailing comma
+
+fun a(i: Int,
+    b: Boolean)<caret> = Unit

idea/testData/inspectionsLocal/trailingComma/addComma2.kt.after

@@ -0,0 +1,7 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Add trailing comma
+
+fun a(
+        i: Int,
+        b: Boolean,
+) = Unit

idea/testData/inspectionsLocal/trailingComma/addComma3.kt

@@ -0,0 +1,5 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// PROBLEM: none
+
+fun a(i: Int,
+    b: Boolea<caret>n) = Unit

idea/testData/inspectionsLocal/trailingComma/addComma4.kt

@@ -0,0 +1,5 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// PROBLEM: none
+
+fun a(i: Int,
+    b: Boolean) <caret>= Unit

idea/testData/inspectionsLocal/trailingComma/addComma5.kt

@@ -0,0 +1,5 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Add trailing comma
+
+fun a(i: Int,
+    b: Boolean<caret> ) = Unit

idea/testData/inspectionsLocal/trailingComma/addComma5.kt.after

@@ -0,0 +1,7 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Add trailing comma
+
+fun a(
+        i: Int,
+        b: Boolean,
+) = Unit

idea/testData/inspectionsLocal/trailingComma/addComma6.kt

@@ -0,0 +1,8 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Add trailing comma
+
+val x = {
+        x: String, y
+    : String<caret>->
+                val a = 2
+}

idea/testData/inspectionsLocal/trailingComma/addComma6.kt.after

@@ -0,0 +1,10 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Add trailing comma
+
+val x = {
+    x: String,
+    y
+    : String,
+    ->
+                val a = 2
+}

idea/testData/inspectionsLocal/trailingComma/changeCommaPosition2.kt

@@ -0,0 +1,5 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Fix comma position
+
+fun a(i: Int /*
+*/,<caret> b: Boolean) = Unit

idea/testData/inspectionsLocal/trailingComma/changeCommaPosition2.kt.after

@@ -0,0 +1,8 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Fix comma position
+
+fun a(
+        i: Int, /*
+*/
+        b: Boolean,
+) = Unit

idea/testData/inspectionsLocal/trailingComma/changeCommaPosition3.kt

@@ -0,0 +1,5 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// PROBLEM: none
+
+fun a(i: Int /*
+*<caret>/, b: Boolean) = Unit

idea/testData/inspectionsLocal/trailingComma/changeCommaPosition4.kt

@@ -0,0 +1,5 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// PROBLEM: none
+
+fun a(i: Int /*
+*/, <caret>b: Boolean) = Unit

idea/testData/inspectionsLocal/trailingComma/removeComma2.kt

@@ -0,0 +1,4 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Remove trailing comma
+
+fun a(i: Int, b: Boolean<caret>,) = Unit

idea/testData/inspectionsLocal/trailingComma/removeComma2.kt.after

@@ -0,0 +1,4 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// FIX: Remove trailing comma
+
+fun a(i: Int, b: Boolean<caret>) = Unit

idea/testData/inspectionsLocal/trailingComma/removeComma3.kt

@@ -0,0 +1,4 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// PROBLEM: none
+
+fun a(i: Int, b: Boolea<caret>n,) = Unit

idea/testData/inspectionsLocal/trailingComma/removeComma4.kt

@@ -0,0 +1,4 @@
+// COMPILER_ARGUMENTS: -XXLanguage:+TrailingCommas
+// PROBLEM: none
+
+fun a(i: Int, b: Boolean,)<caret> = Unit

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -12271,6 +12271,31 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             runTest("idea/testData/inspectionsLocal/trailingComma/addComma.kt");
         }
 
+        @TestMetadata("addComma2.kt")
+        public void testAddComma2() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/addComma2.kt");
+        }
+
+        @TestMetadata("addComma3.kt")
+        public void testAddComma3() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/addComma3.kt");
+        }
+
+        @TestMetadata("addComma4.kt")
+        public void testAddComma4() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/addComma4.kt");
+        }
+
+        @TestMetadata("addComma5.kt")
+        public void testAddComma5() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/addComma5.kt");
+        }
+
+        @TestMetadata("addComma6.kt")
+        public void testAddComma6() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/addComma6.kt");
+        }
+
         public void testAllFilesPresentInTrailingComma() throws Exception {
             KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/inspectionsLocal/trailingComma"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), null, true);
         }
@@ -12280,10 +12305,40 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             runTest("idea/testData/inspectionsLocal/trailingComma/changeCommaPosition.kt");
         }
 
+        @TestMetadata("changeCommaPosition2.kt")
+        public void testChangeCommaPosition2() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/changeCommaPosition2.kt");
+        }
+
+        @TestMetadata("changeCommaPosition3.kt")
+        public void testChangeCommaPosition3() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/changeCommaPosition3.kt");
+        }
+
+        @TestMetadata("changeCommaPosition4.kt")
+        public void testChangeCommaPosition4() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/changeCommaPosition4.kt");
+        }
+
         @TestMetadata("removeComma.kt")
         public void testRemoveComma() throws Exception {
             runTest("idea/testData/inspectionsLocal/trailingComma/removeComma.kt");
         }
+
+        @TestMetadata("removeComma2.kt")
+        public void testRemoveComma2() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/removeComma2.kt");
+        }
+
+        @TestMetadata("removeComma3.kt")
+        public void testRemoveComma3() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/removeComma3.kt");
+        }
+
+        @TestMetadata("removeComma4.kt")
+        public void testRemoveComma4() throws Exception {
+            runTest("idea/testData/inspectionsLocal/trailingComma/removeComma4.kt");
+        }
     }
 
     @TestMetadata("idea/testData/inspectionsLocal/unlabeledReturnInsideLambda")