[kotlin] "Redundant 'with' call": fix false negative when with is the return value of a function and lambda has multiple lines

^KTIJ-19621 Fixed

closes https://github.com/JetBrains/intellij-community/pull/2448

GitOrigin-RevId: 738fe2a2156ae66254ec4af9cd956d43eb9b127d

plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/inspections/RedundantWithInspection.kt

@@ -7,23 +7,28 @@ import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElementVisitor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.base.psi.replaced
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.codeinsight.api.classic.inspections.AbstractKotlinInspection
 import org.jetbrains.kotlin.idea.codeinsight.utils.findExistingEditor
 import org.jetbrains.kotlin.idea.core.moveCaret
+import org.jetbrains.kotlin.idea.core.setType
 import org.jetbrains.kotlin.idea.inspections.UnusedLambdaExpressionBodyInspection.Companion.replaceBlockExpressionWithLambdaBody
+import org.jetbrains.kotlin.idea.inspections.collections.isCalling
+import org.jetbrains.kotlin.idea.search.usagesSearch.descriptor
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
+import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.psi2ir.deparenthesize
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
-import org.jetbrains.kotlin.resolve.calls.util.getResolvedCall
-import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 
 class RedundantWithInspection : AbstractKotlinInspection() {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor =
@@ -38,38 +43,25 @@ class RedundantWithInspection : AbstractKotlinInspection() {
             val lambdaBody = lambda.bodyExpression ?: return
 
             val context = callExpression.analyze(BodyResolveMode.PARTIAL_WITH_CFA)
-            if (lambdaBody.statements.size > 1 && callExpression.isUsedAsExpression(context)) return
-            if (callExpression.getResolvedCall(context)?.resultingDescriptor?.fqNameSafe != FqName("kotlin.with")) return
+            if (!callExpression.isCalling(FqName("kotlin.with"), context)) return
 
-            val lambdaDescriptor = context[BindingContext.FUNCTION, lambda.functionLiteral] ?: return
+            if (lambdaBody.statements.size > 1 &&
+                callExpression.isUsedAsExpression(context) &&
+                callExpression.getStrictParentOfType<KtFunction>()?.bodyExpression?.deparenthesize() != callExpression
+            ) return
 
-            var used = false
-            lambda.functionLiteral.acceptChildren(object : KtVisitorVoid() {
-                override fun visitKtElement(element: KtElement) {
-                    if (used) return
-                    element.acceptChildren(this)
-
-                    if (element is KtReturnExpression && element.getLabelName() == "with") {
-                        used = true
-                        return
-                    }
-
-                    if (isUsageOfDescriptor(lambdaDescriptor, element, context)) {
-                        used = true
-                    }
-                }
-            })
+            val functionLiteral = lambda.functionLiteral
+            val lambdaDescriptor = context[BindingContext.FUNCTION, functionLiteral] ?: return
+            val used = functionLiteral.anyDescendantOfType<KtElement> {
+                (it as? KtReturnExpression)?.getLabelName() == "with" || isUsageOfDescriptor(lambdaDescriptor, it, context)
+            }
 
             if (!used) {
                 val quickfix = when (receiver) {
                     is KtSimpleNameExpression, is KtStringTemplateExpression, is KtConstantExpression -> arrayOf(RemoveRedundantWithFix())
                     else -> LocalQuickFix.EMPTY_ARRAY
                 }
-                holder.registerProblem(
-                    callee,
-                    KotlinBundle.message("inspection.redundant.with.display.name"),
-                    *quickfix
-                )
+                holder.registerProblem(callee, KotlinBundle.message("inspection.redundant.with.display.name"), *quickfix)
             }
         })
 }
