From 1d3ba5641d345ec912e4cf7714ad377f8faabd46 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Wed, 18 Apr 2018 11:08:13 +0300
Subject: [PATCH] "Main should return Unit": report on identifier, delete block
 body type

---
 .../MainFunctionReturnUnitInspection.kt       | 31 +++++++------------
 .../mainFunctionReturnUnit/notMain.kt         |  2 +-
 .../mainFunctionReturnUnit/notUnit1.kt        |  2 +-
 .../mainFunctionReturnUnit/notUnit1.kt.after  |  2 +-
 .../mainFunctionReturnUnit/notUnit2.kt        |  2 +-
 .../mainFunctionReturnUnit/notUnit2.kt.after  |  2 +-
 .../mainFunctionReturnUnit/notUnit3.kt        |  2 +-
 .../mainFunctionReturnUnit/notUnit3.kt.after  |  2 +-
 .../mainFunctionReturnUnit/unit.kt            |  2 +-
 9 files changed, 20 insertions(+), 27 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt
index c3162ab24b9..ae5f3b0bb2a 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt
+++ b/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())
+        }
     }
 }
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notMain.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notMain.kt
index 338945e05c4..63c0cb1e10e 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notMain.kt
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notMain.kt
@@ -1,2 +1,2 @@
 // PROBLEM: none
-<caret>fun test(args: Array<String>): Int = 1
\ No newline at end of file
+fun <caret>test(args: Array<String>): Int = 1
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt
index a19f7a6cffc..f5effba3b7a 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt
+++ b/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
\ No newline at end of file
+fun <caret>main(args: Array<String>) = 1
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after
index bff1ac1d1b5..f3bb7eeaee2 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after
+++ b/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
\ No newline at end of file
+fun main(args: Array<String>): Unit = 1
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt
index 2bdf40a4e09..5f1a662ffb3 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt
+++ b/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
\ No newline at end of file
+fun <caret>main(args: Array<String>): Int = 1
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt.after b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt.after
index a7e4e3d933b..7a4fb16785a 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit2.kt.after
+++ b/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
\ No newline at end of file
+fun main(args: Array<String>): Unit = 1
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt
index 313a98e71ae..e6735f997ca 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt
+++ b/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
 }
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt.after b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt.after
index f326f1244ab..8d82f299bc2 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit3.kt.after
+++ b/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
 }
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/unit.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/unit.kt
index 95ed374bd53..3d13cbd9536 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/unit.kt
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/unit.kt
@@ -1,2 +1,2 @@
 // PROBLEM: none
-<caret>fun main(args: Array<String>) {}
\ No newline at end of file
+fun <caret>main(args: Array<String>) {}
\ No newline at end of file