Refactored branched folding intentions into separate files and inlining methods from util

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

@@ -22,7 +22,7 @@ import org.jetbrains.kotlin.psi.*
 
 import java.util.ArrayList
 
-public object BranchedFoldingUtils {
+object BranchedFoldingUtils {
     private val CHECK_ASSIGNMENT = object : Predicate<JetElement> {
         override fun apply(input: JetElement): Boolean {
             if (!JetPsiUtil.isAssignment(input)) return false
@@ -41,175 +41,17 @@ public object BranchedFoldingUtils {
         }
     }
 
-    private fun getFoldableBranchedAssignment(branch: JetExpression?): JetBinaryExpression? {
+    public fun getFoldableBranchedAssignment(branch: JetExpression?): JetBinaryExpression? {
         return JetPsiUtil.getOutermostLastBlockElement(branch, CHECK_ASSIGNMENT) as JetBinaryExpression?
     }
 
-    private fun getFoldableBranchedReturn(branch: JetExpression?): JetReturnExpression? {
+    public fun getFoldableBranchedReturn(branch: JetExpression?): JetReturnExpression? {
         return JetPsiUtil.getOutermostLastBlockElement(branch) {
             (it as? JetReturnExpression)?.getReturnedExpression() != null
         } as JetReturnExpression?
     }
 
-    private fun checkAssignmentsMatch(a1: JetBinaryExpression, a2: JetBinaryExpression): Boolean {
+    public fun checkAssignmentsMatch(a1: JetBinaryExpression, a2: JetBinaryExpression): Boolean {
         return a1.getLeft()?.getText() == a2.getLeft()?.getText() && a1.getOperationToken() == a2.getOperationToken()
     }
-
-    public fun checkFoldableIfExpressionWithAssignments(ifExpression: JetIfExpression): Boolean {
-        val thenBranch = ifExpression.getThen()
-        val elseBranch = ifExpression.getElse()
-
-        val thenAssignment = getFoldableBranchedAssignment(thenBranch)
-        val elseAssignment = getFoldableBranchedAssignment(elseBranch)
-
-        if (thenAssignment == null || elseAssignment == null) return false
-
-        return checkAssignmentsMatch(thenAssignment, elseAssignment)
-    }
-
-    public fun checkFoldableWhenExpressionWithAssignments(whenExpression: JetWhenExpression): Boolean {
-        if (!JetPsiUtil.checkWhenExpressionHasSingleElse(whenExpression)) return false
-
-        val entries = whenExpression.getEntries()
-
-        if (entries.isEmpty()) return false
-
-        val assignments = ArrayList<JetBinaryExpression>()
-        for (entry in entries) {
-            val assignment = getFoldableBranchedAssignment(entry.getExpression()) ?: return false
-            assignments.add(assignment)
-        }
-
-        assert(!assignments.isEmpty())
-
-        val firstAssignment = assignments.get(0)
-        for (assignment in assignments) {
-            if (!checkAssignmentsMatch(assignment, firstAssignment)) return false
-        }
-
-        return true
-    }
-
-    public fun checkFoldableIfExpressionWithReturns(ifExpression: JetIfExpression): Boolean {
-        return getFoldableBranchedReturn(ifExpression.getThen()) != null && getFoldableBranchedReturn(ifExpression.getElse()) != null
-    }
-
-    public fun checkFoldableWhenExpressionWithReturns(whenExpression: JetWhenExpression): Boolean {
-        if (!JetPsiUtil.checkWhenExpressionHasSingleElse(whenExpression)) return false
-
-        val entries = whenExpression.getEntries()
-
-        if (entries.isEmpty()) return false
-
-        for (entry in entries) {
-            if (getFoldableBranchedReturn(entry.getExpression()) == null) return false
-        }
-
-        return true
-    }
-
-    public fun checkFoldableIfExpressionWithAsymmetricReturns(ifExpression: JetIfExpression): Boolean {
-        if (getFoldableBranchedReturn(ifExpression.getThen()) == null || ifExpression.getElse() != null) {
-            return false
-        }
-
-        val nextElement = JetPsiUtil.skipTrailingWhitespacesAndComments(ifExpression)
-        return (nextElement is JetExpression) && getFoldableBranchedReturn(nextElement) != null
-    }
-
-    public fun foldIfExpressionWithAssignments(ifExpression: JetIfExpression) {
-        var thenAssignment = getFoldableBranchedAssignment(ifExpression.getThen()!!)!!
-
-        val op = thenAssignment.getOperationReference().getText()
-        val lhs = thenAssignment.getLeft() as JetSimpleNameExpression
-
-        val assignment = JetPsiFactory(ifExpression).createExpressionByPattern("$0 $1 $2", lhs, op, ifExpression)
-        val newIfExpression = (assignment as JetBinaryExpression).getRight() as JetIfExpression
-
-        thenAssignment = getFoldableBranchedAssignment(newIfExpression.getThen()!!)!!
-        val elseAssignment = getFoldableBranchedAssignment(newIfExpression.getElse()!!)!!
-
-        val thenRhs = thenAssignment.getRight()!!
-        val elseRhs = elseAssignment.getRight()!!
-
-        thenAssignment.replace(thenRhs)
-        elseAssignment.replace(elseRhs)
-
-        ifExpression.replace(assignment)
-    }
-
-    public fun foldIfExpressionWithReturns(ifExpression: JetIfExpression) {
-        val newReturnExpression = JetPsiFactory(ifExpression).createReturn(ifExpression)
-        val newIfExpression = newReturnExpression.getReturnedExpression() as JetIfExpression
-
-        val thenReturn = getFoldableBranchedReturn(newIfExpression.getThen()!!)!!
-        val elseReturn = getFoldableBranchedReturn(newIfExpression.getElse()!!)!!
-
-        val thenExpr = thenReturn.getReturnedExpression()!!
-        val elseExpr = elseReturn.getReturnedExpression()!!
-
-        thenReturn.replace(thenExpr)
-        elseReturn.replace(elseExpr)
-
-        ifExpression.replace(newReturnExpression)
-    }
-
-    public fun foldIfExpressionWithAsymmetricReturns(ifExpression: JetIfExpression) {
-        val condition = ifExpression.getCondition()!!
-        val thenRoot = ifExpression.getThen()!!
-        val elseRoot = JetPsiUtil.skipTrailingWhitespacesAndComments(ifExpression) as JetExpression
-
-        val psiFactory = JetPsiFactory(ifExpression)
-        var newIfExpression = psiFactory.createIf(condition, thenRoot, elseRoot)
-        val newReturnExpression = psiFactory.createReturn(newIfExpression)
-
-        newIfExpression = newReturnExpression.getReturnedExpression() as JetIfExpression
-
-        val thenReturn = getFoldableBranchedReturn(newIfExpression.getThen()!!)!!
-        val elseReturn = getFoldableBranchedReturn(newIfExpression.getElse()!!)!!
-
-        val thenExpr = thenReturn.getReturnedExpression()!!
-        val elseExpr = elseReturn.getReturnedExpression()!!
-
-        thenReturn.replace(thenExpr)
-        elseReturn.replace(elseExpr)
-
-        elseRoot.delete()
-        ifExpression.replace(newReturnExpression)
-    }
-
-    public fun foldWhenExpressionWithAssignments(whenExpression: JetWhenExpression) {
-        assert(!whenExpression.getEntries().isEmpty())
-
-        val firstAssignment = getFoldableBranchedAssignment(whenExpression.getEntries().get(0).getExpression()!!)!!
-
-        val op = firstAssignment.getOperationReference().getText()
-        val lhs = firstAssignment.getLeft() as JetSimpleNameExpression
-
-        val assignment = JetPsiFactory(whenExpression).createExpressionByPattern("$0 $1 $2", lhs, op, whenExpression)
-        val newWhenExpression = (assignment as JetBinaryExpression).getRight() as JetWhenExpression
-
-        for (entry in newWhenExpression.getEntries()) {
-            val currAssignment = getFoldableBranchedAssignment(entry.getExpression()!!)!!
-            val currRhs = currAssignment.getRight()!!
-            currAssignment.replace(currRhs)
-        }
-
-        whenExpression.replace(assignment)
-    }
-
-    public fun foldWhenExpressionWithReturns(whenExpression: JetWhenExpression) {
-        assert(!whenExpression.getEntries().isEmpty())
-
-        val newReturnExpression = JetPsiFactory(whenExpression).createReturn(whenExpression)
-        val newWhenExpression = newReturnExpression.getReturnedExpression() as JetWhenExpression
-
-        for (entry in newWhenExpression.getEntries()) {
-            val currReturn = getFoldableBranchedReturn(entry.getExpression()!!)!!
-            val currExpr = currReturn.getReturnedExpression()!!
-            currReturn.replace(currExpr)
-        }
-
-        whenExpression.replace(newReturnExpression)
-    }
 }

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

@@ -1,75 +0,0 @@
-/*
- * Copyright 2010-2015 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-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.branchedTransformations.BranchedFoldingUtils
-import org.jetbrains.kotlin.psi.JetExpression
-import org.jetbrains.kotlin.psi.JetFile
-import org.jetbrains.kotlin.psi.JetIfExpression
-import org.jetbrains.kotlin.psi.JetWhenExpression
-
-public class FoldIfToAssignmentIntention : JetSelfTargetingOffsetIndependentIntention<JetIfExpression>(javaClass(), "Replace 'if' expression with assignment") {
-    override fun isApplicableTo(element: JetIfExpression): Boolean {
-        return BranchedFoldingUtils.checkFoldableIfExpressionWithAssignments(element)
-    }
-
-    override fun applyTo(element: JetIfExpression, editor: Editor) {
-        BranchedFoldingUtils.foldIfExpressionWithAssignments(element)
-    }
-}
-        
-public class FoldIfToReturnAsymmetricallyIntention : JetSelfTargetingOffsetIndependentIntention<JetIfExpression>(javaClass(), "Replace 'if' expression with return") {
-    override fun isApplicableTo(element: JetIfExpression): Boolean {
-        return BranchedFoldingUtils.checkFoldableIfExpressionWithAsymmetricReturns(element)
-    }
-
-    override fun applyTo(element: JetIfExpression, editor: Editor) {
-        BranchedFoldingUtils.foldIfExpressionWithAsymmetricReturns(element)
-    }
-}
-        
-public class FoldIfToReturnIntention : JetSelfTargetingOffsetIndependentIntention<JetIfExpression>(javaClass(), "Replace 'if' expression with return") {
-    override fun isApplicableTo(element: JetIfExpression): Boolean {
-        return BranchedFoldingUtils.checkFoldableIfExpressionWithReturns(element)
-    }
-
-    override fun applyTo(element: JetIfExpression, editor: Editor) {
-        BranchedFoldingUtils.foldIfExpressionWithReturns(element)
-    }
-}
-        
-public class FoldWhenToAssignmentIntention : JetSelfTargetingOffsetIndependentIntention<JetWhenExpression>(javaClass(), "Replace 'when' expression with assignment") {
-    override fun isApplicableTo(element: JetWhenExpression): Boolean {
-        return BranchedFoldingUtils.checkFoldableWhenExpressionWithAssignments(element)
-    }
-
-    override fun applyTo(element: JetWhenExpression, editor: Editor) {
-        BranchedFoldingUtils.foldWhenExpressionWithAssignments(element)
-    }
-}
-        
-public class FoldWhenToReturnIntention : JetSelfTargetingOffsetIndependentIntention<JetWhenExpression>(javaClass(), "Replace 'when' expression with return") {
-    override fun isApplicableTo(element: JetWhenExpression): Boolean {
-        return BranchedFoldingUtils.checkFoldableWhenExpressionWithReturns(element)
-    }
-
-    override fun applyTo(element: JetWhenExpression, editor: Editor) {
-        BranchedFoldingUtils.foldWhenExpressionWithReturns(element)
-    }
-}

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

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.psi
+import org.jetbrains.kotlin.psi.JetBinaryExpression
+import org.jetbrains.kotlin.psi.JetIfExpression
+import org.jetbrains.kotlin.psi.JetSimpleNameExpression
+import org.jetbrains.kotlin.psi.createExpressionByPattern
+
+public class FoldIfToAssignmentIntention : JetSelfTargetingOffsetIndependentIntention<JetIfExpression>(javaClass(), "Replace 'if' expression with assignment") {
+    override fun isApplicableTo(element: JetIfExpression): Boolean {
+        val thenBranch = element.getThen()
+        val elseBranch = element.getElse()
+
+        val thenAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(thenBranch)
+        val elseAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(elseBranch)
+
+        if (thenAssignment == null || elseAssignment == null) return false
+
+        return BranchedFoldingUtils.checkAssignmentsMatch(thenAssignment, elseAssignment)
+    }
+
+    override fun applyTo(element: JetIfExpression, editor: Editor) {
+        var thenAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.getThen()!!)!!
+
+        val op = thenAssignment.getOperationReference().getText()
+        val lhs = thenAssignment.getLeft() as JetSimpleNameExpression
+
+        val assignment = psi.JetPsiFactory(element).createExpressionByPattern("$0 $1 $2", lhs, op, element)
+        val newIfExpression = (assignment as JetBinaryExpression).getRight() as JetIfExpression
+
+        thenAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(newIfExpression.getThen()!!)!!
+        val elseAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(newIfExpression.getElse()!!)!!
+
+        val thenRhs = thenAssignment.getRight()!!
+        val elseRhs = elseAssignment.getRight()!!
+
+        thenAssignment.replace(thenRhs)
+        elseAssignment.replace(elseRhs)
+
+        element.replace(assignment)
+    }
+}

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

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.psi
+import org.jetbrains.kotlin.psi.JetExpression
+import org.jetbrains.kotlin.psi.JetIfExpression
+import org.jetbrains.kotlin.psi.JetPsiUtil
+
+public class FoldIfToReturnAsymmetricallyIntention : JetSelfTargetingOffsetIndependentIntention<JetIfExpression>(javaClass(), "Replace 'if' expression with return") {
+    override fun isApplicableTo(element: JetIfExpression): Boolean {
+        if (BranchedFoldingUtils.getFoldableBranchedReturn(element.getThen()) == null || element.getElse() != null) {
+            return false
+        }
+
+        val nextElement = JetPsiUtil.skipTrailingWhitespacesAndComments(element)
+        return nextElement is JetExpression && BranchedFoldingUtils.getFoldableBranchedReturn(nextElement) != null
+    }
+
+    override fun applyTo(element: JetIfExpression, editor: Editor) {
+        val condition = element.getCondition()!!
+        val thenRoot = element.getThen()!!
+        val elseRoot = JetPsiUtil.skipTrailingWhitespacesAndComments(element) as JetExpression
+
+        val psiFactory = psi.JetPsiFactory(element)
+        var newIfExpression = psiFactory.createIf(condition, thenRoot, elseRoot)
+        val newReturnExpression = psiFactory.createReturn(newIfExpression)
+
+        newIfExpression = newReturnExpression.getReturnedExpression() as JetIfExpression
+
+        val thenReturn = BranchedFoldingUtils.getFoldableBranchedReturn(newIfExpression.getThen()!!)!!
+        val elseReturn = BranchedFoldingUtils.getFoldableBranchedReturn(newIfExpression.getElse()!!)!!
+
+        val thenExpr = thenReturn.getReturnedExpression()!!
+        val elseExpr = elseReturn.getReturnedExpression()!!
+
+        thenReturn.replace(thenExpr)
+        elseReturn.replace(elseExpr)
+
+        elseRoot.delete()
+        element.replace(newReturnExpression)
+    }
+}

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.psi.JetIfExpression
+import org.jetbrains.kotlin.psi.JetPsiFactory
+
+public class FoldIfToReturnIntention : JetSelfTargetingOffsetIndependentIntention<JetIfExpression>(javaClass(), "Replace 'if' expression with return") {
+    override fun isApplicableTo(element: JetIfExpression): Boolean {
+        return BranchedFoldingUtils.getFoldableBranchedReturn(element.getThen()) != null
+               && BranchedFoldingUtils.getFoldableBranchedReturn(element.getElse()) != null
+    }
+
+    override fun applyTo(element: JetIfExpression, editor: Editor) {
+        val newReturnExpression = JetPsiFactory(element).createReturn(element)
+        val newIfExpression = newReturnExpression.getReturnedExpression() as JetIfExpression
+
+        val thenReturn = BranchedFoldingUtils.getFoldableBranchedReturn(newIfExpression.getThen()!!)!!
+        val elseReturn = BranchedFoldingUtils.getFoldableBranchedReturn(newIfExpression.getElse()!!)!!
+
+        val thenExpr = thenReturn.getReturnedExpression()!!
+        val elseExpr = elseReturn.getReturnedExpression()!!
+
+        thenReturn.replace(thenExpr)
+        elseReturn.replace(elseExpr)
+
+        element.replace(newReturnExpression)
+    }
+}

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

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.psi.*
+import java.util.ArrayList
+
+public class FoldWhenToAssignmentIntention : JetSelfTargetingOffsetIndependentIntention<JetWhenExpression>(javaClass(), "Replace 'when' expression with assignment") {
+    override fun isApplicableTo(element: JetWhenExpression): Boolean {
+        if (!JetPsiUtil.checkWhenExpressionHasSingleElse(element)) return false
+
+        val entries = element.getEntries()
+
+        if (entries.isEmpty()) return false
+
+        val assignments = ArrayList<JetBinaryExpression>()
+        for (entry in entries) {
+            val assignment = BranchedFoldingUtils.getFoldableBranchedAssignment(entry.getExpression()) ?: return false
+            assignments.add(assignment)
+        }
+
+        assert(!assignments.isEmpty())
+
+        val firstAssignment = assignments.get(0)
+        for (assignment in assignments) {
+            if (!BranchedFoldingUtils.checkAssignmentsMatch(assignment, firstAssignment)) return false
+        }
+
+        return true
+    }
+
+    override fun applyTo(element: JetWhenExpression, editor: Editor) {
+        assert(!element.getEntries().isEmpty())
+
+        val firstAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.getEntries().get(0).getExpression()!!)!!
+
+        val op = firstAssignment.getOperationReference().getText()
+        val lhs = firstAssignment.getLeft() as JetSimpleNameExpression
+
+        val assignment = JetPsiFactory(element).createExpressionByPattern("$0 $1 $2", lhs, op, element)
+        val newWhenExpression = (assignment as JetBinaryExpression).getRight() as JetWhenExpression
+
+        for (entry in newWhenExpression.getEntries()) {
+            val currAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(entry.getExpression()!!)!!
+            val currRhs = currAssignment.getRight()!!
+            currAssignment.replace(currRhs)
+        }
+
+        element.replace(assignment)
+    }
+}

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.psi
+import org.jetbrains.kotlin.psi.JetPsiUtil
+import org.jetbrains.kotlin.psi.JetWhenExpression
+
+public class FoldWhenToReturnIntention : JetSelfTargetingOffsetIndependentIntention<JetWhenExpression>(javaClass(), "Replace 'when' expression with return") {
+    override fun isApplicableTo(element: JetWhenExpression): Boolean {
+        if (!JetPsiUtil.checkWhenExpressionHasSingleElse(element)) return false
+
+        val entries = element.getEntries()
+
+        if (entries.isEmpty()) return false
+
+        for (entry in entries) {
+            if (BranchedFoldingUtils.getFoldableBranchedReturn(entry.getExpression()) == null) return false
+        }
+
+        return true
+    }
+
+    override fun applyTo(element: JetWhenExpression, editor: Editor) {
+        assert(!element.getEntries().isEmpty())
+
+        val newReturnExpression = psi.JetPsiFactory(element).createReturn(element)
+        val newWhenExpression = newReturnExpression.getReturnedExpression() as JetWhenExpression
+
+        for (entry in newWhenExpression.getEntries()) {
+            val currReturn = BranchedFoldingUtils.getFoldableBranchedReturn(entry.getExpression()!!)!!
+            val currExpr = currReturn.getReturnedExpression()!!
+            currReturn.replace(currExpr)
+        }
+
+        element.replace(newReturnExpression)
+    }
+}