From 3178dee759151dc518467f9e2308dc562065b950 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Thu, 21 Dec 2017 17:09:25 +0100
Subject: [PATCH] Insert line break before elvis when converting if to elvis

This ensures correct indentation
---
 .../FoldInitializerAndIfToElvisIntention.kt   | 19 ++++++++++++++++++-
 .../CommentInBlock.kt.after                   |  3 ++-
 .../Comments2.kt.after                        |  3 ++-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/FoldInitializerAndIfToElvisIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/FoldInitializerAndIfToElvisIntention.kt
index 734e2a1f777..a6086f7d117 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/FoldInitializerAndIfToElvisIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/FoldInitializerAndIfToElvisIntention.kt
@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.search.LocalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
@@ -70,7 +72,12 @@ class FoldInitializerAndIfToElvisIntention : SelfTargetingRangeIntention<KtIfExp
         val commentSaver = CommentSaver(childRangeBefore)
         val childRangeAfter = childRangeBefore.withoutLastStatement()
 
-        val elvis = factory.createExpressionByPattern("$0 ?: $1", initializer, ifNullExpr) as KtBinaryExpression
+        val pattern = if (element.then?.hasComments() == true)
+            "$0\n?: $1"
+        else
+            "$0 ?: $1"
+
+        val elvis = factory.createExpressionByPattern(pattern, initializer, ifNullExpr) as KtBinaryExpression
         if (typeReference != null) {
             elvis.left!!.replace(factory.createExpressionByPattern("$0 as? $1", initializer, typeReference))
         }
@@ -86,6 +93,16 @@ class FoldInitializerAndIfToElvisIntention : SelfTargetingRangeIntention<KtIfExp
         editor?.caretModel?.moveToOffset(newElvis.right!!.textOffset)
     }
 
+    private fun PsiElement.hasComments(): Boolean {
+        var result = false
+        accept(object : KtTreeVisitorVoid() {
+            override fun visitComment(comment: PsiComment?) {
+                result = true
+            }
+        })
+        return result
+    }
+
     private data class Data(
             val initializer: KtExpression,
             val declaration: KtVariableDeclaration,
diff --git a/idea/testData/intentions/foldInitializerAndIfToElvis/CommentInBlock.kt.after b/idea/testData/intentions/foldInitializerAndIfToElvis/CommentInBlock.kt.after
index ac0a8592d10..9633f1b8a91 100644
--- a/idea/testData/intentions/foldInitializerAndIfToElvis/CommentInBlock.kt.after
+++ b/idea/testData/intentions/foldInitializerAndIfToElvis/CommentInBlock.kt.after
@@ -1,5 +1,6 @@
 fun foo(p: List<String?>): Int {
-    val v = p[0] ?: // return -1 if null
+    val v = p[0]
+            ?: // return -1 if null
             <caret>return -1
     return v.length
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/foldInitializerAndIfToElvis/Comments2.kt.after b/idea/testData/intentions/foldInitializerAndIfToElvis/Comments2.kt.after
index 876c50fc712..85bc36dd462 100644
--- a/idea/testData/intentions/foldInitializerAndIfToElvis/Comments2.kt.after
+++ b/idea/testData/intentions/foldInitializerAndIfToElvis/Comments2.kt.after
@@ -1,5 +1,6 @@
 fun foo(p: List<String?>): Int {
-    val v = p[0] ?: // v is null
+    val v = p[0]
+            ?: // v is null
             // we should do something with it
             return -1 // let's return -1
     // end of if