Do not remove return type from recursive functions when converting to…

… expression body (#905)

### What's done:
* Changed logic
* Added test

diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/files/NewlinesRule.kt

@@ -80,16 +80,20 @@ 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
 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

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"

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"