Redundant Unit inspection: fix false positive for single expression

So #KT-22097 Fixed

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

@@ -9,30 +9,58 @@ import com.intellij.codeInspection.*
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.previousStatement
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 
 class RedundantUnitExpressionInspection : AbstractKotlinInspection(), CleanupLocalInspectionTool {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean, session: LocalInspectionToolSession): PsiElementVisitor {
         return referenceExpressionVisitor(fun(expression) {
-            if (KotlinBuiltIns.FQ_NAMES.unit.shortName() != (expression as? KtNameReferenceExpression)?.getReferencedNameAsName()) {
+            if (expression.isRedundantUnit()) {
-                return
+                holder.registerProblem(
+                    expression,
+                    "Redundant 'Unit'",
+                    ProblemHighlightType.LIKE_UNUSED_SYMBOL,
+                    RemoveRedundantUnitFix()
+                )
             }
-
-            val parent = expression.parent
-            if (parent !is KtReturnExpression && parent !is KtBlockExpression) return
-
-            // Do not report just 'Unit' in function literals (return@label Unit is OK even in literals)
-            if (parent is KtBlockExpression && parent.getParentOfType<KtFunctionLiteral>(strict = true) != null) return
-
-            holder.registerProblem(expression,
-                                   "Redundant 'Unit'",
-                                   ProblemHighlightType.LIKE_UNUSED_SYMBOL,
-                                   RemoveRedundantUnitFix())
         })
     }
 }
 
+private fun KtReferenceExpression.isRedundantUnit(): Boolean {
+    if (!isUnitLiteral()) return false
+    val parent = this.parent ?: return false
+    if (parent is KtReturnExpression) return true
+    if (parent is KtBlockExpression) {
+        // Do not report just 'Unit' in function literals (return@label Unit is OK even in literals)
+        if (parent.getParentOfType<KtFunctionLiteral>(strict = true) != null) return false
+
+        if (this == parent.lastBlockStatementOrThis()) {
+            val prev = this.previousStatement() ?: return true
+            if (prev.isUnitLiteral()) return true
+            if (prev.getResolvedCall(analyze())?.resultingDescriptor?.returnType?.isUnit() == true) return true
+            if (prev is KtDeclaration) {
+                return if (prev is KtFunction)
+                    parent.parent?.parent?.let { it is KtIfExpression || it is KtWhenExpression } != true
+                else
+                    true
+            }
+            return false
+        }
+
+        return true
+    }
+    return false
+}
+
+private fun KtExpression.isUnitLiteral(): Boolean =
+    KotlinBuiltIns.FQ_NAMES.unit.shortName() == (this as? KtNameReferenceExpression)?.getReferencedNameAsName()
+
 private class RemoveRedundantUnitFix : LocalQuickFix {
     override fun getName() = "Remove redundant 'Unit'"
 

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterClass.kt

@@ -0,0 +1,5 @@
+fun test(b: Boolean): Unit = if (b) {
+    class A
+    <caret>Unit
+} else {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterClass.kt.after

@@ -0,0 +1,4 @@
+fun test(b: Boolean): Unit = if (b) {
+    class A
+} else {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterConstant.kt

@@ -0,0 +1,6 @@
+// PROBLEM: none
+fun test(b: Boolean): Unit = if (b) {
+    1
+    <caret>Unit
+} else {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterFunInIf.kt

@@ -0,0 +1,6 @@
+// PROBLEM: none
+fun test(b: Boolean): Unit = if (b) {
+    fun a() = 1
+    <caret>Unit
+} else {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterFunInTry.kt

@@ -0,0 +1,5 @@
+fun test(b: Boolean): Unit = try {
+    fun a() {}
+    <caret>Unit
+} catch (e: Exception) {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterFunInTry.kt.after

@@ -0,0 +1,4 @@
+fun test(b: Boolean): Unit = try {
+    fun a() {}
+} catch (e: Exception) {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterFunInWhen.kt

@@ -0,0 +1,9 @@
+// PROBLEM: none
+fun test(b: Boolean): Unit = when (b) {
+    true -> {
+        fun a() {}
+        <caret>Unit
+    }
+    else -> {
+    }
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterIntFunctionCall.kt

@@ -0,0 +1,8 @@
+// PROBLEM: none
+fun test(b: Boolean): Unit = if (b) {
+    int()
+    <caret>Unit
+} else {
+}
+
+fun int() = 1

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterUnit.kt

@@ -0,0 +1,5 @@
+fun test(b: Boolean): Unit = if (b) {
+    Unit
+    <caret>Unit
+} else {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterUnit.kt.after

@@ -0,0 +1,4 @@
+fun test(b: Boolean): Unit = if (b) {
+    Unit
+} else {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterUnitFunctionCall.kt

@@ -0,0 +1,8 @@
+fun test(b: Boolean): Unit = if (b) {
+    unit()
+    <caret>Unit
+} else {
+}
+
+fun unit() {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterUnitFunctionCall.kt.after

@@ -0,0 +1,7 @@
+fun test(b: Boolean): Unit = if (b) {
+    unit()
+} else {
+}
+
+fun unit() {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterVal.kt

@@ -0,0 +1,5 @@
+fun test(b: Boolean): Unit = if (b) {
+    val a = 1
+    <caret>Unit
+} else {
+}

idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterVal.kt.after

@@ -0,0 +1,4 @@
+fun test(b: Boolean): Unit = if (b) {
+    val a = 1
+} else {
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -2988,6 +2988,60 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/inspectionsLocal/redundantUnitExpression"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), TargetBackend.ANY, true);
         }
 
+        @TestMetadata("atLastAfterClass.kt")
+        public void testAtLastAfterClass() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterClass.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterConstant.kt")
+        public void testAtLastAfterConstant() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterConstant.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterFunInIf.kt")
+        public void testAtLastAfterFunInIf() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterFunInIf.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterFunInTry.kt")
+        public void testAtLastAfterFunInTry() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterFunInTry.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterFunInWhen.kt")
+        public void testAtLastAfterFunInWhen() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterFunInWhen.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterIntFunctionCall.kt")
+        public void testAtLastAfterIntFunctionCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterIntFunctionCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterUnit.kt")
+        public void testAtLastAfterUnit() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterUnit.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterUnitFunctionCall.kt")
+        public void testAtLastAfterUnitFunctionCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterUnitFunctionCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("atLastAfterVal.kt")
+        public void testAtLastAfterVal() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/atLastAfterVal.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("lambda.kt")
         public void testLambda() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/redundantUnitExpression/lambda.kt");