Extract Function: Generate variables when boxing multiline expressions

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/ExtractableCodeDescriptor.kt

@@ -155,6 +155,8 @@ abstract class OutputValueBoxer(val outputValues: List<OutputValue>) {
 
     protected abstract fun getBoxingExpressionText(arguments: List<String>): String?
 
+    abstract val boxingRequired: Boolean
+
     fun getReturnExpression(arguments: List<String>, psiFactory: JetPsiFactory): JetReturnExpression? {
         return getBoxingExpressionText(arguments)?.let { psiFactory.createReturn(it) }
     }
@@ -208,6 +210,8 @@ abstract class OutputValueBoxer(val outputValues: List<OutputValue>) {
             getType()
         }
 
+        override val boxingRequired: Boolean = outputValues.size > 1
+
         override fun getBoxingExpressionText(arguments: List<String>): String? {
             return when (arguments.size) {
                 0 -> null
@@ -245,6 +249,8 @@ abstract class OutputValueBoxer(val outputValues: List<OutputValue>) {
             )
         }
 
+        override val boxingRequired: Boolean = outputValues.size > 0
+
         override fun getBoxingExpressionText(arguments: List<String>): String? {
             if (arguments.isEmpty()) return null
             return arguments.joinToString(prefix = "kotlin.listOf(", separator = ", ", postfix = ")")

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractorUtil.kt

@@ -43,7 +43,6 @@ import org.jetbrains.jet.lang.psi.psiUtil.prependElement
 import org.jetbrains.jet.lang.psi.psiUtil.appendElement
 import org.jetbrains.jet.lang.psi.psiUtil.replaced
 import org.jetbrains.jet.plugin.intentions.declarations.DeclarationUtils
-import org.jetbrains.jet.lang.psi.psiUtil.getParentByTypesAndPredicate
 import org.jetbrains.jet.lang.psi.JetBlockExpression
 import org.jetbrains.jet.plugin.refactoring.extractFunction.OutputValue.ParameterUpdate
 import org.jetbrains.jet.plugin.refactoring.extractFunction.OutputValue.Jump
@@ -52,6 +51,7 @@ import org.jetbrains.jet.plugin.refactoring.extractFunction.OutputValue.Expressi
 import org.jetbrains.jet.lang.psi.JetReturnExpression
 import org.jetbrains.jet.plugin.refactoring.JetNameValidatorImpl
 import org.jetbrains.jet.plugin.refactoring.JetNameSuggester
+import org.jetbrains.jet.plugin.refactoring.isMultiLine
 
 fun ExtractableCodeDescriptor.getDeclarationText(
         options: ExtractionGeneratorOptions = ExtractionGeneratorOptions.DEFAULT,
@@ -140,7 +140,8 @@ fun ExtractableCodeDescriptor.generateDeclaration(options: ExtractionGeneratorOp
 
     fun replaceWithReturn(
             originalExpression: JetExpression,
-            replacingExpression: JetReturnExpression
+            replacingExpression: JetReturnExpression,
+            expressionToUnifyWith: JetExpression?
     ) {
         val currentResultExpression =
                 if (originalExpression is JetReturnExpression) originalExpression.getReturnedExpression() else originalExpression
@@ -154,7 +155,7 @@ fun ExtractableCodeDescriptor.generateDeclaration(options: ExtractionGeneratorOp
             throw AssertionError("Can' replace '${originalExpression.getText()}' with '${replacingExpression.getText()}'")
         }
 
-        val counterpartMap = createNameCounterpartMap(currentResultExpression, newResultExpression)
+        val counterpartMap = createNameCounterpartMap(currentResultExpression, expressionToUnifyWith ?: newResultExpression)
         nameByOffset.entrySet().forEach { e -> counterpartMap[e.getValue()]?.let { e.setValue(it) } }
     }
 
@@ -225,19 +226,32 @@ fun ExtractableCodeDescriptor.generateDeclaration(options: ExtractionGeneratorOp
             }
         }
 
+        val defaultValue = controlFlow.defaultOutputValue
+
         val lastExpression = body.getStatements().lastOrNull() as? JetExpression
         if (lastExpression is JetReturnExpression) return
 
-        val returnExpression = controlFlow.outputValueBoxer.getReturnExpression(getReturnArguments(lastExpression), psiFactory)
+        val (defaultExpression, expressionToUnifyWith) =
+                if (!options.inTempFile && defaultValue != null && controlFlow.outputValueBoxer.boxingRequired && lastExpression!!.isMultiLine()) {
+                    val varNameValidator = JetNameValidatorImpl(body, lastExpression, JetNameValidatorImpl.Target.PROPERTIES)
+                    val resultVal = JetNameSuggester.suggestNames(defaultValue.valueType, varNameValidator, null).first()
+                    val newDecl = body.addBefore(psiFactory.createDeclaration("val $resultVal = ${lastExpression!!.getText()}"), lastExpression) as JetProperty
+                    body.addBefore(psiFactory.createNewLine(), lastExpression)
+                    psiFactory.createExpression(resultVal) to newDecl.getInitializer()!!
+                }
+                else {
+                    lastExpression to null
+                }
+
+        val returnExpression = controlFlow.outputValueBoxer.getReturnExpression(getReturnArguments(defaultExpression), psiFactory)
         if (returnExpression == null) return
 
-        val defaultValue = controlFlow.defaultOutputValue
         when {
             defaultValue == null ->
                 body.appendElement(returnExpression)
 
             !defaultValue.callSiteReturn ->
-                replaceWithReturn(lastExpression!!, returnExpression)
+                replaceWithReturn(lastExpression!!, returnExpression, expressionToUnifyWith)
         }
     }
 
@@ -386,13 +400,12 @@ fun ExtractableCodeDescriptor.generateDeclaration(options: ExtractionGeneratorOp
         }
     }
 
-    val declaration = createDeclaration()
+    val declaration = createDeclaration().let { if (options.inTempFile) it else insertDeclaration(it) }
     adjustDeclarationBody(declaration)
 
     if (options.inTempFile) return ExtractionResult(declaration, nameByOffset)
 
-    val declarationInPlace = insertDeclaration(declaration)
+    makeCall(declaration)
-    makeCall(declarationInPlace)
+    ShortenReferences.process(declaration)
-    ShortenReferences.process(declarationInPlace)
     return ExtractionResult(declaration, nameByOffset)
 }

idea/src/org/jetbrains/jet/plugin/refactoring/jetRefactoringUtil.kt

@@ -271,18 +271,7 @@ fun PsiElement.getLineCount(): Int {
         return endLine - startLine
     }
 
-    var lineCount = 1
+    return (getText() ?: "").count { it == '\n' } + 1
-    accept(
-            object: JetTreeVisitorVoid() {
-                override fun visitWhiteSpace(space: PsiWhiteSpace) {
-                    if ("\n" in space.getText() ?: "") {
-                        lineCount++
-                    }
-                }
-            }
-    )
-
-    return lineCount
 }
 
 fun PsiElement.isMultiLine(): Boolean = getLineCount() > 1

