Replace all function bodies with body expressions in a single run (#2395)

If a file contains multiple function bodies that have to be replaced with a body expression, then all should be replaced with a single run of ktlint.

Closes #2394

Author: paul-dingemans

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

@@ -23,7 +23,8 @@ import com.pinterest.ktlint.rule.engine.core.api.editorconfig.INDENT_STYLE_PROPE
 import com.pinterest.ktlint.rule.engine.core.api.editorconfig.MAX_LINE_LENGTH_PROPERTY
 import com.pinterest.ktlint.rule.engine.core.api.firstChildLeafOrSelf
 import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpace
-import com.pinterest.ktlint.rule.engine.core.api.leavesIncludingSelf
+import com.pinterest.ktlint.rule.engine.core.api.lastChildLeafOrSelf
+import com.pinterest.ktlint.rule.engine.core.api.leavesInClosedRange
 import com.pinterest.ktlint.rule.engine.core.api.nextSibling
 import com.pinterest.ktlint.rule.engine.core.api.prevSibling
 import com.pinterest.ktlint.ruleset.standard.StandardRule
@@ -111,7 +112,7 @@ public class FunctionExpressionBodyRule :
         require(block.elementType == BLOCK)
         block
             .takeIf { it.containingOnly(RETURN) }
-            ?.takeUnless { it.containsMultipleReturns() }
+            ?.takeUnless { it.countReturnKeywords() > 1 }
             ?.findChildByType(RETURN)
             ?.findChildByType(RETURN_KEYWORD)
             ?.nextSibling { !it.isWhiteSpace() }
@@ -164,8 +165,9 @@ public class FunctionExpressionBodyRule :
                 .singleOrNull()
                 ?.elementType
 
-    private fun ASTNode.containsMultipleReturns() =
-        firstChildLeafOrSelf().leavesIncludingSelf().count { it.elementType == RETURN_KEYWORD } > 1
+    private fun ASTNode.countReturnKeywords() =
+        leavesInClosedRange(this.firstChildLeafOrSelf(), this.lastChildLeafOrSelf())
+            .count { it.elementType == RETURN_KEYWORD }
 
     private fun ASTNode.createUnitTypeReference() =
         PsiFileFactory

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

@@ -1,6 +1,7 @@
 package com.pinterest.ktlint.ruleset.standard.rules
 
 import com.pinterest.ktlint.test.KtLintAssertThat.Companion.assertThatRule
+import com.pinterest.ktlint.test.LintViolation
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.ValueSource
@@ -168,4 +169,27 @@ class FunctionExpressionBodyRuleTest {
             """.trimIndent()
         functionExpressionBodyRule(code).hasNoLintViolations()
     }
+
+    @Test
+    fun `Issue 2394 - Given multiple function bodies with a single return statement`() {
+        val code =
+            """
+            fun foo1(): String {
+                return "foo1"
+            }
+            fun foo2(): String {
+                return "foo2"
+            }
+            """.trimIndent()
+        val formattedCode =
+            """
+            fun foo1(): String = "foo1"
+            fun foo2(): String = "foo2"
+            """.trimIndent()
+        functionExpressionBodyRule(code)
+            .hasLintViolations(
+                LintViolation(1, 20, "Function body should be replaced with body expression"),
+                LintViolation(4, 20, "Function body should be replaced with body expression"),
+            ).isFormattedAs(formattedCode)
+    }
 }