FunctionWithLambdaExpressionBodyInspection: fix ClassCastException for "Remove braces" quick fix

#KT-32580 Fixed

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

@@ -10,7 +10,6 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.searches.ReferencesSearch
-import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
 import org.jetbrains.kotlin.idea.core.replaced
@@ -20,6 +19,7 @@ import org.jetbrains.kotlin.idea.quickfix.SpecifyTypeExplicitlyFix
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.types.typeUtil.isNothing
 
 class FunctionWithLambdaExpressionBodyInspection : AbstractKotlinInspection() {
@@ -73,7 +73,7 @@ class FunctionWithLambdaExpressionBodyInspection : AbstractKotlinInspection() {
             val lambda = descriptor.psiElement as? KtLambdaExpression ?: return
             val body = lambda.functionLiteral.bodyExpression ?: return
             val replaced = lambda.replaced(body)
-            replaced.parentOfType<KtCallableDeclaration>()?.setTypeIfNeed()
+            replaced.setTypeIfNeed()
         }
     }
 
@@ -86,14 +86,15 @@ class FunctionWithLambdaExpressionBodyInspection : AbstractKotlinInspection() {
             val lambda = descriptor.psiElement as? KtLambdaExpression ?: return
             val body = lambda.functionLiteral.bodyExpression ?: return
             val replaced = lambda.replaced(KtPsiFactory(lambda).createExpressionByPattern("run { $0 }", body))
-            replaced.parentOfType<KtCallableDeclaration>()?.setTypeIfNeed()
+            replaced.setTypeIfNeed()
         }
     }
 }
 
-private fun KtCallableDeclaration.setTypeIfNeed() {
-    val type = (resolveToDescriptorIfAny() as? CallableDescriptor)?.returnType
+private fun KtExpression.setTypeIfNeed() {
+    val declaration = getStrictParentOfType<KtCallableDeclaration>() ?: return
+    val type = (declaration.resolveToDescriptorIfAny() as? CallableDescriptor)?.returnType
     if (type?.isNothing() == true) {
-        this.setType(type)
+        declaration.setType(type)
     }
 }

idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/removeBraces/kt32580.kt

@@ -0,0 +1,6 @@
+// FIX: Remove braces
+class C {
+    fun f4() = {<caret>
+        "single-expression function which returns lambda"
+    }
+}

idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/removeBraces/kt32580.kt.after

@@ -0,0 +1,4 @@
+// FIX: Remove braces
+class C {
+    fun f4() = "single-expression function which returns lambda"
+}

idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/wrapRun/kt32580.kt

@@ -0,0 +1,7 @@
+// FIX: Convert to run { ... }
+// WITH_RUNTIME
+class C {
+    fun f4() = {<caret>
+        "single-expression function which returns lambda"
+    }
+}

idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/wrapRun/kt32580.kt.after

@@ -0,0 +1,5 @@
+// FIX: Convert to run { ... }
+// WITH_RUNTIME
+class C {
+    fun f4() = run { "single-expression function which returns lambda" }
+}

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

@@ -4534,6 +4534,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             public void testGetterReturnsNothing() throws Exception {
                 runTest("idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/removeBraces/getterReturnsNothing.kt");
             }
+
+            @TestMetadata("kt32580.kt")
+            public void testKt32580() throws Exception {
+                runTest("idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/removeBraces/kt32580.kt");
+            }
         }
 
         @TestMetadata("idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/specifyType")
@@ -4590,6 +4595,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             public void testGetterReturnsNothing() throws Exception {
                 runTest("idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/wrapRun/getterReturnsNothing.kt");
             }
+
+            @TestMetadata("kt32580.kt")
+            public void testKt32580() throws Exception {
+                runTest("idea/testData/inspectionsLocal/functionWithLambdaExpressionBody/wrapRun/kt32580.kt");
+            }
         }
     }