Convert to Expression Body Intention: Do not remove declaration type if it affects the type of body expression

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

@@ -17,17 +17,22 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.codeInsight.intention.PsiElementBaseIntentionAction
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.editor.Editor
-import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.idea.JetBundle
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.lexer.JetTokens
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analyzer.analyzeInContext
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithVisibility
+import org.jetbrains.kotlin.idea.JetBundle
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.lexer.JetTokens
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 
 public class ConvertToExpressionBodyAction : PsiElementBaseIntentionAction() {
     override fun getFamilyName(): String = JetBundle.message("convert.to.expression.body.action.family.name")
@@ -52,11 +57,15 @@ public class ConvertToExpressionBodyAction : PsiElementBaseIntentionAction() {
             }
         }
 
+        val omitType = declaration.hasDeclaredReturnType() && declaration is JetCallableDeclaration && canOmitType(declaration, value)
+
         val body = declaration.getBodyExpression()!!
         declaration.addBefore(JetPsiFactory(declaration).createEQ(), body)
         body.replace(value)
 
-        if (declaration.hasDeclaredReturnType() && declaration is JetCallableDeclaration && canOmitType(declaration)) {
+        if (omitType) {
+            declaration as JetCallableDeclaration
+
             val typeRef = declaration.getTypeReference()!!
             val colon = declaration.getColon()!!
             if (editor != null) {
@@ -74,10 +83,22 @@ public class ConvertToExpressionBodyAction : PsiElementBaseIntentionAction() {
         invoke(element, editor)
     }
 
-    private fun canOmitType(declaration: JetCallableDeclaration): Boolean {
+    private fun canOmitType(declaration: JetCallableDeclaration, expression: JetExpression): Boolean {
         if (declaration.getModifierList()?.hasModifier(JetTokens.OVERRIDE_KEYWORD) ?: false) return true
+
         val descriptor = declaration.resolveToDescriptor()
-        return !((descriptor as? DeclarationDescriptorWithVisibility)?.getVisibility()?.isPublicAPI() ?: false)
+        if ((descriptor as? DeclarationDescriptorWithVisibility)?.getVisibility()?.isPublicAPI() ?: false) return false
+
+        // Workaround for anonymous objects and similar expressions without resolution scope
+        // TODO: This should probably be fixed in front-end so that resolution scope is recorded for anonymous objects as well
+        val scopeExpression = ((declaration as? JetDeclarationWithBody)?.getBodyExpression() as? JetBlockExpression)
+                                 ?.getStatements()?.singleOrNull() as? JetExpression
+                         ?: return false
+
+        val declaredType = (descriptor as? CallableDescriptor)?.getReturnType() ?: return false
+        val scope = scopeExpression.analyze()[BindingContext.RESOLUTION_SCOPE, scopeExpression] ?: return false
+        val expressionType = expression.analyzeInContext(scope)[BindingContext.EXPRESSION_TYPE, expression] ?: return false
+        return expressionType.isSubtypeOf(declaredType)
     }
 
     private data class Data(val declaration: JetDeclarationWithBody, val value: JetExpression)

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaEmptyArgList.kt.after

@@ -5,4 +5,4 @@ fun foo(t: Array<Int>) {
     t.check(function())
 }
 
-private fun function() = { it + 1 > 1 }
+private fun function(): (Int) -> Boolean = { it + 1 > 1 }

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNoArgList.kt.after

@@ -5,4 +5,4 @@ fun foo(t: Array<Int>) {
     t.check(function())
 }
 
-private fun function() = { it + 1 > 1 }
+private fun function(): (Int) -> Boolean = { it + 1 > 1 }

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNonEmptyArgList.kt.after

@@ -5,4 +5,4 @@ fun foo(t: Array<Int>) {
     t.check(1, 2, function())
 }
 
-private fun function() = { it + 1 > 1 }
+private fun function(): (Int) -> Boolean = { it + 1 > 1 }

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNonEmptyArgListWithNamedArgs.kt.after

@@ -5,4 +5,4 @@ fun foo(t: Array<Int>) {
     t.check(a = 1, b = 2, f = function())
 }
 
-private fun function() = { it + 1 > 1 }
+private fun function(): (Int) -> Boolean = { it + 1 > 1 }