KT-9284 Intention to convert to expression body doesn't work in particular case

#KT-9284 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToExpressionBodyIntention.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.core.CommentSaver
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -36,8 +37,10 @@ public class ConvertToExpressionBodyIntention : JetSelfTargetingOffsetIndependen
         javaClass(), "Convert to expression body"
 ) {
     override fun isApplicableTo(element: JetDeclarationWithBody): Boolean {
-        val value = calcValue(element)
+        val value = calcValue(element) ?: return false
-        return value != null && !containsReturn(value)
+        return !value.anyDescendantOfType<JetReturnExpression>(
+                canGoInside = { it !is JetFunctionLiteral && it !is JetNamedFunction && it !is JetPropertyAccessor }
+        )
     }
 
     override fun allowCaretInsideElement(element: PsiElement) = element !is JetDeclaration
@@ -121,18 +124,4 @@ public class ConvertToExpressionBodyIntention : JetSelfTargetingOffsetIndependen
             }
         }
     }
-
-    private fun containsReturn(element: PsiElement): Boolean {
-        if (element is JetReturnExpression) return true
-        //TODO: would be better to have some interface of declaration where return can be used
-        if (element is JetNamedFunction || element is JetPropertyAccessor) return false // can happen inside
-
-        var child = element.getFirstChild()
-        while (child != null) {
-            if (containsReturn(child)) return true
-            child = child.getNextSibling()
-        }
-
-        return false
-    }
 }

idea/testData/intentions/convertToExpressionBody/returnFromLambda.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+
+public fun List<String>.fn() : List<String> {
+    <caret>return map {
+        if (it.isEmpty()) return@map "<empty>"
+        it
+    }
+}

idea/testData/intentions/convertToExpressionBody/returnFromLambda.kt.after

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+
+public fun List<String>.fn() <selection>: List<String></selection><caret> = map {
+    if (it.isEmpty()) return@map "<empty>"
+    it
+}

idea/testData/intentions/convertToExpressionBody/returnFromLambda2.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+
+public fun List<String>.fn() : List<String> {
+    <caret>return map {
+        if (it.isEmpty()) return emptyList()
+        it
+    }
+}

idea/testData/intentions/convertToExpressionBody/returnFromLambda2.kt.after

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+
+public fun List<String>.fn() : List<String> = map {
+    if (it.isEmpty()) return emptyList()
+    it
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -4180,6 +4180,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("returnFromLambda.kt")
+        public void testReturnFromLambda() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/returnFromLambda.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("returnFromLambda2.kt")
+        public void testReturnFromLambda2() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/returnFromLambda2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("returnWithNoValue.kt")
         public void testReturnWithNoValue() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/returnWithNoValue.kt");