@@ -88,17 +80,27 @@ private class RemoveRedundantWithFix : LocalQuickFix {
         val lambdaBody = lambdaExpression.bodyExpression ?: return
 
         val function = callExpression.getStrictParentOfType<KtFunction>()
-        val functionBody = KtPsiUtil.deparenthesize(function?.bodyExpression)
+        val functionBody = function?.bodyExpression
 
-        val replaced = if (function?.equalsToken != null && functionBody == callExpression) {
-            val singleStatement = lambdaBody.statements.singleOrNull()?.let {
-                (it as? KtReturnExpression)?.returnedExpression ?: it
-            }
+        val replaced = if (functionBody?.deparenthesize() == callExpression) {
+            val singleStatement = lambdaBody.statements.singleOrNull()
             if (singleStatement != null) {
-                callExpression.replaced(singleStatement)
+                callExpression.replaced(
+                    (singleStatement as? KtReturnExpression)?.returnedExpression ?: singleStatement
+                )
             } else {
+                val returnType = (function.descriptor as? FunctionDescriptor)?.returnType
+                if (returnType != null && !returnType.isUnit()) {
+                    function.setType(returnType, shortenReferences = true)
+                }
+                val lastStatement = lambdaBody.statements.lastOrNull()
+                if (lastStatement != null && lastStatement !is KtReturnExpression) {
+                    lastStatement.replaced(
+                        KtPsiFactory(project).createExpressionByPattern("return $0", lastStatement)
+                    )
+                }
                 function.replaceBlockExpressionWithLambdaBody(lambdaBody)
-                function.bodyExpression
+                functionBody
             }
         } else {
             val result = lambdaBody.allChildren.takeUnless { it.isEmpty }?.let { range ->

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

@@ -10412,24 +10412,29 @@ public abstract class LocalInspectionTestGenerated extends AbstractLocalInspecti
             KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
         }
 
-        @TestMetadata("asInitializer.kt")
-        public void testAsInitializer() throws Exception {
-            runTest("testData/inspectionsLocal/redundantWith/asInitializer.kt");
-        }
-
-        @TestMetadata("asInitializerWithSingleReturn.kt")
-        public void testAsInitializerWithSingleReturn() throws Exception {
-            runTest("testData/inspectionsLocal/redundantWith/asInitializerWithSingleReturn.kt");
-        }
-
         @TestMetadata("emptyExpressionInReturn.kt")
         public void testEmptyExpressionInReturn() throws Exception {
             runTest("testData/inspectionsLocal/redundantWith/emptyExpressionInReturn.kt");
         }
 
-        @TestMetadata("expressionBodyFunction.kt")
-        public void testExpressionBodyFunction() throws Exception {
-            runTest("testData/inspectionsLocal/redundantWith/expressionBodyFunction.kt");
+        @TestMetadata("functionBody.kt")
+        public void testFunctionBody() throws Exception {
+            runTest("testData/inspectionsLocal/redundantWith/functionBody.kt");
+        }
+
+        @TestMetadata("functionBodyWithMultiStatement.kt")
+        public void testFunctionBodyWithMultiStatement() throws Exception {
+            runTest("testData/inspectionsLocal/redundantWith/functionBodyWithMultiStatement.kt");
+        }
+
+        @TestMetadata("functionBodyWithMultiStatementAndReturn.kt")
+        public void testFunctionBodyWithMultiStatementAndReturn() throws Exception {
+            runTest("testData/inspectionsLocal/redundantWith/functionBodyWithMultiStatementAndReturn.kt");
+        }
+
+        @TestMetadata("functionBodyWithReturn.kt")
+        public void testFunctionBodyWithReturn() throws Exception {
+            runTest("testData/inspectionsLocal/redundantWith/functionBodyWithReturn.kt");
         }
 
         @TestMetadata("nested.kt")
@@ -10477,11 +10482,6 @@ public abstract class LocalInspectionTestGenerated extends AbstractLocalInspecti
             runTest("testData/inspectionsLocal/redundantWith/notApplicable_inBinaryExpression.kt");
         }
 
-        @TestMetadata("notApplicable_inFunctionBody.kt")
-        public void testNotApplicable_inFunctionBody() throws Exception {
-            runTest("testData/inspectionsLocal/redundantWith/notApplicable_inFunctionBody.kt");
-        }
-
         @TestMetadata("notApplicable_inProperty.kt")
         public void testNotApplicable_inProperty() throws Exception {
             runTest("testData/inspectionsLocal/redundantWith/notApplicable_inProperty.kt");

plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantWith/functionBodyWithMultiStatement.kt

@@ -1,4 +1,3 @@
-// PROBLEM: none
 // WITH_STDLIB
 fun test() = <caret>with("") {
     println()

plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantWith/functionBodyWithMultiStatement.kt.after

@@ -0,0 +1,5 @@
+// WITH_STDLIB
+fun test(): Int {
+    println()
+    return 1
+}

plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantWith/functionBodyWithMultiStatementAndReturn.kt

@@ -1,6 +1,4 @@
-// PROBLEM: none
 // WITH_STDLIB
-// See KTIJ-19621
 fun test(): Int = <caret>with("") {
     println()
     return 42

plugins/kotlin/idea/tests/testData/inspectionsLocal/redundantWith/functionBodyWithMultiStatementAndReturn.kt.after

@@ -0,0 +1,5 @@
+// WITH_STDLIB
+fun test(): Int {
+    println()
+    return 42
+}