From 2201861e8c01ae0effd47e7e20a215834339f3db Mon Sep 17 00:00:00 2001
From: Peter Trifanov <peter.trifanov@mail.ru>
Date: Thu, 3 Jun 2021 12:06:22 +0300
Subject: [PATCH] Do not remove return type from recursive functions when
 converting to expression body (#905)

### What's done:
* Changed logic
* Added test
---
 .../rules/chapter3/files/NewlinesRule.kt      | 25 ++++++++++++++++++-
 .../newlines/ExpressionBodyExpected.kt        | 12 +++++++++
 .../paragraph3/newlines/ExpressionBodyTest.kt | 22 ++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/files/NewlinesRule.kt b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/files/NewlinesRule.kt
index 149d0a9878..dde96f9f0d 100644
--- a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/files/NewlinesRule.kt
+++ b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/files/NewlinesRule.kt
@@ -80,6 +80,7 @@ import com.pinterest.ktlint.core.ast.isWhiteSpaceWithNewline
 import com.pinterest.ktlint.core.ast.nextCodeSibling
 import com.pinterest.ktlint.core.ast.parent
 import com.pinterest.ktlint.core.ast.prevCodeSibling
+import com.pinterest.ktlint.core.ast.prevSibling
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafPsiElement
@@ -87,9 +88,12 @@ import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.PsiWhiteSpaceImpl
 import org.jetbrains.kotlin.com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
 import org.jetbrains.kotlin.psi.KtBinaryExpression
+import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtParameterList
+import org.jetbrains.kotlin.psi.KtReferenceExpression
 import org.jetbrains.kotlin.psi.KtSuperTypeList
 import org.jetbrains.kotlin.psi.KtValueArgumentList
+import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.children
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.psi.psiUtil.siblings
@@ -335,9 +339,15 @@ class NewlinesRule(configRules: List<RulesConfig>) : DiktatRule(
                     // if return type is not Unit, then there should be type specification
                     // otherwise code won't compile and colon being null is correctly invalid
                     val colon = funNode.findChildByType(COLON)!!
+                    val returnType = (funNode.psi as? KtNamedFunction)?.typeReference?.node
+                    val needsExplicitType = returnType != null && (funNode.psi as? KtNamedFunction)?.isRecursive() == true
                     val expression = node.findChildByType(RETURN_KEYWORD)!!.nextCodeSibling()!!
                     funNode.apply {
-                        removeRange(if (colon.treePrev.elementType == WHITE_SPACE) colon.treePrev else colon, null)
+                        if (needsExplicitType) {
+                            removeRange(returnType!!.treeNext, null)
+                        } else {
+                            removeRange(if (colon.treePrev.elementType == WHITE_SPACE) colon.treePrev else colon, null)
+                        }
                         addChild(PsiWhiteSpaceImpl(" "), null)
                         addChild(LeafPsiElement(EQ, "="), null)
                         addChild(PsiWhiteSpaceImpl(" "), null)
@@ -627,3 +637,16 @@ class NewlinesRule(configRules: List<RulesConfig>) : DiktatRule(
         private val parenthesesTypes = TokenSet.create(CONDITION, WHEN, VALUE_ARGUMENT)
     }
 }
+
+/**
+ * Checks whether [this] function is recursive, i.e. calls itself inside from it's body
+ *
+ * @return true if function is recursive, false otherwise
+ */
+fun KtNamedFunction.isRecursive() = bodyBlockExpression
+    ?.statements?.any { statement ->
+    statement.anyDescendantOfType<KtReferenceExpression> {
+        it.text == this@isRecursive.name
+    }
+}
+    ?: false
diff --git a/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyExpected.kt b/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyExpected.kt
index 171b20ffd8..a52c30ce8d 100644
--- a/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyExpected.kt
+++ b/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyExpected.kt
@@ -1,3 +1,15 @@
 package test.paragraph3.newlines
 
 fun foo() = "lorem ipsum"
+
+fun foo() = "lorem ipsum"
+
+fun foo() = "lorem ipsum"
+
+fun recFoo(): String = "lorem " + recFoo()
+
+fun recFoo():String = "lorem " + recFoo()
+
+fun recFoo(): String = "lorem " + recFoo()
+
+fun foo() = "lorem ipsum"
\ No newline at end of file
diff --git a/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyTest.kt b/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyTest.kt
index f0123d18e6..3f8de25113 100644
--- a/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyTest.kt
+++ b/diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyTest.kt
@@ -3,3 +3,25 @@ package test.paragraph3.newlines
 fun foo(): String {
     return "lorem ipsum"
 }
+
+fun foo():String{
+    return "lorem ipsum"
+}
+
+fun foo() : String {
+    return "lorem ipsum"
+}
+
+fun recFoo(): String {
+    return "lorem " + recFoo()
+}
+
+fun recFoo():String {
+    return "lorem " + recFoo()
+}
+
+fun recFoo(): String{
+    return "lorem " + recFoo()
+}
+
+fun foo() = "lorem ipsum"
\ No newline at end of file