if to when intention now detects effectively else branches in subsequent code + performs more accurate comment handling #KT-10750 Fixed

Also #KT-11424 Fixed

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

@@ -18,12 +18,18 @@ package org.jetbrains.kotlin.idea.intentions.branchedTransformations.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 org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.getSubjectToIntroduce
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.introduceSubject
 import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
+import org.jetbrains.kotlin.psi.psiUtil.getNextSiblingIgnoringWhitespaceAndComments
+import org.jetbrains.kotlin.psi.psiUtil.siblings
 import java.util.*
 
 class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpression::class.java, "Replace 'if' with 'when'") {
@@ -32,15 +38,70 @@ class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpres
         return element.ifKeyword.textRange
     }
 
-    override fun applyTo(element: KtIfExpression, editor: Editor?) {
+    private fun canPassThrough(expression: KtExpression?): Boolean =
-        val commentSaver = CommentSaver(element)
+            when (expression) {
+                is KtReturnExpression, is KtThrowExpression ->
+                    false
+                is KtBlockExpression ->
+                    expression.statements.all { canPassThrough(it) }
+                is KtIfExpression ->
+                    canPassThrough(expression.then) || canPassThrough(expression.`else`)
+                else ->
+                    true
+            }
 
+    private fun buildNextBranch(ifExpression: KtIfExpression): KtExpression? {
+        var nextSibling = ifExpression.getNextSiblingIgnoringWhitespaceAndComments() ?: return null
+        return when (nextSibling) {
+            is KtIfExpression ->
+                if (nextSibling.then == null) null else nextSibling
+            else -> {
+                val builder = StringBuilder()
+                while (true) {
+                    builder.append(nextSibling.text)
+                    nextSibling = nextSibling.nextSibling ?: break
+                }
+                KtPsiFactory(ifExpression).createBlock(builder.toString())
+            }
+        }
+    }
+
+    private fun KtExpression?.unwrapBlockIfNeeded(): KtExpression? =
+            (this as? KtBlockExpression)?.statements?.singleOrNull() ?: this
+
+    private fun KtIfExpression.siblingsUpTo(other: KtExpression): List<PsiElement> {
+        val result = ArrayList<PsiElement>()
+        var nextSibling = nextSibling
+        // We delete elements up to the next if (or up to the end of the surrounding block)
+        while (nextSibling != null && nextSibling != other) {
+            // RBRACE closes the surrounding block, so it should not be copied / deleted
+            if (nextSibling !is PsiWhiteSpace && nextSibling.node.elementType != KtTokens.RBRACE) {
+                result.add(nextSibling)
+            }
+            nextSibling = nextSibling.nextSibling
+        }
+        return result
+    }
+
+    private fun BuilderByPattern<*>.appendElseBlock(block: KtExpression?) {
+        appendFixedText("else->")
+        appendExpression(block.unwrapBlockIfNeeded())
+        appendFixedText("\n")
+    }
+
+    override fun applyTo(element: KtIfExpression, editor: Editor?) {
+        val siblings = element.siblings()
+        val commentSaver = CommentSaver(PsiChildRange(element, siblings.last()), saveLineBreaks = true)
+
+        val toDelete = ArrayList<PsiElement>()
         var whenExpression = KtPsiFactory(element).buildExpression {
             appendFixedText("when {\n")
 
-            var ifExpression = element
+            var currentIfExpression = element
+            var baseIfExpressionForSyntheticBranch = currentIfExpression
+            var canPassThrough = false
             while (true) {
-                val condition = ifExpression.condition
+                val condition = currentIfExpression.condition
                 val orBranches = ArrayList<KtExpression>()
                 if (condition != null) {
                     orBranches.addOrBranches(condition)
@@ -50,18 +111,33 @@ class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpres
 
                 appendFixedText("->")
 
-                val thenBranch = ifExpression.then
+                val currentThenBranch = currentIfExpression.then
-                appendExpression(thenBranch)
+                appendExpression(currentThenBranch.unwrapBlockIfNeeded())
                 appendFixedText("\n")
 
-                val elseBranch = ifExpression.`else` ?: break
+                canPassThrough = canPassThrough || canPassThrough(currentThenBranch)
-                if (elseBranch is KtIfExpression) {
+
-                    ifExpression = elseBranch
+                val currentElseBranch = currentIfExpression.`else`
+                if (currentElseBranch == null) {
+                    // Try to build synthetic if / else according to KT-10750
+                    val syntheticElseBranch = if (canPassThrough) break else buildNextBranch(baseIfExpressionForSyntheticBranch) ?: break
+                    toDelete.addAll(baseIfExpressionForSyntheticBranch.siblingsUpTo(syntheticElseBranch))
+                    if (syntheticElseBranch is KtIfExpression) {
+                        baseIfExpressionForSyntheticBranch = syntheticElseBranch
+                        currentIfExpression = syntheticElseBranch
+                        toDelete.add(syntheticElseBranch)
+                    }
+                    else {
+                        appendElseBlock(syntheticElseBranch)
+                        break
+                    }
+                }
+                else if (currentElseBranch is KtIfExpression) {
+                    currentIfExpression = currentElseBranch
                 }
                 else {
-                    appendFixedText("else->")
+                    toDelete.addAll(baseIfExpressionForSyntheticBranch.siblingsUpTo(currentElseBranch))
-                    appendExpression(elseBranch)
+                    appendElseBlock(currentElseBranch)
-                    appendFixedText("\n")
                     break
                 }
             }
@@ -76,6 +152,9 @@ class IfToWhenIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpres
 
         val result = element.replace(whenExpression)
         commentSaver.restore(result)
+        toDelete.forEach {
+            it.delete()
+        }
     }
 
     private fun MutableList<KtExpression>.addOrBranches(expression: KtExpression): List<KtExpression> {

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

@@ -0,0 +1,8 @@
+fun bar(arg: Int): String {
+    <caret>if (arg == 1) return "One"
+    else if (arg == 2) return "Two"
+    if (arg == 0) return "Zero"
+    if (arg == -1) return "Minus One"
+    else if (arg == -2) return "Minus Two"
+    return "Something Complex"
+}

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

@@ -0,0 +1,10 @@
+fun bar(arg: Int): String {
+    when (arg) {
+        1 -> return "One"
+        2 -> return "Two"
+        0 -> return "Zero"
+        -1 -> return "Minus One"
+        -2 -> return "Minus Two"
+        else -> return "Something Complex"
+    }
+}

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

@@ -0,0 +1,6 @@
+fun foo(arg: Int): Int {
+    // 1
+    <caret>if (arg == 0) return 0
+    // 2
+    else return -1
+}

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

@@ -0,0 +1,8 @@
+fun foo(arg: Int): Int {
+    // 1
+    when (arg) {
+        0 -> return 0
+// 2
+        else -> return -1
+    }
+}

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

@@ -0,0 +1,4 @@
+fun foo(arg: Int): Int {
+    <caret>if (arg == 0) return 0 // 1
+    else return -1 // 2
+}

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

@@ -0,0 +1,6 @@
+fun foo(arg: Int): Int {
+    when (arg) {
+        0 -> return 0 // 1
+        else -> return -1 // 2
+    }
+}

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

@@ -0,0 +1,9 @@
+fun bar(arg: Int): Int {
+    <caret>if (arg < 0) {
+        if (arg == -3) return 0
+    }
+    if (arg > 0) {
+        return 1
+    }
+    return -1
+}

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

@@ -0,0 +1,9 @@
+fun bar(arg: Int): Int {
+    when {
+        arg < 0 -> if (arg == -3) return 0
+    }
+    if (arg > 0) {
+        return 1
+    }
+    return -1
+}

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

@@ -0,0 +1,15 @@
+fun foo(arg: Int): Int {
+    <caret>if (arg < 0) {
+        var x = arg + 1
+        x++
+        return x
+    }
+    if (arg > 0) {
+        var y = arg - 1
+        y--
+        return y
+    }
+    var z = arg
+    z *= 2
+    return z
+}

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

@@ -0,0 +1,19 @@
+fun foo(arg: Int): Int {
+    when {
+        arg < 0 -> {
+            var x = arg + 1
+            x++
+            return x
+        }
+        arg > 0 -> {
+            var y = arg - 1
+            y--
+            return y
+        }
+        else -> {
+            var z = arg
+            z *= 2
+            return z
+        }
+    }
+}

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

@@ -0,0 +1,17 @@
+fun foo(arg: Any?): Int {
+    // 1
+    <caret>if (arg is Int) {
+        return arg // 2
+    }
+    // 3
+    if (arg is String) {
+        return 42 // 4
+    }
+    // 5
+    if (arg == null) {
+        // 6
+        return 0
+    }
+    // 7
+    return -1 // 8
+}

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

@@ -0,0 +1,14 @@
+fun foo(arg: Any?): Int {
+    // 1
+    when (arg) {
+        is Int -> return arg // 2
+// 3
+        is String -> return 42 // 4
+// 5
+        null -> // 6
+            return 0
+// 7
+// 8
+        else -> return -1
+    }
+}

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

@@ -0,0 +1,6 @@
+fun bar(arg: Int): Int {
+    <caret>if (arg == 1) return 2
+    if (arg == 2) return 5
+    if (arg == 3) return 9
+    return 13
+}

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

@@ -0,0 +1,8 @@
+fun bar(arg: Int): Int {
+    when (arg) {
+        1 -> return 2
+        2 -> return 5
+        3 -> return 9
+        else -> return 13
+    }
+}

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

@@ -0,0 +1,4 @@
+fun foo(arg: Int) {
+    <caret>if (arg == 0) return
+    if (arg == 1) else return
+}

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

@@ -0,0 +1,6 @@
+fun foo(arg: Int) {
+    when (arg) {
+        0 -> return
+    }
+    if (arg == 1) else return
+}

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

@@ -1495,12 +1495,30 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/ifWhen/ifToWhen"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
                 }
 
+                @TestMetadata("combinedIf.kt")
+                public void testCombinedIf() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/combinedIf.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("comment.kt")
                 public void testComment() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/comment.kt");
                     doTest(fileName);
                 }
 
+                @TestMetadata("ifElseSwallowComments.kt")
+                public void testIfElseSwallowComments() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/ifElseSwallowComments.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifElseSwallowReturnComment.kt")
+                public void testIfElseSwallowReturnComment() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/ifElseSwallowReturnComment.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("ifWithEqualityTests.kt")
                 public void testIfWithEqualityTests() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/ifWithEqualityTests.kt");
@@ -1555,6 +1573,36 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("multipleIfFake.kt")
+                public void testMultipleIfFake() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/multipleIfFake.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("multipleIfWithLongBranches.kt")
+                public void testMultipleIfWithLongBranches() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/multipleIfWithLongBranches.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("multipleIfWithReturns.kt")
+                public void testMultipleIfWithReturns() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/multipleIfWithReturns.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("multipleIfWithSingleReturns.kt")
+                public void testMultipleIfWithSingleReturns() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/multipleIfWithSingleReturns.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("secondIfNoThen.kt")
+                public void testSecondIfNoThen() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/secondIfNoThen.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("whenWithMultipleConditionTypes.kt")
                 public void testWhenWithMultipleConditionTypes() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifWhen/ifToWhen/whenWithMultipleConditionTypes.kt");