idea/testData/refactoring/extractFunction/controlFlow/outputValues/nestedNonInlinableCall.kt.after

@@ -17,8 +17,9 @@ fun foo(a: Int): Int {
 
 private fun pair(a: Int, b: Int): Pair<Int, Int> {
     var b1 = b
-    return Pair(if (a > 0) {
+    val i = if (a > 0) {
         b1 += a
         a + 1
-    } else a, b1)
+    } else a
+    return Pair(i, b1)
 }

idea/testData/refactoring/extractFunction/controlFlow/outputValues/outputValueWithExpression.kt.after

@@ -17,10 +17,11 @@ fun foo(a: Int): Int {
 
 private fun pair(a: Int, b: Int): Pair<Int, Int> {
     var b1 = b
-    return Pair(if (a > 0) {
+    val i = if (a > 0) {
         b1 += a
         b1
     } else {
         a
-    }, b1)
+    }
+    return Pair(i, b1)
 }

idea/testData/refactoring/extractFunction/controlFlow/outputValues/outputValuesWithExpression.kt.after

@@ -22,11 +22,12 @@ fun foo(a: Int): Int {
 private fun triple(a: Int, b: Int, c: Int): Triple<Int, Int, Int> {
     var b1 = b
     var c1 = c
-    return Triple(if (a > 0) {
+    val i = if (a > 0) {
         b1 += a
         c1 -= b1
         b1
     } else {
         a
-    }, b1, c1)
+    }
+    return Triple(i, b1, c1)
 }

idea/tests/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/JetExtractionTestGenerated.java

@@ -747,7 +747,8 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 
                 @TestMetadata("outputValueWithSingleLineExpression.kt")
                 public void testOutputValueWithSingleLineExpression() throws Exception {
-                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/controlFlow/outputValues/outputValueWithSingleLineExpression.kt");
+                    String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/controlFlow/outputValues/outputValueWithSingleLineExpression.kt");
+                    doExtractFunctionTest(fileName);
                 }
                 
                 @TestMetadata("outputValuesWithExpression.kt")