Replace if with when: fix behavior on 'else if' #KT-4645 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfToWhenIntention.kt

@@ -121,19 +121,31 @@ class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpres
         appendFixedText("\n")
     }
 
+    private fun KtIfExpression.topmostIfExpression(): KtIfExpression {
+        var target = this
+        while (true) {
+            val container = target.parent as? KtContainerNodeForControlStructureBody ?: break
+            val parent = container.parent as? KtIfExpression ?: break
+            if (parent.`else` != target) break
+            target = parent
+        }
+        return target
+    }
+
     override fun applyTo(element: KtIfExpression, editor: Editor?) {
-        val siblings = element.siblings()
-        val elementCommentSaver = CommentSaver(element)
-        val fullCommentSaver = CommentSaver(PsiChildRange(element, siblings.last()), saveLineBreaks = true)
+        val ifExpression = element.topmostIfExpression()
+        val siblings = ifExpression.siblings()
+        val elementCommentSaver = CommentSaver(ifExpression)
+        val fullCommentSaver = CommentSaver(PsiChildRange(ifExpression, siblings.last()), saveLineBreaks = true)
 
         val toDelete = ArrayList<PsiElement>()
         var applyFullCommentSaver = true
-        val loop = element.getStrictParentOfType<KtLoopExpression>()
+        val loop = ifExpression.getStrictParentOfType<KtLoopExpression>()
         val loopJumpVisitor = LabelLoopJumpVisitor(loop)
-        var whenExpression = KtPsiFactory(element).buildExpression {
+        var whenExpression = KtPsiFactory(ifExpression).buildExpression {
             appendFixedText("when {\n")
 
-            var currentIfExpression = element
+            var currentIfExpression = ifExpression
             var baseIfExpressionForSyntheticBranch = currentIfExpression
             var canPassThrough = false
             while (true) {
@@ -183,7 +195,8 @@ class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpres
             whenExpression = whenExpression.introduceSubject()
         }
 
-        val result = element.replaced(whenExpression)
+        val result = ifExpression.replaced(whenExpression)
+        editor?.caretModel?.moveToOffset(result.startOffset)
 
         (if (applyFullCommentSaver) fullCommentSaver else elementCommentSaver).restore(result)
         toDelete.forEach(PsiElement::delete)

idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf.kt

@@ -0,0 +1,9 @@
+fun test(n: Int): String {
+    return if (n == 0)
+        "zero"
+    else <caret>if (n == 1)
+        "one"
+    else if (n == 2)
+        "two"
+    else "unknown"
+}

idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf.kt.after

@@ -0,0 +1,8 @@
+fun test(n: Int): String {
+    return <caret>when (n) {
+        0 -> "zero"
+        1 -> "one"
+        2 -> "two"
+        else -> "unknown"
+    }
+}

idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf2.kt

@@ -0,0 +1,11 @@
+fun test(n: Int): String {
+    return if (n == 0) {
+        "zero"
+    } else if (n == 1) {
+        "one"
+    } else <caret>if (n == 2) {
+        "two"
+    } else {
+        "unknown"
+    }
+}

idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf2.kt.after

@@ -0,0 +1,8 @@
+fun test(n: Int): String {
+    return <caret>when (n) {
+        0 -> "zero"
+        1 -> "one"
+        2 -> "two"
+        else -> "unknown"
+    }
+}

idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf3.kt

@@ -0,0 +1,13 @@
+fun test(b: Boolean, n: Int): String {
+    return if (b) {
+        if (n == 0)
+            "zero"
+        else if (n == 1)
+            "one"
+        else <caret>if (n == 2)
+            "two"
+        else "unknown"
+    } else {
+        ""
+    }
+}

idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf3.kt.after

@@ -0,0 +1,12 @@
+fun test(b: Boolean, n: Int): String {
+    return if (b) {
+        <caret>when (n) {
+            0 -> "zero"
+            1 -> "one"
+            2 -> "two"
+            else -> "unknown"
+        }
+    } else {
+        ""
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -2659,6 +2659,21 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     runTest("idea/testData/intentions/branched/ifWhen/ifToWhen/multipleIfWithSingleReturns.kt");
                 }
 
+                @TestMetadata("onElseIf.kt")
+                public void testOnElseIf() throws Exception {
+                    runTest("idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf.kt");
+                }
+
+                @TestMetadata("onElseIf2.kt")
+                public void testOnElseIf2() throws Exception {
+                    runTest("idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf2.kt");
+                }
+
+                @TestMetadata("onElseIf3.kt")
+                public void testOnElseIf3() throws Exception {
+                    runTest("idea/testData/intentions/branched/ifWhen/ifToWhen/onElseIf3.kt");
+                }
+
                 @TestMetadata("secondIfNoThen.kt")
                 public void testSecondIfNoThen() throws Exception {
                     runTest("idea/testData/intentions/branched/ifWhen/ifToWhen/secondIfNoThen.kt");