Create From Usage: Join local "vars" with their initializers

idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt

@@ -58,6 +58,7 @@ import org.jetbrains.jet.lexer.JetTokens
 import org.jetbrains.jet.plugin.quickfix.createFromUsage.createFunction
 import org.jetbrains.jet.plugin.util.application.runWriteAction
 import org.jetbrains.jet.plugin.refactoring.isMultiLine
+import org.jetbrains.jet.plugin.intentions.declarations.DeclarationUtils
 
 private val TYPE_PARAMETER_LIST_VARIABLE_NAME = "typeParameterList"
 private val TEMPLATE_FROM_USAGE_FUNCTION_BODY = "New Kotlin Function Body.kt"
@@ -92,6 +93,7 @@ class TypeCandidate(val theType: JetType, scope: JetScope? = null) {
 
 class CallableBuilderConfiguration(
         val callableInfo: CallableInfo,
+        val originalExpression: JetExpression,
         val currentFile: JetFile,
         val currentEditor: Editor
 )
@@ -241,6 +243,12 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
 
         private fun createDeclarationSkeleton(): JetCallableDeclaration {
             with (config) {
+                val assignmentToReplace =
+                        if (containingElement is JetBlockExpression && (config.callableInfo as? PropertyInfo)?.writable ?: false) {
+                            originalExpression as JetBinaryExpression
+                        }
+                        else null
+
                 val ownerTypeString = if (isExtension) "${receiverTypeCandidate!!.renderedType!!}." else ""
                 val paramList = when (callableInfo.kind) {
                     CallableKind.FUNCTION ->
@@ -248,7 +256,7 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                     CallableKind.PROPERTY ->
                         ""
                 }
-                val returnTypeString = if (isUnit) "" else ": Any"
+                val returnTypeString = if (isUnit || assignmentToReplace != null) "" else ": Any"
                 val header = "$ownerTypeString${callableInfo.name}$paramList$returnTypeString"
 
                 val psiFactory = JetPsiFactory(currentFile)
@@ -260,6 +268,12 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                         psiFactory.createProperty("$valVar $header")
                     }
                 }
+
+                if (assignmentToReplace != null) {
+                    (declaration as JetProperty).setInitializer(assignmentToReplace.getRight())
+                    return assignmentToReplace.replace(declaration) as JetCallableDeclaration
+                }
+
                 val newLine = psiFactory.createNewLine()
 
                 fun prepend(element: PsiElement, elementBeforeStart: PsiElement, skipInitial: Boolean): PsiElement {
@@ -364,8 +378,7 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                 }
             }
 
-            if (!isUnit) {
-                returnTypeExpression!!
+            if (returnTypeExpression != null) {
                 val returnTypeRef = declaration.getReturnTypeRef()
                 if (returnTypeRef != null) {
                     val returnType = returnTypeExpression.getTypeFromSelection(
@@ -430,8 +443,8 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
             func.getBodyExpression()!!.replace(newBodyExpression)
         }
 
-        private fun setupReturnTypeTemplate(builder: TemplateBuilder, declaration: JetCallableDeclaration): TypeExpression {
-            val returnTypeRef = declaration.getReturnTypeRef()!!
+        private fun setupReturnTypeTemplate(builder: TemplateBuilder, declaration: JetCallableDeclaration): TypeExpression? {
+            val returnTypeRef = declaration.getReturnTypeRef() ?: return null
             val returnTypeExpression = TypeExpression(typeCandidates[config.callableInfo.returnTypeInfo]!!)
             builder.replaceElement(returnTypeRef, returnTypeExpression)
             return returnTypeExpression

idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/createFunction/CreateFunctionFromUsageFix.kt

@@ -1,6 +1,5 @@
 package org.jetbrains.jet.plugin.quickfix.createFromUsage.createFunction
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.jet.plugin.quickfix.createFromUsage.CreateFromUsageFixBase
 import org.jetbrains.jet.lang.psi.JetFile
 import com.intellij.openapi.editor.Editor
@@ -14,8 +13,11 @@ import org.jetbrains.jet.plugin.refactoring.chooseContainerElementIfNecessary
 import org.jetbrains.jet.plugin.refactoring.getExtractionContainers
 import org.jetbrains.jet.lang.psi.JetClassBody
 import org.jetbrains.jet.plugin.quickfix.createFromUsage.callableBuilder.*
+import org.jetbrains.jet.lang.psi.JetExpression
 
-public class CreateFunctionFromUsageFix(element: PsiElement, val functionInfo: CallableInfo) : CreateFromUsageFixBase(element) {
+public class CreateFunctionFromUsageFix(
+        originalExpression: JetExpression,
+        val functionInfo: CallableInfo) : CreateFromUsageFixBase(originalExpression) {
     override fun getText(): String {
         val key = when (functionInfo.kind) {
             CallableKind.FUNCTION -> "create.function.from.usage"
@@ -25,7 +27,7 @@ public class CreateFunctionFromUsageFix(element: PsiElement, val functionInfo: C
     }
 
     override fun invoke(project: Project, editor: Editor?, file: JetFile?) {
-        val functionBuilder = CallableBuilderConfiguration(functionInfo, file!!, editor!!).createBuilder()
+        val functionBuilder = CallableBuilderConfiguration(functionInfo, element as JetExpression, file!!, editor!!).createBuilder()
 
         fun runBuilder(placement: CallablePlacement) {
             functionBuilder.placement = placement

idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt

@@ -37,7 +37,8 @@ object CreateLocalVariableActionFactory: JetSingleIntentionActionFactory() {
                 .filter { it is JetBlockExpression || it is JetDeclarationWithBody }
                 .firstOrNull() as? JetElement ?: return null
 
-        val varExpected = refExpr.getAssignmentByLHS() != null
+        val assignment = refExpr.getAssignmentByLHS()
+        val varExpected = assignment != null
         val typeInfo = TypeInfo(
                 refExpr.getExpressionForTypeGuess(),
                 if (varExpected) Variance.INVARIANT else Variance.OUT_VARIANCE
@@ -49,7 +50,7 @@ object CreateLocalVariableActionFactory: JetSingleIntentionActionFactory() {
             override fun getText(): String = JetBundle.message("create.local.variable.from.usage", propertyInfo.name)
 
             override fun invoke(project: Project, editor: Editor?, file: JetFile?) {
-                with (CallableBuilderConfiguration(propertyInfo, file!!, editor!!).createBuilder()) {
+                with (CallableBuilderConfiguration(propertyInfo, assignment ?: refExpr, file!!, editor!!).createBuilder()) {
                     val actualContainer = when (container) {
                         is JetBlockExpression -> container
                         else -> ConvertToBlockBodyAction().convert(container as JetDeclarationWithBody).getBodyExpression()!!

idea/testData/quickfix/createFromUsage/createVariable/localVariable/afterAssignedInFun.kt

@@ -1,7 +1,5 @@
 // "Create local variable 'foo'" "true"
 
 fun test(n: Int) {
-    var foo: Int
-
-    foo = n + 1
+    var foo = n + 1
 }