"Redundant Unit" inspection: new tests, simplification & fixes

Related to KT-23977 and KT-24066

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

@@ -11,7 +11,6 @@ import com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
-import org.jetbrains.kotlin.idea.intentions.getCallableDescriptor
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.previousStatement
 import org.jetbrains.kotlin.js.descriptorUtils.nameIfStandardType
 import org.jetbrains.kotlin.psi.*
@@ -19,7 +18,9 @@ import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypesAndPredicate
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
-import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+import org.jetbrains.kotlin.resolve.bindingContextUtil.getTargetFunctionDescriptor
+import org.jetbrains.kotlin.resolve.calls.model.ArgumentMatch
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.isUnit
@@ -43,7 +44,7 @@ private fun KtReferenceExpression.isRedundantUnit(): Boolean {
     if (!isUnitLiteral()) return false
     val parent = this.parent ?: return false
     if (parent is KtReturnExpression) {
-        return parent.returnType()?.nameIfStandardType != KotlinBuiltIns.FQ_NAMES.any.shortName()
+        return parent.expectedReturnType()?.nameIfStandardType != KotlinBuiltIns.FQ_NAMES.any.shortName()
     }
     if (parent is KtBlockExpression) {
         // Do not report just 'Unit' in function literals (return@label Unit is OK even in literals)
@@ -68,17 +69,17 @@ private fun KtReferenceExpression.isRedundantUnit(): Boolean {
 private fun KtExpression.isUnitLiteral(): Boolean =
     KotlinBuiltIns.FQ_NAMES.unit.shortName() == (this as? KtNameReferenceExpression)?.getReferencedNameAsName()
 
-private fun KtReturnExpression.returnType(): KotlinType? {
-    if (labeledExpression != null) {
-        val functionLiteral = getStrictParentOfType<KtFunctionLiteral>() ?: return null
+private fun KtReturnExpression.expectedReturnType(): KotlinType? {
+    val functionDescriptor = getTargetFunctionDescriptor(analyze()) ?: return null
+    val functionLiteral = DescriptorToSourceUtils.descriptorToDeclaration(functionDescriptor) as? KtFunctionLiteral
+    if (functionLiteral != null) {
         val callExpression = functionLiteral.getStrictParentOfType<KtCallExpression>() ?: return null
-        val valueArgument = functionLiteral.getStrictParentOfType<KtValueArgument>()
-        val valueArguments = callExpression.valueArguments
-        val index = if (valueArgument != null) valueArguments.indexOf(valueArgument) else valueArguments.size - 1
-        val parameter = callExpression.getCallableDescriptor()?.valueParameters?.getOrNull(index)
-        return parameter?.returnType?.arguments?.firstOrNull()?.type
+        val resolvedCall = callExpression.resolveToCall() ?: return null
+        val valueArgument = functionLiteral.getStrictParentOfType<KtValueArgument>() ?: return null
+        val mapping = resolvedCall.getArgumentMapping(valueArgument) as? ArgumentMatch ?: return null
+        return mapping.valueParameter.returnType?.arguments?.lastOrNull()?.type
     }
-    return getStrictParentOfType<KtNamedFunction>()?.typeReference?.let { analyze(BodyResolveMode.PARTIAL)[BindingContext.TYPE, it] }
+    return functionDescriptor.returnType
 }
 
 private class RemoveRedundantUnitFix : LocalQuickFix {

idea/testData/inspectionsLocal/redundantUnitExpression/labeledReturnAnyWithParameters.kt

@@ -0,0 +1,9 @@
+// PROBLEM: none
+
+fun foo(f: (Unit, Int, String) -> Any) {}
+
+fun test() {
+    foo { _, _, _ ->
+        return@foo Unit<caret>
+    }
+}

idea/testData/inspectionsLocal/redundantUnitExpression/returnAsNullableAny.kt

@@ -0,0 +1,5 @@
+// PROBLEM: none
+
+fun f(): Any? {
+    return Unit<caret>
+}

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

@@ -3368,6 +3368,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             runTest("idea/testData/inspectionsLocal/redundantUnitExpression/labeledReturnAnyInValueArgument.kt");
         }
 
+        @TestMetadata("labeledReturnAnyWithParameters.kt")
+        public void testLabeledReturnAnyWithParameters() throws Exception {
+            runTest("idea/testData/inspectionsLocal/redundantUnitExpression/labeledReturnAnyWithParameters.kt");
+        }
+
         @TestMetadata("labeledReturnGenericType.kt")
         public void testLabeledReturnGenericType() throws Exception {
             runTest("idea/testData/inspectionsLocal/redundantUnitExpression/labeledReturnGenericType.kt");
@@ -3417,6 +3422,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
         public void testReturnAsAny() throws Exception {
             runTest("idea/testData/inspectionsLocal/redundantUnitExpression/returnAsAny.kt");
         }
+
+        @TestMetadata("returnAsNullableAny.kt")
+        public void testReturnAsNullableAny() throws Exception {
+            runTest("idea/testData/inspectionsLocal/redundantUnitExpression/returnAsNullableAny.kt");
+        }
     }
 
     @TestMetadata("idea/testData/inspectionsLocal/removeRedundantBackticks")