From af0c4015f3004769e69bbd71b9263a90fcf22e35 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Wed, 27 Mar 2024 16:52:04 +0100
Subject: [PATCH] KTIJ-29230 [kotlin] Do not report `Unit?` as redundant return
 type in K2

`Unit?` return type does not allow empty and
implicit returns, so it cannot be safely removed

^KTIJ-29230 Fixed

GitOrigin-RevId: c9b593e99402cc8b5ce87dcef75fc4a97251b93b
---
 .../declarations/RedundantUnitReturnTypeInspection.kt | 11 ++++++++---
 .../tests/K2LocalInspectionTestGenerated.java         |  5 +++++
 .../inspections/LocalInspectionTestGenerated.java     |  5 +++++
 .../redundantUnitReturnType/nullableUnit.kt           |  4 ++++
 4 files changed, 22 insertions(+), 3 deletions(-)
 create mode 100644 plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantUnitReturnType/nullableUnit.kt

diff --git a/plugins/kotlin/code-insight/inspections-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/inspections/declarations/RedundantUnitReturnTypeInspection.kt b/plugins/kotlin/code-insight/inspections-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/inspections/declarations/RedundantUnitReturnTypeInspection.kt
index 6f36fe70f062..2a6430483aed 100644
--- a/plugins/kotlin/code-insight/inspections-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/inspections/declarations/RedundantUnitReturnTypeInspection.kt
+++ b/plugins/kotlin/code-insight/inspections-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/inspections/declarations/RedundantUnitReturnTypeInspection.kt
@@ -42,9 +42,14 @@ internal class RedundantUnitReturnTypeInspection :
     }
 
     context(KtAnalysisSession)
-    override fun prepareContext(element: KtNamedFunction): TypeInfo? = when {
-        element.getFunctionLikeSymbol().returnType.isUnit -> TypeInfo(TypeInfo.UNIT)
-        else -> null
+    override fun prepareContext(element: KtNamedFunction): TypeInfo? {
+        val returnType = element.getFunctionLikeSymbol().returnType
+
+        if (!returnType.isMarkedNullable && returnType.isUnit) {
+            return TypeInfo(TypeInfo.UNIT)
+        }
+
+        return null
     }
 
     override fun createQuickFix(
diff --git a/plugins/kotlin/code-insight/inspections-k2/tests/test/org/jetbrains/kotlin/idea/k2/inspections/tests/K2LocalInspectionTestGenerated.java b/plugins/kotlin/code-insight/inspections-k2/tests/test/org/jetbrains/kotlin/idea/k2/inspections/tests/K2LocalInspectionTestGenerated.java
index 39b9d26540c2..831fc542cc94 100644
--- a/plugins/kotlin/code-insight/inspections-k2/tests/test/org/jetbrains/kotlin/idea/k2/inspections/tests/K2LocalInspectionTestGenerated.java
+++ b/plugins/kotlin/code-insight/inspections-k2/tests/test/org/jetbrains/kotlin/idea/k2/inspections/tests/K2LocalInspectionTestGenerated.java
@@ -4918,6 +4918,11 @@ public abstract class K2LocalInspectionTestGenerated extends AbstractK2LocalInsp
             runTest("../../../idea/tests/testData/inspectionsLocal/redundantUnitReturnType/expressionBodyUnitType.kt");
         }
 
+        @TestMetadata("nullableUnit.kt")
+        public void testNullableUnit() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/redundantUnitReturnType/nullableUnit.kt");
+        }
+
         @TestMetadata("simple.kt")
         public void testSimple() throws Exception {
             runTest("../../../idea/tests/testData/inspectionsLocal/redundantUnitReturnType/simple.kt");
diff --git a/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java b/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
index fd3b83a2d6d5..51c70751bd13 100644
--- a/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
+++ b/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
@@ -10742,6 +10742,11 @@ public abstract class LocalInspectionTestGenerated extends AbstractLocalInspecti
             runTest("testData/inspectionsLocal/redundantUnitReturnType/expressionBodyUnitType.kt");
         }
 
+        @TestMetadata("nullableUnit.kt")
+        public void testNullableUnit() throws Exception {
+            runTest("testData/inspectionsLocal/redundantUnitReturnType/nullableUnit.kt");
+        }
+
         @TestMetadata("simple.kt")
         public void testSimple() throws Exception {
             runTest("testData/inspectionsLocal/redundantUnitReturnType/simple.kt");
diff --git a/plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantUnitReturnType/nullableUnit.kt b/plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantUnitReturnType/nullableUnit.kt
new file mode 100644
index 000000000000..36e916e81f43
--- /dev/null
+++ b/plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantUnitReturnType/nullableUnit.kt
@@ -0,0 +1,4 @@
+// PROBLEM: none
+fun test(): <caret>Unit? {
+    return null
+}
\ No newline at end of file