Refactored (and more correct) getOutermostLastBlockElement and its usages

compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiUtil.java

@@ -600,21 +600,6 @@ public class JetPsiUtil {
                ((JetBinaryExpression) element).getOperationToken().equals(JetTokens.EQ);
     }
 
-    @Nullable
-    public static JetElement getOutermostLastBlockElement(@Nullable JetElement element, @NotNull Predicate<JetElement> checkElement) {
-        if (element == null) return null;
-
-        if (!(element instanceof JetBlockExpression)) return checkElement.apply(element) ? element : null;
-
-        JetBlockExpression block = (JetBlockExpression)element;
-        int n = block.getStatements().size();
-
-        if (n == 0) return null;
-
-        JetElement lastElement = block.getStatements().get(n - 1);
-        return checkElement.apply(lastElement) ? lastElement : null;
-    }
-
     public static boolean checkVariableDeclarationInBlock(@NotNull JetBlockExpression block, @NotNull String varName) {
         for (JetElement element : block.getStatements()) {
             if (element instanceof JetVariableDeclaration) {

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/jetPsiUtil.kt

@@ -44,6 +44,7 @@ import com.intellij.openapi.util.TextRange
 import com.intellij.psi.stubs.StubElement
 import org.jetbrains.kotlin.JetNodeTypes
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull
 
 public fun JetCallElement.getCallNameExpression(): JetSimpleNameExpression? {
     val calleeExpression = getCalleeExpression() ?: return null
@@ -135,8 +136,11 @@ public fun <T: PsiElement> T.copied(): T = copy() as T
 public fun JetElement.blockExpressionsOrSingle(): Sequence<JetElement> =
         if (this is JetBlockExpression) getStatements().asSequence() else sequenceOf(this)
 
-public fun JetElement.outermostLastBlockElement(predicate: (JetElement) -> Boolean = { true }): JetElement? {
+public fun JetExpression.lastBlockStatementOrThis(): JetExpression {
-    return JetPsiUtil.getOutermostLastBlockElement(this) { e -> e != null && predicate(e) }
+    return if (this is JetBlockExpression)
+        this.getStatements().lastIsInstanceOrNull<JetExpression>() ?: this
+    else
+        this
 }
 
 public fun JetBlockExpression.appendElement(element: JetElement): JetElement {

compiler/util/src/org/jetbrains/kotlin/utils/addToStdlib.kt

@@ -53,6 +53,24 @@ public inline fun <reified T> Array<*>.firstIsInstance(): T {
     throw NoSuchElementException("No element of given type found")
 }
 
+public inline fun <reified T : Any> Iterable<*>.lastIsInstanceOrNull(): T? {
+    when (this) {
+        is List<*> -> {
+            for (i in this.indices.reversed()) {
+                val element = this[i]
+                if (element is T) return element
+            }
+            return null
+        }
+
+        else -> {
+            return reverse().firstIsInstanceOrNull<T>()
+        }
+    }
+}
+
 public fun <T> streamOfLazyValues(vararg elements: () -> T): Stream<T> = elements.stream().map { it() }
 
-public fun <T1, T2> Pair<T1, T2>.swap(): Pair<T2, T1> = Pair(second, first)
+public fun <T1, T2> Pair<T1, T2>.swap(): Pair<T2, T1> = Pair(second, first)
+
+public fun <T: Any> T.check(predicate: (T) -> Boolean): T? = if (predicate(this)) this else null

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

@@ -16,36 +16,31 @@
 
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
-import com.google.common.base.Predicate
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
+import org.jetbrains.kotlin.utils.addToStdlib.check
 
 object BranchedFoldingUtils {
-    private val CHECK_ASSIGNMENT = object : Predicate<JetElement> {
+    public fun getFoldableBranchedAssignment(branch: JetExpression?): JetBinaryExpression? {
-        override fun apply(input: JetElement): Boolean {
+        fun checkAssignment(expression: JetBinaryExpression): Boolean {
-            if (input !is JetBinaryExpression) return false
+            if (expression.getOperationToken() !in JetTokens.ALL_ASSIGNMENTS) return false
-            if (input.getOperationToken() !in JetTokens.ALL_ASSIGNMENTS) return false
 
-            val left = input.getLeft() as? JetSimpleNameExpression ?: return false
+            val left = expression.getLeft() as? JetSimpleNameExpression ?: return false
-            if (input.getRight() == null) return false
+            if (expression.getRight() == null) return false
 
-            val parent = input.getParent()
+            val parent = expression.getParent()
             if (parent is JetBlockExpression) {
                 return !JetPsiUtil.checkVariableDeclarationInBlock(parent, left.getText())
             }
 
             return true
         }
-    }
+        return (branch?.lastBlockStatementOrThis() as? JetBinaryExpression)?.check(::checkAssignment)
-
-    public fun getFoldableBranchedAssignment(branch: JetExpression?): JetBinaryExpression? {
-        return JetPsiUtil.getOutermostLastBlockElement(branch, CHECK_ASSIGNMENT) as JetBinaryExpression?
     }
 
     public fun getFoldableBranchedReturn(branch: JetExpression?): JetReturnExpression? {
-        return JetPsiUtil.getOutermostLastBlockElement(branch) {
+        return (branch?.lastBlockStatementOrThis() as? JetReturnExpression)?.check { it.getReturnedExpression() != null }
-            (it as? JetReturnExpression)?.getReturnedExpression() != null
-        } as JetReturnExpression?
     }
 
     public fun checkAssignmentsMatch(a1: JetBinaryExpression, a2: JetBinaryExpression): Boolean {

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

@@ -19,12 +19,9 @@ package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 import com.intellij.openapi.editor.Editor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.copied
+import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
 
 public object BranchedUnfoldingUtils {
-    public fun getOutermostLastBlockElement(expression: JetExpression?): JetExpression {
-        return JetPsiUtil.getOutermostLastBlockElement(expression, JetPsiUtil.ANY_JET_ELEMENT) as JetExpression
-    }
-
     public fun unfoldAssignmentToIf(assignment: JetBinaryExpression, editor: Editor) {
         val op = assignment.getOperationReference().getText()
         val left = assignment.getLeft()!!
@@ -32,8 +29,8 @@ public object BranchedUnfoldingUtils {
 
         val newIfExpression = ifExpression.copied()
 
-        val thenExpr = getOutermostLastBlockElement(newIfExpression.getThen())
+        val thenExpr = newIfExpression.getThen()!!.lastBlockStatementOrThis()
-        val elseExpr = getOutermostLastBlockElement(newIfExpression.getElse())
+        val elseExpr = newIfExpression.getElse()!!.lastBlockStatementOrThis()
 
         val psiFactory = JetPsiFactory(assignment)
         thenExpr.replace(psiFactory.createExpressionByPattern("$0 $1 $2", left, op, thenExpr))
@@ -52,7 +49,7 @@ public object BranchedUnfoldingUtils {
         val newWhenExpression = whenExpression.copied()
 
         for (entry in newWhenExpression.getEntries()) {
-            val expr = getOutermostLastBlockElement(entry.getExpression())
+            val expr = entry.getExpression()!!.lastBlockStatementOrThis()
             expr.replace(JetPsiFactory(assignment).createExpressionByPattern("$0 $1 $2", left, op, expr))
         }
 

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

@@ -54,7 +54,7 @@ public class MergeWhenIntention : JetSelfTargetingRangeIntention<JetWhenExpressi
         val names2 = e2.declarationNames()
         if (names1.any { it in names2 }) return false
 
-        return when (e1.getExpression()?.outermostLastBlockElement()) {
+        return when (e1.getExpression()?.lastBlockStatementOrThis()) {
             is JetReturnExpression, is JetThrowExpression, is JetBreakExpression, is JetContinueExpression -> false
             else -> true
         }

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

@@ -33,6 +33,7 @@ public class UnfoldAssignmentToWhenIntention : JetSelfTargetingRangeIntention<Je
         if (element.getLeft() == null) return null
         val right = element.getRight() as? JetWhenExpression ?: return null
         if (!JetPsiUtil.checkWhenExpressionHasSingleElse(right)) return null
+        if (right.getEntries().any { it.getExpression() == null }) return null
         return TextRange(element.startOffset, right.getWhenKeyword().endOffset)
     }
 

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

@@ -32,6 +32,7 @@ public class UnfoldPropertyToWhenIntention : JetSelfTargetingRangeIntention<JetP
         if (!element.isLocal()) return null
         val initializer = element.getInitializer() as? JetWhenExpression ?: return null
         if (!JetPsiUtil.checkWhenExpressionHasSingleElse(initializer)) return null
+        if (initializer.getEntries().any { it.getExpression() == null }) return null
         return TextRange(element.startOffset, initializer.getWhenKeyword().endOffset)
     }
 

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

@@ -19,13 +19,10 @@ package org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.idea.intentions.JetSelfTargetingRangeIntention
-import org.jetbrains.kotlin.idea.intentions.branchedTransformations.BranchedUnfoldingUtils
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.JetIfExpression
-import org.jetbrains.kotlin.psi.JetPsiFactory
-import org.jetbrains.kotlin.psi.JetReturnExpression
-import org.jetbrains.kotlin.psi.createExpressionByPattern
 import org.jetbrains.kotlin.psi.psiUtil.copied
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 public class UnfoldReturnToIfIntention : JetSelfTargetingRangeIntention<JetReturnExpression>(javaClass(), "Replace return with 'if' expression") {
@@ -37,8 +34,8 @@ public class UnfoldReturnToIfIntention : JetSelfTargetingRangeIntention<JetRetur
     override fun applyTo(element: JetReturnExpression, editor: Editor) {
         val ifExpression = element.getReturnedExpression() as JetIfExpression
         val newIfExpression = ifExpression.copied()
-        val thenExpr = BranchedUnfoldingUtils.getOutermostLastBlockElement(newIfExpression.getThen())
+        val thenExpr = newIfExpression.getThen()!!.lastBlockStatementOrThis()
-        val elseExpr = BranchedUnfoldingUtils.getOutermostLastBlockElement(newIfExpression.getElse())
+        val elseExpr = newIfExpression.getElse()!!.lastBlockStatementOrThis()
 
         val psiFactory = JetPsiFactory(element)
         thenExpr.replace(psiFactory.createExpressionByPattern("return $0", thenExpr))

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

@@ -23,12 +23,14 @@ import org.jetbrains.kotlin.idea.intentions.branchedTransformations.BranchedUnfo
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.copied
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 public class UnfoldReturnToWhenIntention : JetSelfTargetingRangeIntention<JetReturnExpression>(javaClass(), "Replace return with 'when' expression") {
     override fun applicabilityRange(element: JetReturnExpression): TextRange? {
         val whenExpr = element.getReturnedExpression() as? JetWhenExpression ?: return null
         if (!JetPsiUtil.checkWhenExpressionHasSingleElse(whenExpr)) return null
+        if (whenExpr.getEntries().any { it.getExpression() == null }) return null
         return TextRange(element.startOffset, whenExpr.getWhenKeyword().endOffset)
     }
 
@@ -37,8 +39,8 @@ public class UnfoldReturnToWhenIntention : JetSelfTargetingRangeIntention<JetRet
         val newWhenExpression = whenExpression.copied()
 
         for (entry in newWhenExpression.getEntries()) {
-            val currExpr = BranchedUnfoldingUtils.getOutermostLastBlockElement(entry.getExpression())
+            val expr = entry.getExpression()!!.lastBlockStatementOrThis()
-            currExpr.replace(JetPsiFactory(element).createExpressionByPattern("return $0", currExpr))
+            expr.replace(JetPsiFactory(element).createExpressionByPattern("return $0", expr))
         }
 
         element.replace(newWhenExpression)