Use expression body: add new line more accurately

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 {

idea/testData/refactoring/extractFunction/basic/callableReferenceSelector.kt.after

@@ -10,7 +10,6 @@ class Foo {
             arguments = map(arg)
         }
 
-        private fun map(arg: String) =
+        private fun map(arg: String) = mapOf(Foo::arguments.name to arg)
-                mapOf(Foo::arguments.name to arg)
     }
 }