From b9865c7225bd6a31b599301287e3e77860d6d7ca Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Fri, 25 Mar 2016 18:50:52 +0300
Subject: [PATCH] Minor

---
 .../SoftKeywordsHighlightingVisitor.kt        | 26 +++++++++++--------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/SoftKeywordsHighlightingVisitor.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/SoftKeywordsHighlightingVisitor.kt
index 3a0b78ec024..24cffb8934b 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/SoftKeywordsHighlightingVisitor.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/SoftKeywordsHighlightingVisitor.kt
@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.idea.highlighter
 import com.intellij.lang.annotation.AnnotationHolder
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.psi.PsiElement
-import com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtVisitorVoid
@@ -27,29 +26,34 @@ import org.jetbrains.kotlin.psi.KtVisitorVoid
 internal class SoftKeywordsHighlightingVisitor(private val holder: AnnotationHolder) : KtVisitorVoid() {
 
     override fun visitElement(element: PsiElement) {
-        if (element is LeafPsiElement) {
-            val elementType = element.elementType
-            if (KtTokens.SOFT_KEYWORDS.contains(elementType)) {
-                var attributes = KotlinHighlightingColors.KEYWORD
-                if (KtTokens.MODIFIER_KEYWORDS.contains(elementType)) {
-                    attributes = KotlinHighlightingColors.BUILTIN_ANNOTATION
+        val elementType = element.node.elementType
+        val attributes = when (elementType) {
+            in KtTokens.SOFT_KEYWORDS -> {
+                when (elementType) {
+                    in KtTokens.MODIFIER_KEYWORDS -> KotlinHighlightingColors.BUILTIN_ANNOTATION
+                    else -> KotlinHighlightingColors.KEYWORD
                 }
-                holder.createInfoAnnotation(element as PsiElement, null).textAttributes = attributes
-            }
-            if (KtTokens.SAFE_ACCESS == elementType) {
-                holder.createInfoAnnotation(element as PsiElement, null).textAttributes = KotlinHighlightingColors.SAFE_ACCESS
             }
+
+            KtTokens.SAFE_ACCESS -> KotlinHighlightingColors.SAFE_ACCESS
+
+            else -> return
         }
+
+        holder.createInfoAnnotation(element, null).textAttributes = attributes
     }
 
     override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression) {
         if (ApplicationManager.getApplication().isUnitTestMode) return
+
         val functionLiteral = lambdaExpression.functionLiteral
         holder.createInfoAnnotation(functionLiteral.lBrace, null).textAttributes = KotlinHighlightingColors.FUNCTION_LITERAL_BRACES_AND_ARROW
+
         val closingBrace = functionLiteral.rBrace
         if (closingBrace != null) {
             holder.createInfoAnnotation(closingBrace, null).textAttributes = KotlinHighlightingColors.FUNCTION_LITERAL_BRACES_AND_ARROW
         }
+
         val arrow = functionLiteral.arrow
         if (arrow != null) {
             holder.createInfoAnnotation(arrow, null).textAttributes = KotlinHighlightingColors.FUNCTION_LITERAL_BRACES_AND_ARROW