IfNullToElvisIntention - better behavior for var's + slightly refactored utility functions

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

@@ -16,14 +16,13 @@
 
 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.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.types.JetType
 
 public class ConvertToBlockBodyIntention : JetSelfTargetingIntention<JetDeclarationWithBody>(
         javaClass(), "Convert to block body", firstElementOfTypeOnly = true
@@ -33,7 +32,7 @@ public class ConvertToBlockBodyIntention : JetSelfTargetingIntention<JetDeclarat
 
         when (element) {
             is JetNamedFunction -> {
-                val returnType = functionReturnType(element) ?: return false
+                val returnType = element.returnType() ?: return false
                 if (!element.hasDeclaredReturnType() && returnType.isError()) return false// do not convert when type is implicit and unknown
                 return true
             }
@@ -53,7 +52,7 @@ public class ConvertToBlockBodyIntention : JetSelfTargetingIntention<JetDeclarat
             val body = declaration.getBodyExpression()!!
 
             fun generateBody(returnsValue: Boolean): JetExpression {
-                val bodyType = expressionType(body)
+                val bodyType = body.analyze().getType(body)
                 val needReturn = returnsValue &&
                                  (bodyType == null || (!KotlinBuiltIns.isUnit(bodyType) && !KotlinBuiltIns.isNothing(bodyType)))
 
@@ -64,9 +63,9 @@ public class ConvertToBlockBodyIntention : JetSelfTargetingIntention<JetDeclarat
 
             val newBody = when (declaration) {
                 is JetNamedFunction -> {
-                    val returnType = functionReturnType(declaration)!!
+                    val returnType = declaration.returnType()!!
                     if (!declaration.hasDeclaredReturnType() && !KotlinBuiltIns.isUnit(returnType)) {
-                        specifyTypeExplicitly(declaration, returnType)
+                        declaration.setType(returnType)
                     }
                     generateBody(!KotlinBuiltIns.isUnit(returnType) && !KotlinBuiltIns.isNothing(returnType))
                 }
@@ -80,5 +79,10 @@ public class ConvertToBlockBodyIntention : JetSelfTargetingIntention<JetDeclarat
             body.replace(newBody)
             return declaration
         }
+
+        private fun JetNamedFunction.returnType(): JetType? {
+            val descriptor = analyze()[BindingContext.DECLARATION_TO_DESCRIPTOR, this] ?: return null
+            return (descriptor as FunctionDescriptor).getReturnType()
+        }
     }
 }

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

@@ -52,9 +52,9 @@ public class ConvertToExpressionBodyIntention : JetSelfTargetingOffsetIndependen
         val value = calcValue(declaration)!!
 
         if (!declaration.hasDeclaredReturnType() && declaration is JetNamedFunction) {
-            val valueType = expressionType(value)
+            val valueType = value.analyze().getType(value)
             if (valueType == null || !KotlinBuiltIns.isUnit(valueType)) {
-                specifyTypeExplicitly(declaration, "Unit")
+                declaration.setType(KotlinBuiltIns.getInstance().getUnitType())
             }
         }
 
@@ -114,7 +114,7 @@ public class ConvertToExpressionBodyIntention : JetSelfTargetingOffsetIndependen
 
             is JetExpression -> {
                 if (statement is JetBinaryExpression && statement.getOperationToken() == JetTokens.EQ) return null // assignment does not have value
-                val expressionType = expressionType(statement) ?: return null
+                val expressionType = statement.analyze().getType(statement) ?: return null
                 if (!KotlinBuiltIns.isUnit(expressionType) && !KotlinBuiltIns.isNothing(expressionType)) return null
                 return statement
             }

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

@@ -37,21 +37,24 @@ public class IfNullToElvisIntention : JetSelfTargetingOffsetIndependentIntention
         val type = data.ifNullExpression.analyze().getType(data.ifNullExpression) ?: return false
         if (!type.isNothing()) return false
 
-        //TODO: what if implicit var's type will change?
-
         return true
     }
 
     override fun applyTo(element: JetIfExpression, editor: Editor) {
-        val (initializer, statement, ifNullExpr) = calcData(element)!!
+        val (initializer, declaration, ifNullExpr) = calcData(element)!!
         val factory = JetPsiFactory(element)
 
+        val explicitTypeToAdd = if (declaration.isVar() && declaration.getTypeReference() == null)
+            initializer.analyze().getType(initializer)
+        else
+            null
+
         // do not loose any comments!
         val comments = element.extractComments(ifNullExpr)
 
         for (comment in comments) {
-            statement.add(factory.createWhiteSpace())
-            statement.add(comment)
+            declaration.add(factory.createWhiteSpace())
+            declaration.add(comment)
         }
 
         val elvis = factory.createExpression("a ?: b") as JetBinaryExpression
@@ -60,12 +63,16 @@ public class IfNullToElvisIntention : JetSelfTargetingOffsetIndependentIntention
         val newElvis = initializer.replaced(elvis)
         element.delete()
 
+        if (explicitTypeToAdd != null && !explicitTypeToAdd.isError()) {
+            declaration.setType(explicitTypeToAdd)
+        }
+
         editor.getCaretModel().moveToOffset(newElvis.getRight()!!.getTextOffset())
     }
 
     private data class Data(
             val initializer: JetExpression,
-            val statement: JetExpression,
+            val declaration: JetVariableDeclaration,
             val ifNullExpression: JetExpression
     )
 

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

