Fix malformed AST when && or || is at start of line chain-wrapping

CHANGELOG.md

@@ -22,6 +22,7 @@ This project adheres to [Semantic Versioning](https://semver.org/).
 * Ignore function naming in Kotest classes `function-naming` [#2289](https://github.com/pinterest/ktlint/issue/2289)
 * Prevent wrapping of nested multiline binary expression before operation reference as it results in a compilation error `multiline-expression-wrapping` [#2286](https://github.com/pinterest/ktlint/issue/2286)
 * Force blank line before object declaration if preceded by another declaration `blank-line-before-declaration` [#2284](https://github.com/pinterest/ktlint/issues/2284)
+* Fix malformed AST when `&&` or `||` is at start of line `chain-wrapping` [#2297](https://github.com/pinterest/ktlint/issues/2297) 
 
 ### Changed
 

ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/ChainWrappingRule.kt

@@ -10,12 +10,12 @@ 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.MINUS
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.MUL
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.OPERATION_REFERENCE
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.OROR
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.PERC
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.PLUS
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.PREFIX_EXPRESSION
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.SAFE_ACCESS
-import com.pinterest.ktlint.rule.engine.core.api.ElementType.WHITE_SPACE
 import com.pinterest.ktlint.rule.engine.core.api.IndentConfig
 import com.pinterest.ktlint.rule.engine.core.api.IndentConfig.Companion.DEFAULT_INDENT_CONFIG
 import com.pinterest.ktlint.rule.engine.core.api.RuleId
@@ -29,13 +29,14 @@ import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpaceWithNewline
 import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpaceWithoutNewline
 import com.pinterest.ktlint.rule.engine.core.api.nextCodeLeaf
 import com.pinterest.ktlint.rule.engine.core.api.nextLeaf
+import com.pinterest.ktlint.rule.engine.core.api.nextSibling
 import com.pinterest.ktlint.rule.engine.core.api.prevCodeLeaf
+import com.pinterest.ktlint.rule.engine.core.api.prevCodeSibling
 import com.pinterest.ktlint.rule.engine.core.api.prevLeaf
 import com.pinterest.ktlint.rule.engine.core.api.upsertWhitespaceAfterMe
 import com.pinterest.ktlint.rule.engine.core.api.upsertWhitespaceBeforeMe
 import com.pinterest.ktlint.ruleset.standard.StandardRule
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
-import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafElement
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
@@ -105,14 +106,16 @@ public class ChainWrappingRule :
                 return
             }
             val prevLeaf = node.prevLeaf()
-            if (
-                prevLeaf?.elementType == WHITE_SPACE &&
-                prevLeaf.textContains('\n') &&
+            if (node.elementType != MUL && prevLeaf?.isPartOfSpread() == true) {
                 // fn(*typedArray<...>()) case
-                (elementType != MUL || !prevLeaf.isPartOfSpread()) &&
+                return
+            }
+            if (prefixTokens.contains(elementType) && node.isInPrefixPosition()) {
                 // unary +/-
-                (!prefixTokens.contains(elementType) || !node.isInPrefixPosition())
-            ) {
+                return
+            }
+
+            if (prevLeaf != null && prevLeaf.isWhiteSpaceWithNewline()) {
                 emit(node.startOffset, "Line must not begin with \"${node.text}\"", true)
                 if (autoCorrect) {
                     // rewriting
@@ -122,13 +125,35 @@ public class ChainWrappingRule :
                     // <insertionPoint><spaceBeforeComment><comment><prevLeaf="\n"><node="&&"><nextLeaf=" "> to
                     // <insertionPoint><space if needed><node="&&"><spaceBeforeComment><comment><prevLeaf="\n"><delete node="&&"><delete nextLeaf=" ">
                     val nextLeaf = node.nextLeaf()
-                    if (nextLeaf is PsiWhiteSpace) {
-                        nextLeaf.node.treeParent.removeChild(nextLeaf.node)
+                    val whiteSpaceToBeDeleted =
+                        when {
+                            nextLeaf.isWhiteSpaceWithNewline() -> {
+                                // Node is preceded and followed by whitespace. Prefer to remove the whitespace before the node as this will
+                                // change the indent of the next line
+                                prevLeaf
+                            }
+                            nextLeaf.isWhiteSpaceWithoutNewline() -> nextLeaf
+                            else -> null
+                        }
+
+                    if (node.treeParent.elementType == OPERATION_REFERENCE) {
+                        val operationReference = node.treeParent
+                        val insertBeforeSibling =
+                            operationReference
+                                .prevCodeSibling()
+                                ?.nextSibling()
+                        operationReference.treeParent.removeChild(operationReference)
+                        insertBeforeSibling?.treeParent?.addChild(operationReference, insertBeforeSibling)
+                        node.treeParent.upsertWhitespaceBeforeMe(" ")
+                    } else {
+                        val insertionPoint = prevLeaf.prevCodeLeaf() as LeafPsiElement
+                        (node as LeafPsiElement).treeParent.removeChild(node)
+                        insertionPoint.rawInsertAfterMe(node)
+                        (insertionPoint as ASTNode).upsertWhitespaceAfterMe(" ")
                     }
-                    val insertionPoint = prevLeaf.prevCodeLeaf() as LeafPsiElement
-                    (node as LeafPsiElement).treeParent.removeChild(node)
-                    insertionPoint.rawInsertAfterMe(node)
-                    (insertionPoint as ASTNode).upsertWhitespaceAfterMe(" ")
+                    whiteSpaceToBeDeleted
+                        ?.treeParent
+                        ?.removeChild(whiteSpaceToBeDeleted)
                 }
             }
         }

ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/ChainWrappingRuleTest.kt

@@ -272,4 +272,61 @@ class ChainWrappingRuleTest {
             .hasLintViolation(5, 13, "Line must not begin with \"&&\"")
             .isFormattedAs(formattedCode)
     }
+
+    @Nested
+    inner class `Issue 2297 - Given an && starting on new line should not result in malformed AST which causes NPE in multiline-expression rule` {
+        @Test
+        fun `No EOL comment on previous line`() {
+            val code =
+                """
+                fun foo(): Boolean {
+                    return true
+                            &&
+                        columns.all { col ->
+                            false
+                        }
+                }
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(): Boolean {
+                    return true &&
+                        columns.all { col ->
+                            false
+                        }
+                }
+                """.trimIndent()
+            chainWrappingRuleAssertThat(code)
+                .addAdditionalRuleProvider { MultilineExpressionWrappingRule() }
+                .hasLintViolation(3, 13, "Line must not begin with \"&&\"")
+                .isFormattedAs(formattedCode)
+        }
+
+        @Test
+        fun `EOL comment on previous line`() {
+            val code =
+                """
+                fun foo(): Boolean {
+                    return true // some comment
+                            &&
+                        columns.all { col ->
+                            false
+                        }
+                }
+                """.trimIndent()
+            val formattedCode =
+                """
+                fun foo(): Boolean {
+                    return true && // some comment
+                        columns.all { col ->
+                            false
+                        }
+                }
+                """.trimIndent()
+            chainWrappingRuleAssertThat(code)
+                .addAdditionalRuleProvider { MultilineExpressionWrappingRule() }
+                .hasLintViolation(3, 13, "Line must not begin with \"&&\"")
+                .isFormattedAs(formattedCode)
+        }
+    }
 }