Keep return type when converting function to function with expression…

… body (#1017)

### What's done:
* fixed bug in WRONG_NEWLINES
* now don't remove return type
Closes #965

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

@@ -318,7 +318,7 @@ class NewlinesRule(configRules: List<RulesConfig>) : DiktatRule(
         }
     }
 
-    @Suppress("UnsafeCallOnNullableType")
+    @Suppress("UnsafeCallOnNullableType", "AVOID_NULL_CHECKS")
     private fun handleReturnStatement(node: ASTNode) {
         val blockNode = node.treeParent.takeIf { it.elementType == BLOCK && it.treeParent.elementType == FUN }
         val returnsUnit = node.children().count() == 1  // the only child is RETURN_KEYWORD
@@ -335,19 +335,16 @@ class NewlinesRule(configRules: List<RulesConfig>) : DiktatRule(
                 WRONG_NEWLINES.warnAndFix(configRules, emitWarn, isFixMode,
                     "functions with single return statement should be simplified to expression body", node.startOffset, node) {
                     val funNode = blockNode.treeParent
-                    // 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()!!
+                    val blockNode = funNode.findChildByType(BLOCK)
                     funNode.apply {
-                        if (needsExplicitType) {
-                            removeRange(returnType!!.treeNext, null)
-                        } else {
-                            removeRange(if (colon.treePrev.elementType == WHITE_SPACE) colon.treePrev else colon, null)
+                        if (returnType != null) {
+                            removeRange(returnType.treeNext, null)
+                            addChild(PsiWhiteSpaceImpl(" "), null)
+                        } else if (blockNode != null) {
+                            removeChild(blockNode)
                         }
-                        addChild(PsiWhiteSpaceImpl(" "), null)
                         addChild(LeafPsiElement(EQ, "="), null)
                         addChild(PsiWhiteSpaceImpl(" "), null)
                         addChild(expression.clone() as ASTNode, null)

diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/files/NewlinesRuleWarnTest.kt

@@ -1117,4 +1117,14 @@ class NewlinesRuleWarnTest : LintTestBase(::NewlinesRule) {
             rulesConfigList = rulesConfigListShort
         )
     }
+
+    @Test
+    @Tag(WarningNames.WRONG_NEWLINES)
+    fun `not complaining on fun without return type`() {
+        lintMethod(
+            """
+                |fun foo(x: Int, y: Int) = x + y
+            """.trimMargin(),
+        )
+    }
 }

diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/smoke/DiktatSmokeTest.kt

@@ -188,8 +188,8 @@ class DiktatSmokeTest : FixTestBase("test/smoke/src/main/kotlin",
             LintError(9, 3, "$DIKTAT_RULE_SET_ID:empty-block-structure", EMPTY_BLOCK_STRUCTURE_ERROR.warnText() +
                     " empty blocks are forbidden unless it is function with override keyword", false),
             LintError(12, 10, "$DIKTAT_RULE_SET_ID:kdoc-formatting", "${KDOC_NO_EMPTY_TAGS.warnText()} @return", false),
-            LintError(14, 3, "$DIKTAT_RULE_SET_ID:kdoc-formatting", "${KDOC_NO_EMPTY_TAGS.warnText()} @return", false),
-            LintError(19, 15, "$DIKTAT_RULE_SET_ID:kdoc-formatting", "${KDOC_NO_EMPTY_TAGS.warnText()} @return", false)
+            LintError(14, 8, "$DIKTAT_RULE_SET_ID:kdoc-formatting", "${KDOC_NO_EMPTY_TAGS.warnText()} @return", false),
+            LintError(19, 20, "$DIKTAT_RULE_SET_ID:kdoc-formatting", "${KDOC_NO_EMPTY_TAGS.warnText()} @return", false)
         )
     }
 

diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyExpected.kt

@@ -1,15 +1,17 @@
 package test.paragraph3.newlines
 
-fun foo() = "lorem ipsum"
+fun foo(): String = "lorem ipsum"
 
-fun foo() = "lorem ipsum"
+fun foo():String = "lorem ipsum"
 
-fun foo() = "lorem ipsum"
+fun foo() : String = "lorem ipsum"
 
 fun recFoo(): String = "lorem " + recFoo()
 
 fun recFoo():String = "lorem " + recFoo()
 
 fun recFoo(): String = "lorem " + recFoo()
 
-fun foo() = "lorem ipsum"
+fun foo() = "lorem ipsum"
+
+fun foo() = println("Logging")

diktat-rules/src/test/resources/test/paragraph3/newlines/ExpressionBodyTest.kt

@@ -24,4 +24,8 @@ fun recFoo(): String{
     return "lorem " + recFoo()
 }
 
-fun foo() = "lorem ipsum"
+fun foo() = "lorem ipsum"
+
+fun foo() {
+    return println("Logging")
+}

diktat-rules/src/test/resources/test/paragraph3/newlines/FunctionalStyleExpected.kt

@@ -11,7 +11,7 @@ fun foo(list: List<Bar>?) {
             ?:foobar
 }
 
-fun bar(x :Int,y:Int) = x+ y
+fun bar(x :Int,y:Int) :Int = x+ y
 
 fun goo() {
     x.map()

diktat-rules/src/test/resources/test/paragraph3/newlines/OneLineFunctionExpected.kt

@@ -1,6 +1,6 @@
 package test.paragraph3.newlines
 
 class Example {
-    fun doubleA() = 2 * a
-    fun doubleA() = 2 * a
+    fun doubleA(): Int = 2 * a
+    fun doubleA(): Int = 2 * a
 }

diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example1Expected.kt

@@ -12,7 +12,7 @@ class Example {
      * @param y
      * @return
      */
-    fun bar(x: Int, y: Int) = x + y
+    fun bar(x: Int, y: Int): Int = x + y
 
     /**
      * @param sub
@@ -36,7 +36,7 @@ class Example {
      * @return
      */
     fun foo(x: Int,
-            y: Int) = x +
+            y: Int): Int = x +
             (y +
                     bar(x, y)
             )