KT-41298 "Remove redundant 'with' call" intention works incorrectly with non-local returns and single-expression functions (#3713)

* Remove redundant 'with' call: remove redundant 'return' keyword

#KT-41298 Fixed

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

@@ -17,6 +17,8 @@ import org.jetbrains.kotlin.idea.core.moveCaret
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
@@ -85,8 +87,22 @@ private class RemoveRedundantWithFix : LocalQuickFix {
 
     override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
         val callExpression = descriptor.psiElement.parent as? KtCallExpression ?: return
-        val lambdaBody = callExpression.valueArguments.getOrNull(1)?.lambdaExpression()?.bodyExpression ?: return
-        val replaced = callExpression.replaced(lambdaBody)
-        replaced.findExistingEditor()?.moveCaret(replaced.startOffset)
+        val lambdaExpression = callExpression.valueArguments.getOrNull(1)?.lambdaExpression() ?: return
+        val lambdaBody = lambdaExpression.bodyExpression ?: return
+        val declaration = callExpression.getStrictParentOfType<KtDeclarationWithBody>()
+        val replaced = if (declaration?.equalsToken != null && KtPsiUtil.deparenthesize(declaration.bodyExpression) == callExpression) {
+            val singleReturnedExpression = (lambdaBody.statements.singleOrNull() as? KtReturnExpression)?.returnedExpression
+            if (singleReturnedExpression != null) {
+                callExpression.replaced(singleReturnedExpression)
+            } else {
+                declaration.equalsToken?.delete()
+                declaration.bodyExpression?.replaced(KtPsiFactory(project).createSingleStatementBlock(lambdaBody))
+            }
+        } else {
+            callExpression.replaced(lambdaBody)
+        }
+        if (replaced != null) {
+            replaced.findExistingEditor()?.moveCaret(replaced.startOffset)
+        }
     }
 }

idea/testData/inspectionsLocal/redundantWith/asInitializer.kt

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+fun test(): Int = <caret>with("") {
+    println()
+    return 42
+}

idea/testData/inspectionsLocal/redundantWith/asInitializer.kt.after

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

idea/testData/inspectionsLocal/redundantWith/asInitializerWithSingleReturn.kt

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+fun foo(s: String): Int =
+    <caret>with("s") {
+        return 42
+    }

idea/testData/inspectionsLocal/redundantWith/asInitializerWithSingleReturn.kt.after

@@ -0,0 +1,3 @@
+// WITH_RUNTIME
+fun foo(s: String): Int =
+    42

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

@@ -9221,6 +9221,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/inspectionsLocal/redundantWith"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), null, true);
         }
 
+        @TestMetadata("asInitializer.kt")
+        public void testAsInitializer() throws Exception {
+            runTest("idea/testData/inspectionsLocal/redundantWith/asInitializer.kt");
+        }
+
+        @TestMetadata("asInitializerWithSingleReturn.kt")
+        public void testAsInitializerWithSingleReturn() throws Exception {
+            runTest("idea/testData/inspectionsLocal/redundantWith/asInitializerWithSingleReturn.kt");
+        }
+
         @TestMetadata("emptyExpressionInReturn.kt")
         public void testEmptyExpressionInReturn() throws Exception {
             runTest("idea/testData/inspectionsLocal/redundantWith/emptyExpressionInReturn.kt");