FlattenWhenIntention - smaller range and refactoring

idea/idea-analysis/src/org/jetbrains/kotlin/idea/JetBundle.properties

@@ -250,8 +250,6 @@ safe.access.to.if.then=Replace safe access expression with 'if' expression
 safe.access.to.if.then.family=Replace Safe Access Expression with 'if' Expression
 when.to.if=Replace 'when' with 'if'
 when.to.if.family=Replace 'when' with 'if'
-flatten.when=Flatten 'when' expression
-flatten.when.family=Flatten 'when' Expression
 merge.when=Merge 'when' expressions
 merge.when.family=Merge 'when' Expression
 introduce.when.subject=Introduce argument to 'when'

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

@@ -16,17 +16,17 @@
 
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.lexer.JetTokens
-import org.jetbrains.kotlin.psi.JetPsiUnparsingUtils.*
-import org.jetbrains.kotlin.psi.psiUtil.*
-import java.util.ArrayList
-import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.PsiWhiteSpace
-import java.util.Collections
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.util.containers.ContainerUtil
-import org.jetbrains.kotlin.idea.util.psi.patternMatching.toRange
 import org.jetbrains.kotlin.idea.util.psi.patternMatching.matches
+import org.jetbrains.kotlin.idea.util.psi.patternMatching.toRange
+import org.jetbrains.kotlin.lexer.JetTokens
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.JetPsiUnparsingUtils.parenthesizeIfNeeded
+import org.jetbrains.kotlin.psi.JetPsiUnparsingUtils.parenthesizeTextIfNeeded
+import org.jetbrains.kotlin.psi.JetPsiUnparsingUtils.toBinaryExpression
+import org.jetbrains.kotlin.psi.psiUtil.*
 
 public val TRANSFORM_WITHOUT_CHECK: String = "Expression must be checked before applying transformation"
 
@@ -62,18 +62,6 @@ fun JetWhenCondition.toExpressionText(subject: JetExpression?): String {
     }
 }
 
