Add new experimental rule to check consistent spacing around the retu…

…rn type of a function

This rule is required for implementing pinterest#1341.

CHANGELOG.md

@@ -5,6 +5,7 @@ This project adheres to [Semantic Versioning](https://semver.org/).
 ## Unreleased
 
 ### Added
+- Add experimental rule for incorrect spacing around the function return type (`function-return-type-spacing`) ([#1341](https://github.com/pinterest/ktlint/pull/1341))
 
 ### Fixed
 

README.md

@@ -88,6 +88,7 @@ by passing the `--experimental` flag to `ktlint`.
 - `experimental:annotation-spacing`: Annotations should be separated by the annotated declaration by a single line break
 - `experimental:double-colon-spacing`: No spaces around `::`
 - `experimental:fun-keyword-spacing`: Consistent spacing after the fun keyword
+- `experimental:function-return-type-spacing`: Consistent spacing around the function return type
 - `experimental:function-type-reference-spacing`: Consistent spacing in the type reference before a function
 - `experimental:modifier-list-spacing`: Consistent spacing between modifiers in and after the last modifier in a modifier list
 - `experimental:spacing-around-angle-brackets`: No spaces around angle brackets

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

@@ -30,6 +30,7 @@ public class ExperimentalRuleSetProvider : RuleSetProvider {
         FunctionTypeReferenceSpacingRule(),
         ModifierListSpacingRule(),
         CommentWrappingRule(),
-        KdocWrappingRule()
+        KdocWrappingRule(),
+        FunctionReturnTypeSpacingRule()
     )
 }

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

@@ -0,0 +1,68 @@
+package com.pinterest.ktlint.ruleset.experimental
+
+import com.pinterest.ktlint.core.Rule
+import com.pinterest.ktlint.core.ast.ElementType.COLON
+import com.pinterest.ktlint.core.ast.ElementType.FUN
+import com.pinterest.ktlint.core.ast.ElementType.WHITE_SPACE
+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("function-return-type-spacing") {
+    override fun visit(
+        node: ASTNode,
+        autoCorrect: Boolean,
+        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit
+    ) {
+        node
+            .takeIf { node.elementType == FUN }
+            ?.let { node.findChildByType(COLON) }
+            ?.let { colonNode ->
+                removeWhiteSpaceBetweenClosingParenthesisAndColon(colonNode, emit, autoCorrect)
+                fixWhiteSpaceBetweenColonAndReturnType(colonNode, emit, autoCorrect)
+            }
+    }
+
+    private fun removeWhiteSpaceBetweenClosingParenthesisAndColon(
+        node: ASTNode,
+        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit,
+        autoCorrect: Boolean
+    ) {
+        require(node.elementType == COLON)
+        node
+            .prevLeaf()
+            ?.takeIf { it.elementType == WHITE_SPACE }
+            ?.let { whitespaceBeforeColonNode ->
+                emit(whitespaceBeforeColonNode.startOffset, "Unexpected whitespace", true)
+                if (autoCorrect) {
+                    whitespaceBeforeColonNode.treeParent?.removeChild(whitespaceBeforeColonNode)
+                }
+            }
+    }
+
+    private fun fixWhiteSpaceBetweenColonAndReturnType(
+        node: ASTNode,
+        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit,
+        autoCorrect: Boolean
+    ) {
+        require(node.elementType == COLON)
+        node
+            .nextLeaf()
+            ?.takeIf { it.elementType == WHITE_SPACE }
+            .let { whiteSpaceAfterColon ->
+                if (whiteSpaceAfterColon == null) {
+                    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(" ")
+                    }
+                }
+            }
+    }
+}

ktlint-ruleset-experimental/src/test/kotlin/com/pinterest/ktlint/ruleset/experimental/FunctionReturnTypeSpacingRuleTest.kt

@@ -0,0 +1,101 @@
+package com.pinterest.ktlint.ruleset.experimental
+
+import com.pinterest.ktlint.core.LintError
+import com.pinterest.ktlint.test.format
+import com.pinterest.ktlint.test.lint
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+
+class FunctionReturnTypeSpacingRuleTest {
+    @Test
+    fun `Given a function signature without whitespace between the closing parenthesis and the colon of the return type then do not reformat`() {
+        val code =
+            """
+            fun foo(): String = "some-result"
+            """.trimIndent()
+        assertThat(FunctionReturnTypeSpacingRule().lint(code)).isEmpty()
+        assertThat(FunctionReturnTypeSpacingRule().format(code)).isEqualTo(code)
+    }
+
+    @Test
+    fun `Given a function signature with at least one space between the closing parenthesis and the colon of the return type then reformat`() {
+        val code =
+            """
+            fun foo() : String = "some-result"
+            """.trimIndent()
+        val formattedCode =
+            """
+            fun foo(): String = "some-result"
+            """.trimIndent()
+        assertThat(FunctionReturnTypeSpacingRule().lint(code)).containsExactly(
+            LintError(1, 10, "function-return-type-spacing", "Unexpected whitespace")
+        )
+        assertThat(FunctionReturnTypeSpacingRule().format(code)).isEqualTo(formattedCode)
+    }
+
+    @Test
+    fun `Given a function signature with a newline between the closing parenthesis and the colon of the return type then reformat`() {
+        val code =
+            """
+            fun foo()
+                : String = "some-result"
+            """.trimIndent()
+        val formattedCode =
+            """
+            fun foo(): String = "some-result"
+            """.trimIndent()
+        assertThat(FunctionReturnTypeSpacingRule().lint(code)).containsExactly(
+            LintError(1, 10, "function-return-type-spacing", "Unexpected whitespace")
+        )
+        assertThat(FunctionReturnTypeSpacingRule().format(code)).isEqualTo(formattedCode)
+    }
+
+    @Test
+    fun `Given a function signature without space between the colon and the return type then reformat`() {
+        val code =
+            """
+            fun foo():String = "some-result"
+            """.trimIndent()
+        val formattedCode =
+            """
+            fun foo(): String = "some-result"
+            """.trimIndent()
+        assertThat(FunctionReturnTypeSpacingRule().lint(code)).containsExactly(
+            LintError(1, 10, "function-return-type-spacing", "Single space expected between colon and return type")
+        )
+        assertThat(FunctionReturnTypeSpacingRule().format(code)).isEqualTo(formattedCode)
+    }
+
+    @Test
+    fun `Given a function signature with multiple spaces between the colon and the return type then reformat`() {
+        val code =
+            """
+            fun foo():  String = "some-result"
+            """.trimIndent()
+        val formattedCode =
+            """
+            fun foo(): String = "some-result"
+            """.trimIndent()
+        assertThat(FunctionReturnTypeSpacingRule().lint(code)).containsExactly(
+            LintError(1, 10, "function-return-type-spacing", "Unexpected whitespace")
+        )
+        assertThat(FunctionReturnTypeSpacingRule().format(code)).isEqualTo(formattedCode)
+    }
+
+    @Test
+    fun `Given a function signature with a new line between the colon and the return type then reformat`() {
+        val code =
+            """
+            fun foo():
+                String = "some-result"
+            """.trimIndent()
+        val formattedCode =
+            """
+            fun foo(): String = "some-result"
+            """.trimIndent()
+        assertThat(FunctionReturnTypeSpacingRule().lint(code)).containsExactly(
+            LintError(1, 10, "function-return-type-spacing", "Unexpected whitespace")
+        )
+        assertThat(FunctionReturnTypeSpacingRule().format(code)).isEqualTo(formattedCode)
+    }
+}