@@ -28,33 +28,11 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.types.JetType
 
-fun specifyTypeExplicitly(declaration: JetNamedFunction, typeText: String) {
-    specifyTypeExplicitly(declaration, JetPsiFactory(declaration).createType(typeText))
-}
-
-fun specifyTypeExplicitly(declaration: JetNamedFunction, type: JetType) {
+fun JetCallableDeclaration.setType(type: JetType) {
     if (type.isError()) return
-    val typeReference = JetPsiFactory(declaration).createType(IdeDescriptorRenderers.SOURCE_CODE.renderType(type))
-    specifyTypeExplicitly(declaration, typeReference)
-    ShortenReferences.DEFAULT.process(declaration.getTypeReference()!!)
-}
-
-fun specifyTypeExplicitly(declaration: JetNamedFunction, typeReference: JetTypeReference) {
-    val anchor = declaration.getValueParameterList() ?: return/*incomplete declaration*/
-    declaration.addAfter(typeReference, anchor)
-    declaration.addAfter(JetPsiFactory(declaration).createColon(), anchor)
-}
-
-fun expressionType(expression: JetExpression): JetType? {
-    val bindingContext = expression.analyze()
-    return bindingContext.getType(expression)
-}
-
-fun functionReturnType(function: JetNamedFunction): JetType? {
-    val bindingContext = function.analyze()
-    val descriptor = bindingContext.get(BindingContext.DECLARATION_TO_DESCRIPTOR, function)
-    if (descriptor == null) return null
-    return (descriptor as FunctionDescriptor).getReturnType()
+    val typeReference = JetPsiFactory(getProject()).createType(IdeDescriptorRenderers.SOURCE_CODE.renderType(type))
+    setTypeReference(typeReference)
+    ShortenReferences.DEFAULT.process(getTypeReference()!!)
 }
 
 fun JetContainerNode.description(): String? {

idea/testData/intentions/ifNullToElvis/Var.kt

@@ -0,0 +1,5 @@
+fun foo(p: List<String?>, b: Boolean) {
+    var v = p[0]
+    <caret>if (v == null) return
+    if (b) v = null
+}

idea/testData/intentions/ifNullToElvis/Var.kt.after

@@ -0,0 +1,4 @@
+fun foo(p: List<String?>, b: Boolean) {
+    var v: String? = p[0] ?: <caret>return
+    if (b) v = null
+}

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

@@ -4173,6 +4173,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/ifNullToElvis/ThrowInBlock.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("Var.kt")
+        public void testVar() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/ifNullToElvis/Var.kt");
+            doTest(fileName);
+        }
     }
 
     @TestMetadata("idea/testData/intentions/insertCurlyBracesToTemplate")