-public fun JetWhenExpression.canFlatten(): Boolean {
-    val subject = getSubjectExpression()
-    if (subject != null && subject !is JetSimpleNameExpression) return false
-
-    if (!JetPsiUtil.checkWhenExpressionHasSingleElse(this)) return false
-
-    val elseBranch = getElseExpression()
-    if (elseBranch !is JetWhenExpression) return false
-
-    return JetPsiUtil.checkWhenExpressionHasSingleElse(elseBranch) && subject.matches(elseBranch.getSubjectExpression())
-}
-
 fun JetWhenExpression.getSubjectCandidate(): JetExpression?  {
     fun JetExpression?.getWhenConditionSubjectCandidate(): JetExpression? {
         return when(this) {
@@ -125,40 +113,6 @@ public fun JetWhenExpression.canIntroduceSubject(): Boolean {
     return getSubjectCandidate() != null
 }
 
-public fun JetWhenExpression.flatten(): JetWhenExpression {
-    val subjectExpression = getSubjectExpression()
-    val elseBranch = getElseExpression()
-
-    assert(elseBranch is JetWhenExpression, TRANSFORM_WITHOUT_CHECK)
-
-    val nestedWhenExpression = elseBranch as JetWhenExpression
-
-    val outerEntries = getEntries()
-    val innerEntries = nestedWhenExpression.getEntries()
-
-    val whenExpression = JetPsiFactory(this).buildExpression {
-        appendFixedText("when")
-        if (subjectExpression != null) {
-            appendFixedText("(").appendExpression(subjectExpression).appendFixedText(")")
-        }
-        appendFixedText("{\n")
-
-        for (entry in outerEntries) {
-            if (entry.isElse()) continue
-            appendNonFormattedText(entry.getText())
-            appendFixedText("\n")
-        }
-        for (entry in innerEntries) {
-            appendNonFormattedText(entry.getText())
-            appendFixedText("\n")
-        }
-
-        appendFixedText("}")
-    } as JetWhenExpression
-
-    return replaced(whenExpression)
-}
-
 public fun JetWhenExpression.introduceSubject(): JetWhenExpression {
     val subject = getSubjectCandidate()!!
 

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

@@ -17,15 +17,59 @@
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions
 
 import com.intellij.openapi.editor.Editor
-import org.jetbrains.kotlin.idea.intentions.JetSelfTargetingOffsetIndependentIntention
+import org.jetbrains.kotlin.idea.intentions.JetSelfTargetingIntention
-import org.jetbrains.kotlin.idea.intentions.branchedTransformations.canFlatten
+import org.jetbrains.kotlin.idea.quickfix.moveCaret
-import org.jetbrains.kotlin.idea.intentions.branchedTransformations.flatten
+import org.jetbrains.kotlin.idea.util.psi.patternMatching.matches
-import org.jetbrains.kotlin.psi.JetWhenExpression
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.replaced
 
-public class FlattenWhenIntention : JetSelfTargetingOffsetIndependentIntention<JetWhenExpression>("flatten.when", javaClass()) {
+public class FlattenWhenIntention : JetSelfTargetingIntention<JetWhenExpression>(javaClass(), "Flatten 'when' expression") {
-    override fun isApplicableTo(element: JetWhenExpression): Boolean = element.canFlatten()
+    override fun isApplicableTo(element: JetWhenExpression, caretOffset: Int): Boolean {
+        val subject = element.getSubjectExpression()
+        if (subject != null && subject !is JetSimpleNameExpression) return false
+
+        if (!JetPsiUtil.checkWhenExpressionHasSingleElse(element)) return false
+
+        val elseEntry = element.getEntries().singleOrNull { it.isElse() } ?: return false
+
+        val innerWhen = elseEntry.getExpression() as? JetWhenExpression ?: return false
+
+        if (!subject.matches(innerWhen.getSubjectExpression())) return false
+        if (!JetPsiUtil.checkWhenExpressionHasSingleElse(innerWhen)) return false
+
+        return elseEntry.getTextRange().getStartOffset() <= caretOffset && caretOffset <= innerWhen.getWhenKeywordElement().getTextRange().getEndOffset()
+    }
 
     override fun applyTo(element: JetWhenExpression, editor: Editor) {
-        element.flatten()
+        val subjectExpression = element.getSubjectExpression()
+        val nestedWhen = element.getElseExpression() as JetWhenExpression
+
+        val outerEntries = element.getEntries()
+        val innerEntries = nestedWhen.getEntries()
+
+        val whenExpression = JetPsiFactory(element).buildExpression {
+            appendFixedText("when")
+            if (subjectExpression != null) {
+                appendFixedText("(").appendExpression(subjectExpression).appendFixedText(")")
+            }
+            appendFixedText("{\n")
+
+            for (entry in outerEntries) {
+                if (entry.isElse()) continue
+                appendNonFormattedText(entry.getText())
+                appendFixedText("\n")
+            }
+            for (entry in innerEntries) {
+                appendNonFormattedText(entry.getText())
+                appendFixedText("\n")
+            }
+
+            appendFixedText("}")
+        } as JetWhenExpression
+
+        val newWhen = element.replaced(whenExpression)
+
+        val firstNewEntry = newWhen.getEntries()[outerEntries.size() - 1]
+        editor.moveCaret(firstNewEntry.getTextOffset())
     }
 }

idea/testData/intentions/branched/when/flatten/flattenWithSubject.kt

@@ -1,8 +1,8 @@
 fun test(n: Int): String {
-    return <caret>when(n) {
+    return when(n) {
         in 0..10 -> "small"
         in 10..100 -> "average"
-        else -> when(n) {
+        <caret>else -> when(n) {
             in 100..1000 -> "big"
             in 1000..10000 -> "very big"
             else -> "unknown"

idea/testData/intentions/branched/when/flatten/flattenWithSubject.kt.after

@@ -1,8 +1,8 @@
 fun test(n: Int): String {
-    return <caret>when (n) {
+    return when (n) {
         in 0..10 -> "small"
         in 10..100 -> "average"
-        in 100..1000 -> "big"
+        <caret>in 100..1000 -> "big"
         in 1000..10000 -> "very big"
         else -> "unknown"
     }

idea/testData/intentions/branched/when/flatten/flattenWithUnmatchedSubjects.kt

@@ -1,9 +1,9 @@
 //IS_APPLICABLE: false
 fun test(n: Int): String {
-    return <caret>when(n) {
+    return when(n) {
         in 0..10 -> "small"
         in 10..100 -> "average"
-        else -> when {
+        <caret>else -> when {
             n in 100..1000 -> "big"
             n in 1000..10000 -> "very big"
             else -> "unknown"

idea/testData/intentions/branched/when/flatten/flattenWithoutSubject.kt

@@ -1,8 +1,8 @@
 fun test(n: Int): String {
-    return <caret>when {
+    return when {
         n in 0..10 -> "small"
         n in 10..100 -> "average"
-        else -> when {
+        else -> when<caret> {
             n in 100..1000 -> "big"
             n in 1000..10000 -> "very big"
             else -> "unknown"

idea/testData/intentions/branched/when/flatten/flattenWithoutSubject.kt.after

@@ -1,8 +1,8 @@
 fun test(n: Int): String {
-    return <caret>when {
+    return when {
         n in 0..10 -> "small"
         n in 10..100 -> "average"
-        n in 100..1000 -> "big"
+        <caret>n in 100..1000 -> "big"
         n in 1000..10000 -> "very big"
         else -> "unknown"
     }