From a960fa78cb2420cf9bfea6490d501169f76512d5 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <Mikhail.Glukhikh@jetbrains.com>
Date: Fri, 7 Jul 2017 17:50:40 +0300
Subject: [PATCH] Use expression body: add new line more accurately

---
 .../inspections/UseExpressionBodyInspection.kt   | 16 ++++++++--------
 .../basic/callableReferenceSelector.kt.after     |  3 +--
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/UseExpressionBodyInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/UseExpressionBodyInspection.kt
index a6e549ad8a5..3d229042e40 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/UseExpressionBodyInspection.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/UseExpressionBodyInspection.kt
@@ -172,6 +172,14 @@ class UseExpressionBodyInspection(private val convertEmptyToUnit: Boolean) : Abs
         declaration.addBefore(factory.createEQ(), body)
         val newBody = body.replaced(value)
 
+        commentSaver.restore(newBody)
+
+        if (deleteTypeHandler != null && declaration is KtCallableDeclaration) {
+            if (declaration.hasDeclaredReturnType() && declaration.canOmitDeclaredType(newBody, canChangeTypeToSubtype = true)) {
+                deleteTypeHandler(declaration)
+            }
+        }
+
         val editor = declaration.findExistingEditor()
         if (editor != null) {
             val startOffset = newBody.startOffset
@@ -182,14 +190,6 @@ class UseExpressionBodyInspection(private val convertEmptyToUnit: Boolean) : Abs
                 declaration.addBefore(factory.createNewLine(), newBody)
             }
         }
-
-        commentSaver.restore(newBody)
-
-        if (deleteTypeHandler != null && declaration is KtCallableDeclaration) {
-            if (declaration.hasDeclaredReturnType() && declaration.canOmitDeclaredType(newBody, canChangeTypeToSubtype = true)) {
-                deleteTypeHandler(declaration)
-            }
-        }
     }
 
     inner class ConvertToExpressionBodyFix : LocalQuickFix {
diff --git a/idea/testData/refactoring/extractFunction/basic/callableReferenceSelector.kt.after b/idea/testData/refactoring/extractFunction/basic/callableReferenceSelector.kt.after
index 010442d23ff..88cf7a7f83b 100644
--- a/idea/testData/refactoring/extractFunction/basic/callableReferenceSelector.kt.after
+++ b/idea/testData/refactoring/extractFunction/basic/callableReferenceSelector.kt.after
@@ -10,7 +10,6 @@ class Foo {
             arguments = map(arg)
         }
 
-        private fun map(arg: String) =
-                mapOf(Foo::arguments.name to arg)
+        private fun map(arg: String) = mapOf(Foo::arguments.name to arg)
     }
 }
\ No newline at end of file