MoveAssignmentToInitializerIntention to select property type when it can be removed

idea/idea-core/src/org/jetbrains/kotlin/idea/core/Utils.kt

@@ -18,8 +18,11 @@ package org.jetbrains.kotlin.idea.core
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.idea.analysis.computeTypeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getFileTopLevelScope
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.resolve.frontendService
@@ -41,9 +44,12 @@ import org.jetbrains.kotlin.resolve.calls.results.ResolutionStatus
 import org.jetbrains.kotlin.resolve.scopes.JetScope
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.receivers.ThisReceiver
+import org.jetbrains.kotlin.resolve.scopes.utils.asJetScope
 import org.jetbrains.kotlin.types.JetType
 import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.checker.JetTypeChecker
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils
+import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 import java.util.*
 
 public fun Call.mapArgumentsToParameters(targetDescriptor: CallableDescriptor): Map<ValueArgument, ValueParameterDescriptor> {
@@ -169,3 +175,11 @@ private fun expectedType(call: Call, bindingContext: BindingContext): JetType {
     } ?: TypeUtils.NO_EXPECTED_TYPE
 }
 
+fun JetCallableDeclaration.canOmitDeclaredType(initializerOrBodyExpression: JetExpression, canChangeTypeToSubtype: Boolean): Boolean {
+    val declaredType = (resolveToDescriptor() as? CallableDescriptor)?.returnType ?: return false
+    val bindingContext = initializerOrBodyExpression.analyze()
+    val scope = initializerOrBodyExpression.getResolutionScope(bindingContext, initializerOrBodyExpression.getResolutionFacade()).asJetScope()
+    val expressionType = initializerOrBodyExpression.computeTypeInContext(scope) ?: return false
+    if (JetTypeChecker.DEFAULT.equalTypes(expressionType, declaredType)) return true
+    return canChangeTypeToSubtype && expressionType.isSubtypeOf(declaredType)
+}

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

@@ -17,21 +17,18 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.descriptors.CallableDescriptor
-import org.jetbrains.kotlin.idea.analysis.computeTypeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.core.CommentSaver
+import org.jetbrains.kotlin.idea.core.canOmitDeclaredType
+import org.jetbrains.kotlin.idea.core.replaced
 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
