"Main should return Unit": report on identifier, delete block body type

idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt

@@ -20,6 +20,8 @@ import org.jetbrains.kotlin.idea.search.usagesSearch.descriptor
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.namedFunctionVisitor
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 class MainFunctionReturnUnitInspection : AbstractKotlinInspection() {
@@ -31,24 +33,11 @@ class MainFunctionReturnUnitInspection : AbstractKotlinInspection() {
             if (!MainFunctionDetector.isMain(descriptor, checkReturnType = false)) return
             if (MainFunctionDetector.isMainReturnType(descriptor)) return
 
-            val funKeyword = function.funKeyword ?: return
-            val valueParameterList = function.valueParameterList ?: return
-            val typeReference = function.typeReference
-            val start = function.startOffset
-            val highlightRange = TextRange(
-                funKeyword.startOffset - start,
-                if (typeReference != null) typeReference.endOffset - start else valueParameterList.endOffset - start
-            )
-
             holder.registerProblem(
-                holder.manager.createProblemDescriptor(
-                    function,
-                    highlightRange,
-                    "main should return Unit",
-                    ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
-                    isOnTheFly,
-                    ChangeMainFunctionReturnTypeToUnitFix(function.typeReference != null)
-                )
+                function.nameIdentifier ?: function,
+                "main should return Unit",
+                ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
+                ChangeMainFunctionReturnTypeToUnitFix(function.typeReference != null)
             )
         })
     }
@@ -60,7 +49,11 @@ private class ChangeMainFunctionReturnTypeToUnitFix(private val hasExplicitRetur
     override fun getFamilyName() = name
 
     override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
-        val function = descriptor.psiElement as? KtNamedFunction ?: return
-        function.setType(KotlinBuiltIns.FQ_NAMES.unit.asString())
+        val function = descriptor.psiElement.getNonStrictParentOfType<KtNamedFunction>() ?: return
+        if (function.hasBlockBody()) {
+            function.typeReference = null
+        } else {
+            function.setType(KotlinBuiltIns.FQ_NAMES.unit.asString())
+        }
     }
 }

idea/testData/inspectionsLocal/mainFunctionReturnUnit/notMain.kt

@@ -1,2 +1,2 @@
 // PROBLEM: none
-<caret>fun test(args: Array<String>): Int = 1
+fun <caret>test(args: Array<String>): Int = 1

idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt

@@ -1,3 +1,3 @@
 // FIX: Add explicit Unit return type
 // DISABLE-ERRORS
-<caret>fun main(args: Array<String>) = 1
+fun <caret>main(args: Array<String>) = 1

idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after

@@ -1,3 +1,3 @@
 // FIX: Add explicit Unit return type
 // DISABLE-ERRORS
-<caret>fun main(args: Array<String>): Unit = 1
+fun main(args: Array<String>): Unit = 1

idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt

@@ -1,3 +1,3 @@
 // FIX: Change return type to Unit
 // DISABLE-ERRORS
-<caret>fun main(args: Array<String>): Int = 1
+fun <caret>main(args: Array<String>): Int = 1

idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt.after

@@ -1,3 +1,3 @@
 // FIX: Change return type to Unit
 // DISABLE-ERRORS
-<caret>fun main(args: Array<String>): Unit = 1
+fun main(args: Array<String>): Unit = 1

idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt

@@ -1,4 +1,4 @@
 // DISABLE-ERRORS
-<caret>fun main(args: Array<String>): Int {
+fun <caret>main(args: Array<String>): Int {
     return 1
 }

idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt.after

@@ -1,4 +1,4 @@
 // DISABLE-ERRORS
-<caret>fun main(args: Array<String>): Unit {
+fun main(args: Array<String>) {
     return 1
 }

idea/testData/inspectionsLocal/mainFunctionReturnUnit/unit.kt

@@ -1,2 +1,2 @@
 // PROBLEM: none
-<caret>fun main(args: Array<String>) {}
+fun <caret>main(args: Array<String>) {}