Add methods "ASTNode.upsertWhitespace*Me"

CHANGELOG.md

@@ -44,6 +44,8 @@ if (node.isRoot()) {
 * Update Kotlin development version to `1.7.20` and Kotlin version to `1.7.20`.
 * CLI options `--debug`, `--trace`, `--verbose` and `-v` are replaced with `--log-level=<level>` or the short version `-l=<level>, see [CLI log-level](https://pinterest.github.io/ktlint/install/cli/#logging). ([#1632](https://github.com/pinterest/ktlint/issue/1632))
 * In CLI, disable logging entirely by setting `--log-level=none` or `-l=none` ([#1652](https://github.com/pinterest/ktlint/issue/1652))
+* Rewrite `indent` rule. Solving problems in the old algorithm was very difficult. With the new algorithm this becomes a lot easier. Although the new implementation of the rule has been compared against several open source projects containing over 400,000 lines of code, it is still likely that new issues will be discovered. Please report your indentation issues so that these can be fixed as well. ([#1682](https://github.com/pinterest/ktlint/pull/1682), [#1321](https://github.com/pinterest/ktlint/issues/1321), [#1200](https://github.com/pinterest/ktlint/issues/1200), [#1562](https://github.com/pinterest/ktlint/issues/1562), [#1563](https://github.com/pinterest/ktlint/issues/1563), [#1639](https://github.com/pinterest/ktlint/issues/1639))
+* Add methods "ASTNode.upsertWhitespaceBeforeMe" and "ASTNode.upsertWhitespaceAfterMe" as replacements for "LeafElement.upsertWhitespaceBeforeMe" and "LeafElement.upsertWhitespaceAfterMe". The new methods are more versatile and allow code to be written more readable in most places. ([#1687](https://github.com/pinterest/ktlint/pull/1687))
 * Rewrite `indent` rule. Solving problems in the old algorithm was very difficult. With the new algorithm this becomes a lot easier. Although the new implementation of the rule has been compared against several open source projects containing over 400,000 lines of code, it is still likely that new issues will be discovered. Please report your indentation issues so that these can be fixed as well. ([#1682](https://github.com/pinterest/ktlint/pull/1682), [#1321](https://github.com/pinterest/ktlint/issues/1321), [#1200](https://github.com/pinterest/ktlint/issues/1200), [#1562](https://github.com/pinterest/ktlint/issues/1562), [#1563](https://github.com/pinterest/ktlint/issues/1563), [#1639](https://github.com/pinterest/ktlint/issues/1639), [#1688](https://github.com/pinterest/ktlint/issues/1688))
 
 ## [0.47.1] - 2022-09-07

ktlint-core/src/main/kotlin/com/pinterest/ktlint/core/ast/package.kt

@@ -212,6 +212,26 @@ public fun ASTNode.isPartOfComment(): Boolean =
 public fun ASTNode.children(): Sequence<ASTNode> =
     generateSequence(firstChildNode) { node -> node.treeNext }
 
+@Deprecated(message = """Marked for removal in KtLint 0.49. See KDOC""")
+/**
+ * Marked for removal in KtLint 0.49.
+ *
+ * Use [ASTNode.upsertWhitespaceBeforeMe] which operates on the [ASTNode] instead of the [LeafElement]. The new method
+ * handles more edge case and as of that a lot of code can be simplified.
+ *
+ * *Code using [LeafElement.upsertWhitespaceBeforeMe]*
+ * ```
+ * if (elementType == WHITE_SPACE) {
+ *     (this as LeafPsiElement).rawReplaceWithText("\n${blockCommentNode.lineIndent()}")
+ * } else {
+ *     (this as LeafPsiElement).upsertWhitespaceBeforeMe("\n${blockCommentNode.lineIndent()}")
+ * }
+ *  ```
+ * *Code using [ASTNode.upsertWhitespaceBeforeMe]*
+ * ```
+ * this.upsertWhitespaceBeforeMe(text)
+ *  ```
+ */
 public fun LeafElement.upsertWhitespaceBeforeMe(text: String): LeafElement {
     val s = treePrev
     return if (s != null && s.elementType == WHITE_SPACE) {
@@ -223,6 +243,14 @@ public fun LeafElement.upsertWhitespaceBeforeMe(text: String): LeafElement {
     }
 }
 
+@Deprecated(
+    message =
+    "Marked for removal in KtLint 0.49. The new insertOrReplaceWhitespaceAfterMe is more versatile as it " +
+        "operates on an AstNode instead of a LeafElement. In a lot of cases the code can be simplified as it is " +
+        "no longer needed to check whether the current node is already a whitespace or a leaf element before " +
+        "calling this method or the rawReplaceWithText.",
+    ReplaceWith("insertOrReplaceWhitespaceBeforeMe"),
+)
 public fun LeafElement.upsertWhitespaceAfterMe(text: String): LeafElement {
     val s = treeNext
     return if (s != null && s.elementType == WHITE_SPACE) {
@@ -234,6 +262,63 @@ public fun LeafElement.upsertWhitespaceAfterMe(text: String): LeafElement {
     }
 }
 
+/**
+ * Updates or inserts a new whitespace element with [text] before the given node. If the node itself is a whitespace
+ * then its contents is replaced with [text]. If the node is a (nested) composite element, the whitespace element is
+ * added after the previous leaf node.
+ */
+public fun ASTNode.upsertWhitespaceBeforeMe(text: String) {
+    if (this is LeafElement) {
+        if (this.elementType == WHITE_SPACE) {
+            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)
+            }
+        }
+    } else {
+        val prevLeaf =
+            requireNotNull(prevLeaf()) {
+                "Can not upsert a whitespace if the first node is a non-leaf node"
+            }
+        prevLeaf.upsertWhitespaceAfterMe(text)
+    }
+}
+
+private fun ASTNode.replaceWhitespaceWith(text: String) {
+    require(this.elementType == WHITE_SPACE)
+    if (this.text != text) {
+        (this.psi as LeafElement).rawReplaceWithText(text)
+    }
+}
+
+/**
+ * Updates or inserts a new whitespace element with [text] after the given node. If the node itself is a whitespace
+ * then its contents is replaced with [text]. If the node is a (nested) composite element, the whitespace element is
+ * added after the last child leaf.
+ */
+public fun ASTNode.upsertWhitespaceAfterMe(text: String) {
+    if (this is LeafElement) {
+        if (this.elementType == WHITE_SPACE) {
+            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)
+            }
+        }
+    } else {
+        lastChildLeafOrSelf().upsertWhitespaceAfterMe(text)
+    }
+}
+
 @Deprecated(message = "Marked for removal in Ktlint 0.48. See Ktlint 0.47.0 changelog for more information.")
 public fun ASTNode.visit(enter: (node: ASTNode) -> Unit) {
     enter(this)

ktlint-core/src/test/kotlin/com/pinterest/ktlint/core/ast/PackageKtTest.kt

@@ -6,6 +6,12 @@ import com.pinterest.ktlint.core.RuleProvider
 import com.pinterest.ktlint.core.ast.ElementType.CLASS
 import com.pinterest.ktlint.core.ast.ElementType.CLASS_BODY
 import com.pinterest.ktlint.core.ast.ElementType.ENUM_ENTRY
+import com.pinterest.ktlint.core.ast.ElementType.FUN
+import com.pinterest.ktlint.core.ast.ElementType.LPAR
+import com.pinterest.ktlint.core.ast.ElementType.RPAR
+import com.pinterest.ktlint.core.ast.ElementType.VALUE_PARAMETER
+import com.pinterest.ktlint.core.ast.ElementType.VALUE_PARAMETER_LIST
+import com.pinterest.ktlint.core.ast.ElementType.WHITE_SPACE
 import com.pinterest.ktlint.core.internal.createRuleExecutionContext
 import org.assertj.core.api.Assertions.assertThat
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
@@ -230,6 +236,230 @@ class PackageKtTest {
         assertThat(actual).isTrue
     }
 
+    @Nested
+    inner class UpsertWhitespaceBeforeMe {
+        @Test
+        fun `Given a whitespace node and upsert a whitespace before the node (RPAR) then replace the current whitespace element`() {
+            val code =
+                """
+                fun foo( ) = 42
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(
+                ) = 42
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        findChildByType(FUN)
+                            ?.findChildByType(VALUE_PARAMETER_LIST)
+                            ?.findChildByType(WHITE_SPACE)
+                            ?.upsertWhitespaceBeforeMe("\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+
+        @Test
+        fun `Given a node (RPAR) which is preceded by a non-whitespace leaf element (LPAR) and upsert a whitespace before the node (RPAR) then create a new whitespace element before the node (RPAR)`() {
+            val code =
+                """
+                fun foo() = 42
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(
+
+                ) = 42
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        nextLeaf { it.elementType == RPAR }
+                            ?.upsertWhitespaceBeforeMe("\n\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+
+        @Test
+        fun `Given a node (RPAR) which is preceded by a whitespace leaf element and upsert a whitespace before the node (RPAR) then replace the whitespace element before the node (RPAR)`() {
+            val code =
+                """
+                fun foo( ) = 42
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(
+
+                ) = 42
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        nextLeaf { it.elementType == RPAR }
+                            ?.upsertWhitespaceBeforeMe("\n\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+
+        @Test
+        fun `Given a node (VALUE_PARAMETER) which is preceded by a non-whitespace leaf element (LPAR) and upsert a whitespace before the node (VALUE_PARAMETER) then create a new whitespace element before the node (VALUE_PARAMETER)`() {
+            val code =
+                """
+                fun foo(string: String) = 42
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(
+                    string: String) = 42
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        findChildByType(FUN)
+                            ?.findChildByType(VALUE_PARAMETER_LIST)
+                            ?.findChildByType(VALUE_PARAMETER)
+                            ?.upsertWhitespaceBeforeMe("\n    ")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+
+        @Test
+        fun `Given a node (FUN bar) which is preceded by a composite element (FUN foo) and upsert a whitespace before the node (FUN bar) then create a new whitespace element before the node (FUN bar)`() {
+            val code =
+                """
+                fun foo() = "foo"
+                fun bar() = "bar"
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo() = "foo"
+
+                fun bar() = "bar"
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        children()
+                            .last { it.elementType == FUN }
+                            .upsertWhitespaceBeforeMe("\n\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+    }
+
+    @Nested
+    inner class UpsertWhitespaceAfterMe {
+        @Test
+        fun `Given a node (LPAR) which is followed by a non-whitespace leaf element (RPAR) and upsert a whitespace after the node (LPAR) then create a new whitespace element after the node (LPAR)`() {
+            val code =
+                """
+                fun foo() = 42
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(
+
+                ) = 42
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        nextLeaf { it.elementType == LPAR }
+                            ?.upsertWhitespaceAfterMe("\n\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+
+        @Test
+        fun `Given a node (LPAR) which is followed by a whitespace leaf element and upsert a whitespace after the node (LPAR) then replace the whitespace element after the node (LPAR)`() {
+            val code =
+                """
+                fun foo( ) = 42
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(
+
+                ) = 42
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        nextLeaf { it.elementType == LPAR }
+                            ?.upsertWhitespaceAfterMe("\n\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+
+        @Test
+        fun `Given a node (VALUE_PARAMETER) which is followed by a non-whitespace leaf element (RPAR) and upsert a whitespace after the node (VALUE_PARAMETER) then create a new whitespace element after the node (VALUE_PARAMETER)`() {
+            val code =
+                """
+                fun foo(string: String) = 42
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(string: String
+                ) = 42
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        findChildByType(FUN)
+                            ?.findChildByType(VALUE_PARAMETER_LIST)
+                            ?.findChildByType(VALUE_PARAMETER)
+                            ?.upsertWhitespaceAfterMe("\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+
+        @Test
+        fun `Given a node (FUN foo) which is followed by a composite element (FUN bar) and upsert a whitespace after the node (FUN foo) then create a new whitespace element after the node (FUN foo)`() {
+            val code =
+                """
+                fun foo() = "foo"
+                fun bar() = "bar"
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo() = "foo"
+
+                fun bar() = "bar"
+                """.trimIndent()
+
+            val actual =
+                code
+                    .transformAst {
+                        children()
+                            .first { it.elementType == FUN }
+                            .upsertWhitespaceAfterMe("\n\n")
+                    }.text
+
+            assertThat(actual).isEqualTo(formattedCode)
+        }
+    }
+
+    private inline fun String.transformAst(block: FileASTNode.() -> Unit): FileASTNode =
+        transformCodeToAST(this)
+            .apply(block)
+
     private fun transformCodeToAST(code: String) =
         createRuleExecutionContext(
             KtLint.ExperimentalParams(

ktlint-ruleset-experimental/src/main/kotlin/com/pinterest/ktlint/ruleset/experimental/CommentWrappingRule.kt

@@ -115,11 +115,7 @@ public class CommentWrappingRule :
     ) {
         emit(startOffset, "A block comment may not be followed by any other element on that same line", true)
         if (autoCorrect) {
-            if (elementType == WHITE_SPACE) {
-                (this as LeafPsiElement).rawReplaceWithText("\n${blockCommentNode.lineIndent()}")
-            } else {
-                (this as LeafPsiElement).upsertWhitespaceBeforeMe("\n${blockCommentNode.lineIndent()}")
-            }
+            this.upsertWhitespaceBeforeMe("\n${blockCommentNode.lineIndent()}")
         }
     }
 

ktlint-ruleset-experimental/src/main/kotlin/com/pinterest/ktlint/ruleset/experimental/FunctionReturnTypeSpacingRule.kt

@@ -8,7 +8,6 @@ import com.pinterest.ktlint.core.ast.nextLeaf
 import com.pinterest.ktlint.core.ast.prevLeaf
 import com.pinterest.ktlint.core.ast.upsertWhitespaceAfterMe
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
-import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafElement
 
 public class FunctionReturnTypeSpacingRule : Rule("$experimentalRulesetId:function-return-type-spacing") {
     override fun beforeVisitChildNodes(
@@ -53,15 +52,10 @@ public class FunctionReturnTypeSpacingRule : Rule("$experimentalRulesetId:functi
             .nextLeaf()
             ?.takeIf { it.elementType == WHITE_SPACE }
             .let { whiteSpaceAfterColon ->
-                if (whiteSpaceAfterColon == null) {
+                if (whiteSpaceAfterColon?.text != " ") {
                     emit(node.startOffset, "Single space expected between colon and return type", true)
                     if (autoCorrect) {
-                        (node as LeafElement).upsertWhitespaceAfterMe(" ")
-                    }
-                } else if (whiteSpaceAfterColon.text != " ") {
-                    emit(node.startOffset, "Unexpected whitespace", true)
-                    if (autoCorrect) {
-                        (whiteSpaceAfterColon as LeafElement).rawReplaceWithText(" ")
+                        node.upsertWhitespaceAfterMe(" ")
                     }
                 }
             }