Do not insert a whitespace element as first or last child inside a co…

…mposite element (#2715) Prevent inserting whitespace element as first or last element in a composite node In such cases insert the whitespace just before or after the composite element (recursively if needed) Closes #2688
pinterest · Jun 25, 2024 · a4ddfaa · a4ddfaa
1 parent 62bb733
commit a4ddfaa
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 11 deletions.
diff --git a/...engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtension.kt b/...engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtension.kt
@@ -238,11 +238,20 @@ public fun ASTNode.upsertWhitespaceBeforeMe(text: String) {
             return replaceWhitespaceWith(text)
         }
         val previous = treePrev ?: prevLeaf()
-        if (previous != null && previous.elementType == WHITE_SPACE) {
-            previous.replaceWhitespaceWith(text)
-        } else {
-            PsiWhiteSpaceImpl(text).also { psiWhiteSpace ->
-                (psi as LeafElement).rawInsertBeforeMe(psiWhiteSpace)
+        when {
+            previous?.elementType == WHITE_SPACE -> {
+                previous.replaceWhitespaceWith(text)
+            }
+
+            treeParent.firstChildNode == this -> {
+                // Never insert a whitespace node as first node in a composite node
+                treeParent.upsertWhitespaceBeforeMe(text)
+            }
+
+            else -> {
+                PsiWhiteSpaceImpl(text).also { psiWhiteSpace ->
+                    (psi as LeafElement).rawInsertBeforeMe(psiWhiteSpace)
+                }
             }
         }
     } else {
@@ -284,18 +293,27 @@ public fun ASTNode.upsertWhitespaceAfterMe(text: String) {
             return replaceWhitespaceWith(text)
         }
         val next = treeNext ?: nextLeaf()
-        if (next != null && next.elementType == WHITE_SPACE) {
-            next.replaceWhitespaceWith(text)
-        } else {
-            PsiWhiteSpaceImpl(text).also { psiWhiteSpace ->
-                (psi as LeafElement).rawInsertAfterMe(psiWhiteSpace)
+        when {
+            next?.elementType == WHITE_SPACE -> {
+                next.replaceWhitespaceWith(text)
+            }
+
+            treeParent.lastChildNode == this -> {
+                // Never insert a whitespace as last node in a composite node
+                treeParent.upsertWhitespaceAfterMe(text)
+            }
+
+            else -> {
+                PsiWhiteSpaceImpl(text).also { psiWhiteSpace ->
+                    (psi as LeafElement).rawInsertAfterMe(psiWhiteSpace)
+                }
             }
         }
     } else {
         when (val nextSibling = nextSibling()) {
             null -> {
                 // Never insert a whitespace element as last child node in a composite node. Instead, upsert just after the composite node
-                treeParent.upsertWhitespaceAfterMe(text)
+                treeParent?.upsertWhitespaceAfterMe(text)
             }
 
             is LeafElement -> {

diff --git a/...ne-core/src/test/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionTest.kt b/...ne-core/src/test/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionTest.kt
@@ -10,6 +10,7 @@ import com.pinterest.ktlint.rule.engine.core.api.ElementType.FILE
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.FUN
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.FUN_KEYWORD
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.IDENTIFIER
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.LBRACE
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.LPAR
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.MODIFIER_LIST
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.PRIVATE_KEYWORD
@@ -457,6 +458,31 @@ class ASTNodeExtensionTest {
 
             assertThat(actual).isEqualTo(formattedCode)
         }
+
+        @Test
+        fun `Issue 2688 - Given a class without a space between the identifier and the left curly brace then insert the whitespace at the correct position in the AST`() {
+            val code =
+                """
+                class Foo{
+                    // some comment
+                }
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        findChildByType(CLASS)
+                            ?.findChildByType(CLASS_BODY)
+                            ?.findChildByType(LBRACE)
+                            ?.upsertWhitespaceBeforeMe(" ")
+                    }.findChildByType(CLASS)
+                    ?.findChildByType(CLASS_BODY)
+                    ?.prevSibling()
+                    ?.let { it.elementType == WHITE_SPACE && it.text == " " }
+                    ?: false
+
+            assertThat(actual).isTrue()
+        }
     }
 
     @Nested
@@ -604,6 +630,31 @@ class ASTNodeExtensionTest {
 
             assertThat(actual).isEqualTo(formattedCode)
         }
+
+        @Test
+        fun `Issue 2688 - Given a class without a space between the identifier and the left curly brace then insert the whitespace at the correct position in the AST`() {
+            val code =
+                """
+                fun foo(){
+                    // some comment
+                }
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        findChildByType(FUN)
+                            ?.findChildByType(VALUE_PARAMETER_LIST)
+                            ?.findChildByType(RPAR)
+                            ?.upsertWhitespaceAfterMe(" ")
+                    }.findChildByType(FUN)
+                    ?.findChildByType(VALUE_PARAMETER_LIST)
+                    ?.nextSibling()
+                    ?.let { it.elementType == WHITE_SPACE && it.text == " " }
+                    ?: false
+
+            assertThat(actual).isTrue()
+        }
     }
 
     @Test