From 2ea2be01cd23ff37ef633eb5d530eddc284c67ed Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Mon, 11 May 2015 15:54:31 +0300
Subject: [PATCH] Minor code refactorings

---
 .../BranchedFoldingUtils.kt                   | 15 +++---
 .../BranchedUnfoldingUtils.kt                 | 46 ++++++++++---------
 2 files changed, 30 insertions(+), 31 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt
index 80f7ad48e6c..1d56d5a08ab 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt
@@ -17,22 +17,19 @@
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
 import com.google.common.base.Predicate
-import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
 
-import java.util.ArrayList
-
 object BranchedFoldingUtils {
     private val CHECK_ASSIGNMENT = object : Predicate<JetElement> {
         override fun apply(input: JetElement): Boolean {
-            if (!JetPsiUtil.isAssignment(input)) return false
+            if (input !is JetBinaryExpression) return false
+            if (input.getOperationToken() !in JetTokens.ALL_ASSIGNMENTS) return false
 
-            val assignment = input as JetBinaryExpression
+            val left = input.getLeft() as? JetSimpleNameExpression ?: return false
+            if (input.getRight() == null) return false
 
-            val left = assignment.getLeft() as? JetSimpleNameExpression ?: return false
-            if (assignment.getRight() == null) return false
-
-            val parent = assignment.getParent()
+            val parent = input.getParent()
             if (parent is JetBlockExpression) {
                 return !JetPsiUtil.checkVariableDeclarationInBlock(parent, left.getText())
             }
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedUnfoldingUtils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedUnfoldingUtils.kt
index d71c544aec1..392c7271269 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedUnfoldingUtils.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedUnfoldingUtils.kt
@@ -17,8 +17,8 @@
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
 import com.intellij.openapi.editor.Editor
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.intentions.declarations.DeclarationUtils
+import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
 
 public object BranchedUnfoldingUtils {
@@ -30,33 +30,35 @@ public object BranchedUnfoldingUtils {
     public fun getUnfoldableExpressionKind(root: JetExpression?): UnfoldableKind? {
         if (root == null) return null
 
-        if (JetPsiUtil.isAssignment(root)) {
-            val assignment = root as JetBinaryExpression
+        when (root) {
+            is JetBinaryExpression -> {
+                if (root.getOperationToken() !in JetTokens.ALL_ASSIGNMENTS) return null
+                if (root.getLeft() == null) return null
 
-            if (assignment.getLeft() == null) return null
-
-            val rhs = assignment.getRight()
-            if (rhs is JetIfExpression) return UnfoldableKind.ASSIGNMENT_TO_IF
-            if (rhs is JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse(rhs)) {
-                return UnfoldableKind.ASSIGNMENT_TO_WHEN
+                val rhs = root.getRight()
+                if (rhs is JetIfExpression) return UnfoldableKind.ASSIGNMENT_TO_IF
+                if (rhs is JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse(rhs)) {
+                    return UnfoldableKind.ASSIGNMENT_TO_WHEN
+                }
             }
-        }
-        else if (root is JetReturnExpression) {
-            val resultExpr = root.getReturnedExpression()
 
-            if (resultExpr is JetIfExpression) return UnfoldableKind.RETURN_TO_IF
-            if (resultExpr is JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse(resultExpr)) {
-                return UnfoldableKind.RETURN_TO_WHEN
+            is JetReturnExpression -> {
+                val resultExpr = root.getReturnedExpression()
+                if (resultExpr is JetIfExpression) return UnfoldableKind.RETURN_TO_IF
+                if (resultExpr is JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse(resultExpr)) {
+                    return UnfoldableKind.RETURN_TO_WHEN
+                }
             }
-        }
-        else if (root is JetProperty) {
-            if (!root.isLocal()) return null
 
-            val initializer = root.getInitializer()
+            is JetProperty -> {
+                if (!root.isLocal()) return null
 
-            if (initializer is JetIfExpression) return UnfoldableKind.PROPERTY_TO_IF
-            if (initializer is JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse(initializer)) {
-                return UnfoldableKind.PROPERTY_TO_WHEN
+                val initializer = root.getInitializer()
+
+                if (initializer is JetIfExpression) return UnfoldableKind.PROPERTY_TO_IF
+                if (initializer is JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse(initializer)) {
+                    return UnfoldableKind.PROPERTY_TO_WHEN
+                }
             }
         }