Allow to convert empty Unit returning functions to expression form #KT-13588 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt

@@ -61,11 +61,11 @@ class ConvertToBlockBodyIntention : SelfTargetingIntention<KtDeclarationWithBody
 
             fun generateBody(returnsValue: Boolean): KtExpression {
                 val bodyType = body.analyze().getType(body)
+                val factory = KtPsiFactory(declaration)
+                if (bodyType != null && bodyType.isUnit() && body is KtNameReferenceExpression) return factory.createEmptyBody()
                 val unitWhenAsResult = (bodyType == null || bodyType.isUnit()) && body.resultingWhens().isNotEmpty()
                 val needReturn = returnsValue &&
                                  (bodyType == null || (!bodyType.isUnit() && !bodyType.isNothing()))
-
-                val factory = KtPsiFactory(declaration)
                 val statement = if (needReturn || unitWhenAsResult) factory.createExpressionByPattern("return $0", body) else body
                 return factory.createSingleStatementBlock(statement)
             }

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToExpressionBodyIntention.kt

@@ -62,9 +62,10 @@ class ConvertToExpressionBodyIntention : SelfTargetingOffsetIndependentIntention
     }
 
     private fun applyTo(declaration: KtDeclarationWithBody, deleteTypeHandler: ((KtCallableDeclaration) -> Unit)?) {
-        val value = calcValue(declaration)!!
+        val block = declaration.blockExpression() ?: return
+        val value = calcValue(block)!!
 
-        if (!declaration.hasDeclaredReturnType() && declaration is KtNamedFunction) {
+        if (!declaration.hasDeclaredReturnType() && declaration is KtNamedFunction && block.statements.isNotEmpty()) {
             val valueType = value.analyze().getType(value)
             if (valueType == null || !KotlinBuiltIns.isUnit(valueType)) {
                 declaration.setType(KotlinBuiltIns.FQ_NAMES.unit.asString(), shortenReferences = true)
@@ -87,13 +88,24 @@ class ConvertToExpressionBodyIntention : SelfTargetingOffsetIndependentIntention
         }
     }
 
-    private fun calcValue(declaration: KtDeclarationWithBody): KtExpression? {
-        if (declaration is KtFunctionLiteral) return null
-        val body = declaration.bodyExpression
-        if (!declaration.hasBlockBody() || body !is KtBlockExpression) return null
+    private fun KtDeclarationWithBody.blockExpression() = when (this) {
+        is KtFunctionLiteral -> null
+        else -> {
+            val body = bodyExpression
+            if (!hasBlockBody() || body !is KtBlockExpression) null else body
+        }
+    }
 
-        val statement = body.statements.singleOrNull() ?: return null
-        when(statement) {
+    private fun calcValue(declaration: KtDeclarationWithBody): KtExpression? {
+        val body = declaration.blockExpression() ?: return null
+        return calcValue(body)
+    }
+
+    private fun calcValue(body: KtBlockExpression): KtExpression? {
+        val bodyStatements = body.statements
+        if (bodyStatements.isEmpty()) return KtPsiFactory(body).createExpression("Unit")
+        val statement = bodyStatements.singleOrNull() ?: return null
+        when (statement) {
             is KtReturnExpression -> {
                 return statement.returnedExpression
             }
@@ -101,7 +113,7 @@ class ConvertToExpressionBodyIntention : SelfTargetingOffsetIndependentIntention
             //TODO: IMO this is not good code, there should be a way to detect that JetExpression does not have value
             is KtDeclaration, is KtLoopExpression -> return null // is JetExpression but does not have value
 
-            else  -> {
+            else -> {
                 if (statement is KtBinaryExpression && statement.operationToken in KtTokens.ALL_ASSIGNMENTS) return null // assignment does not have value
 
                 val context = statement.analyze()
@@ -113,7 +125,7 @@ class ConvertToExpressionBodyIntention : SelfTargetingOffsetIndependentIntention
                         return null
                     }
                     val resultingWhens = statement.resultingWhens()
-                    if (resultingWhens.any { it.elseExpression == null && context.get(BindingContext.EXHAUSTIVE_WHEN, it) != true}) {
+                    if (resultingWhens.any { it.elseExpression == null && context.get(BindingContext.EXHAUSTIVE_WHEN, it) != true }) {
                         return null
                     }
                 }

idea/testData/intentions/convertToBlockBody/funWithCustomUnitClass.kt

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+object Unit
+
+fun <caret>foo() = Unit

idea/testData/intentions/convertToBlockBody/funWithCustomUnitClass.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+object Unit
+
+fun foo(): Unit {
+    return Unit
+}

idea/testData/intentions/convertToBlockBody/funWithUnit.kt

@@ -0,0 +1,3 @@
+// WITH_RUNTIME
+
+fun <caret>foo() = Unit

idea/testData/intentions/convertToBlockBody/funWithUnit.kt.after

@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+
+fun foo() {
+}

idea/testData/intentions/convertToExpressionBody/funWithEmptyBody.kt

@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+
+fun <caret>foo() {
+}

idea/testData/intentions/convertToExpressionBody/funWithEmptyBody.kt.after

@@ -0,0 +1,3 @@
+// WITH_RUNTIME
+
+fun foo() = Unit

idea/testData/intentions/convertToExpressionBody/funWithEmptyBody2.kt

@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+
+fun <caret>foo(): Unit {
+}

idea/testData/intentions/convertToExpressionBody/funWithEmptyBody2.kt.after

@@ -0,0 +1,3 @@
+// WITH_RUNTIME
+
+fun foo(): Unit = Unit

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -4615,12 +4615,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("funWithCustomUnitClass.kt")
+        public void testFunWithCustomUnitClass() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/funWithCustomUnitClass.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("funWithThrow.kt")
         public void testFunWithThrow() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/funWithThrow.kt");
             doTest(fileName);
         }
 
+        @TestMetadata("funWithUnit.kt")
+        public void testFunWithUnit() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/funWithUnit.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("getter.kt")
         public void testGetter() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/getter.kt");
@@ -4933,6 +4945,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("funWithEmptyBody.kt")
+        public void testFunWithEmptyBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/funWithEmptyBody.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("funWithEmptyBody2.kt")
+        public void testFunWithEmptyBody2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/funWithEmptyBody2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("funWithImplicitUnitTypeWithThrow.kt")
         public void testFunWithImplicitUnitTypeWithThrow() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/funWithImplicitUnitTypeWithThrow.kt");