+import org.jetbrains.kotlin.utils.addToStdlib.check
-import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 
 public class ConvertToExpressionBodyIntention : JetSelfTargetingOffsetIndependentIntention<JetDeclarationWithBody>(
         javaClass(), "Convert to expression body"
@@ -46,24 +43,22 @@ public class ConvertToExpressionBodyIntention : JetSelfTargetingOffsetIndependen
     override fun allowCaretInsideElement(element: PsiElement) = element !is JetDeclaration
 
     override fun applyTo(element: JetDeclarationWithBody, editor: Editor) {
-        applyToInternal(element) {
+        applyTo(element) {
-            val typeRef = it.getTypeReference()!!
+            val typeRef = it.typeReference!!
-            val colon = it.getColon()!!
+            val colon = it.colon!!
-            val range = TextRange(colon.startOffset, typeRef.endOffset)
+            editor.selectionModel.setSelection(colon.startOffset, typeRef.endOffset)
-            editor.getSelectionModel().setSelection(range.getStartOffset(), range.getEndOffset())
+            editor.caretModel.moveToOffset(typeRef.endOffset)
-            editor.getCaretModel().moveToOffset(range.getEndOffset())
         }
     }
 
     public fun applyTo(declaration: JetDeclarationWithBody, canDeleteTypeRef: Boolean) {
-        applyToInternal(declaration) {
+        val deleteTypeHandler: (JetCallableDeclaration) -> Unit = {
-            if (canDeleteTypeRef) {
+            it.deleteChildRange(it.getColon()!!, it.getTypeReference()!!)
-                it.deleteChildRange(it.getColon()!!, it.getTypeReference()!!)
-            }
         }
+        applyTo(declaration, deleteTypeHandler.check { canDeleteTypeRef })
     }
 
-    private fun applyToInternal(declaration: JetDeclarationWithBody, onFinish: (JetCallableDeclaration) -> Unit) {
+    private fun applyTo(declaration: JetDeclarationWithBody, deleteTypeHandler: ((JetCallableDeclaration) -> Unit)?) {
         val value = calcValue(declaration)!!
 
         if (!declaration.hasDeclaredReturnType() && declaration is JetNamedFunction) {
@@ -73,35 +68,22 @@ public class ConvertToExpressionBodyIntention : JetSelfTargetingOffsetIndependen
             }
         }
 
-        val omitType = declaration.hasDeclaredReturnType() && declaration is JetCallableDeclaration && canOmitType(declaration, value)
-
         val body = declaration.getBodyExpression()!!
 
         val commentSaver = CommentSaver(body)
 
         declaration.addBefore(JetPsiFactory(declaration).createEQ(), body)
-        val newBody = body.replace(value)
+        val newBody = body.replaced(value)
 
         commentSaver.restore(newBody)
 
-        if (omitType) {
+        if (deleteTypeHandler != null && declaration is JetCallableDeclaration) {
-            onFinish(declaration as JetCallableDeclaration)
+            if (declaration.hasDeclaredReturnType() && declaration.canOmitDeclaredType(newBody, canChangeTypeToSubtype = true)) {
+                deleteTypeHandler(declaration)
+            }
         }
     }
 
-    private fun canOmitType(declaration: JetCallableDeclaration, expression: JetExpression): Boolean {
-        // 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()
-                         ?: return false
-
-        val declaredType = (declaration.resolveToDescriptor() as? CallableDescriptor)?.getReturnType() ?: return false
-        val scope = scopeExpression.analyze()[BindingContext.RESOLUTION_SCOPE, scopeExpression] ?: return false
-        val expressionType = expression.computeTypeInContext(scope)
-        return expressionType?.isSubtypeOf(declaredType) ?: false
-    }
-
     private fun calcValue(declaration: JetDeclarationWithBody): JetExpression? {
         if (declaration is JetFunctionLiteral) return null
         val body = declaration.getBodyExpression()

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

@@ -17,15 +17,16 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.ScrollType
+import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiReferenceService
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.idea.core.canOmitDeclaredType
+import org.jetbrains.kotlin.idea.quickfix.moveCaret
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
+import org.jetbrains.kotlin.psi.psiUtil.*
-import org.jetbrains.kotlin.psi.psiUtil.contentRange
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 public class MoveAssignmentToInitializerIntention :
@@ -51,7 +52,7 @@ public class MoveAssignmentToInitializerIntention :
     override fun applyTo(element: JetBinaryExpression, editor: Editor) {
         val property = findTargetProperty(element) ?: return
         val initializer = element.right ?: return
-        property.setInitializer(initializer)
+        val newInitializer = property.setInitializer(initializer)!!
 
         val initializerBlock = element.getStrictParentOfType<JetClassInitializer>()
         element.delete()
@@ -59,7 +60,17 @@ public class MoveAssignmentToInitializerIntention :
             initializerBlock.delete()
         }
 
-        property.initializer?.navigate(true)
+        PsiDocumentManager.getInstance(property.project).doPostponedOperationsAndUnblockDocument(editor.document)
+
+        val typeRef = property.typeReference
+        if (typeRef != null && property.canOmitDeclaredType(newInitializer, canChangeTypeToSubtype = !property.isVar)) {
+            val colon = property.colon!!
+            editor.selectionModel.setSelection(colon.startOffset, typeRef.endOffset)
+            editor.moveCaret(typeRef.endOffset, ScrollType.CENTER)
+        }
+        else {
+            editor.moveCaret(newInitializer.startOffset, ScrollType.CENTER)
+        }
     }
 
     private fun findTargetProperty(expr: JetBinaryExpression): JetProperty? {

idea/testData/intentions/convertToExpressionBody/emptyList.kt

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+fun foo(): List<String> {
+    return emptyList()<caret>
+}

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

@@ -0,0 +1,3 @@
+// WITH_RUNTIME
+
+fun foo(): List<String> = emptyList()<caret>

idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType.kt

@@ -0,0 +1,7 @@
+class A {
+    var a: String?
+
+    init {
+        <caret>a = null
+    }
+}

idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType.kt.after

@@ -0,0 +1,4 @@
+class A {
+    var a: String? = <caret>null
+
+}

idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType2.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+
+class A {
+    var a: List<String>
+
+    init {
+        <caret>a = emptyList()
+    }
+}

idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType2.kt.after

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+
+class A {
+    var a: List<String> = <caret>emptyList()
+
+}

idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType3.kt

@@ -0,0 +1,7 @@
+class A {
+    val a: String?
+
+    init {
+        <caret>a = null
+    }
+}

idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType3.kt.after

@@ -0,0 +1,4 @@
+class A {
+    val a<selection>: String?</selection><caret> = null
+
+}

idea/testData/intentions/moveAssignmentToInitializer/comment.kt.after

@@ -1,5 +1,5 @@
 class A {
-    var a: Int = 1
+    var a<selection>: Int</selection><caret> = 1
 
     init {
         // Initialize a

idea/testData/intentions/moveAssignmentToInitializer/deleteInitBlock.kt.after

@@ -1,4 +1,4 @@
 class A {
-    var a: Int = 1
+    var a<selection>: Int</selection><caret> = 1
 
 }

idea/testData/intentions/moveAssignmentToInitializer/simple.kt.after

@@ -1,5 +1,5 @@
 class A {
-    var a: Int = 1
+    var a<selection>: Int</selection><caret> = 1
     var b: Int
 
     init {

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

@@ -4078,6 +4078,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("emptyList.kt")
+        public void testEmptyList() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/emptyList.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("expressionWithReturns1.kt")
         public void testExpressionWithReturns1() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/expressionWithReturns1.kt");
@@ -5493,6 +5499,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/moveAssignmentToInitializer"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }
 
+        @TestMetadata("cannotRemoveType.kt")
+        public void testCannotRemoveType() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("cannotRemoveType2.kt")
+        public void testCannotRemoveType2() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType2.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("cannotRemoveType3.kt")
+        public void testCannotRemoveType3() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/moveAssignmentToInitializer/cannotRemoveType3.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("comment.kt")
         public void testComment() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/moveAssignmentToInitializer/comment.kt");