Initial implementation of rainbow highlighting for Kotlin

#KT-12629 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/BeforeResolveHighlightingVisitor.kt

@@ -16,11 +16,16 @@
 
 package org.jetbrains.kotlin.idea.highlighter
 
+import com.intellij.codeHighlighting.RainbowHighlighter
 import com.intellij.lang.annotation.AnnotationHolder
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.editor.colors.EditorColorsManager
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.kdoc.parser.KDocKnownTag
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocLink
+import org.jetbrains.kotlin.kdoc.psi.impl.KDocTag
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtExpressionWithLabel
 import org.jetbrains.kotlin.psi.KtLambdaExpression
@@ -34,7 +39,7 @@ internal class BeforeResolveHighlightingVisitor(private val holder: AnnotationHo
     override fun visitElement(element: PsiElement) {
         val elementType = element.node.elementType
         val attributes = when {
-            element is KDocLink -> KotlinHighlightingColors.KDOC_LINK
+            element is KDocLink && !willApplyRainbowHighlight(element) -> KotlinHighlightingColors.KDOC_LINK
 
             elementType in KtTokens.SOFT_KEYWORDS -> {
                 when (elementType) {
@@ -49,6 +54,14 @@ internal class BeforeResolveHighlightingVisitor(private val holder: AnnotationHo
         holder.createInfoAnnotation(element, null).textAttributes = attributes
     }
 
+    private fun willApplyRainbowHighlight(element: KDocLink): Boolean {
+        if (!RainbowHighlighter.isRainbowEnabledWithInheritance(EditorColorsManager.getInstance().globalScheme, KotlinLanguage.INSTANCE)) {
+            return false
+        }
+        // Can't use resolve because it will access indices
+        return (element.parent as? KDocTag)?.knownTag == KDocKnownTag.PARAM
+    }
+
     override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression) {
         if (ApplicationManager.getApplication().isUnitTestMode) return
 

idea/src/META-INF/plugin.xml

@@ -533,6 +533,8 @@
     <annotator language="kotlin" implementationClass="org.jetbrains.kotlin.idea.highlighter.PlatformHeaderAnnotator"/>
     <problemHighlightFilter implementation="org.jetbrains.kotlin.idea.highlighter.KotlinProblemHighlightFilter"/>
 
+    <highlightVisitor implementation="org.jetbrains.kotlin.idea.highlighter.KotlinRainbowVisitor"/>
+
     <annotator language="JAVA" implementationClass="org.jetbrains.kotlin.idea.java.UnimplementedKotlinInterfaceMemberAnnotator"/>
 
     <extendWordSelectionHandler implementation="org.jetbrains.kotlin.idea.editor.wordSelection.KotlinStatementGroupSelectioner"/>

idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinColorSettingsPage.kt

@@ -22,13 +22,15 @@ import com.intellij.openapi.options.OptionsBundle
 import com.intellij.openapi.options.colors.AttributesDescriptor
 import com.intellij.openapi.options.colors.ColorDescriptor
 import com.intellij.openapi.options.colors.ColorSettingsPage
+import com.intellij.openapi.options.colors.RainbowColorSettingsPage
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.KotlinIcons
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import java.lang.reflect.Modifier
 import java.util.*
 
-class KotlinColorSettingsPage : ColorSettingsPage {
+class KotlinColorSettingsPage : ColorSettingsPage, RainbowColorSettingsPage {
+    override fun getLanguage() = KotlinLanguage.INSTANCE
     override fun getIcon() = KotlinIcons.SMALL_LOGO
     override fun getHighlighter(): SyntaxHighlighter = KotlinHighlighter()
 
@@ -162,4 +164,9 @@ var <PACKAGE_PROPERTY><MUTABLE_VARIABLE>globalCounter</MUTABLE_VARIABLE></PACKAG
 
     override fun getColorDescriptors(): Array<ColorDescriptor> = ColorDescriptor.EMPTY_ARRAY
     override fun getDisplayName(): String = KotlinLanguage.NAME
+
+    override fun isRainbowType(type: TextAttributesKey): Boolean {
+        return type == KotlinHighlightingColors.LOCAL_VARIABLE ||
+               type == KotlinHighlightingColors.PARAMETER
+    }
 }

idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowVisitor.kt

@@ -0,0 +1,100 @@
+/*
+ * Copyright 2010-2017 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.highlighter
+
+import com.intellij.codeInsight.daemon.RainbowVisitor
+import com.intellij.openapi.editor.colors.TextAttributesKey
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors.*
+import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.isAncestor
+import org.jetbrains.kotlin.psi.psiUtil.parents
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+import org.jetbrains.kotlin.resolve.bindingContextUtil.getReferenceTargets
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+
+class KotlinRainbowVisitor : RainbowVisitor() {
+    override fun suitableForFile(file: PsiFile) = file is KtFile
+
+    override fun clone() = KotlinRainbowVisitor()
+
+    override fun visit(element: PsiElement) {
+        if (element.isRainbowDeclaration()) {
+            val rainbowElement = (element as KtNamedDeclaration).nameIdentifier ?: return
+            addRainbowHighlight(element, rainbowElement)
+        }
+        else if (element is KtSimpleNameExpression) {
+            val qualifiedExpression = PsiTreeUtil.getParentOfType(element, KtQualifiedExpression::class.java, true,
+                                                                  KtLambdaExpression::class.java, KtValueArgumentList::class.java)
+            if (qualifiedExpression?.selectorExpression?.isAncestor(element) == true) return
+
+            val bindingContext = element.analyze(BodyResolveMode.PARTIAL_WITH_DIAGNOSTICS)
+            val targets = element.getReferenceTargets(bindingContext)
+            val targetVariable = targets.firstIsInstanceOrNull<VariableDescriptor>()
+            if (targetVariable != null) {
+                val targetElement = DescriptorToSourceUtils.getSourceFromDescriptor(targetVariable)
+                if (targetElement.isRainbowDeclaration()) {
+                    addRainbowHighlight(targetElement!!, element)
+                }
+                else if (targetElement == null && element.getReferencedName() == "it") {
+                    addRainbowHighlight(element, element)
+                }
+            }
+        }
+        else if (element is KDocName) {
+            val target = element.reference?.resolve() ?: return
+            if (target.isRainbowDeclaration()) {
+                addRainbowHighlight(target, element, KDOC_LINK)
+            }
+        }
+    }
+
+    private fun addRainbowHighlight(target: PsiElement, rainbowElement: PsiElement,
+                                    attributesKey: TextAttributesKey? = null) {
+        val lambdaSequenceIterator = target.parents
+                .takeWhile { it !is KtDeclaration || it.isAnonymousFunction() || it is KtFunctionLiteral }
+                .filter { it is KtLambdaExpression || it.isAnonymousFunction() }
+                .iterator()
+
+       val attributesKeyToUse = attributesKey ?: (if (target is KtParameter) PARAMETER else LOCAL_VARIABLE)
+        if (lambdaSequenceIterator.hasNext()) {
+            var lambda = lambdaSequenceIterator.next()
+            var lambdaNestingLevel = 0
+            while (lambdaSequenceIterator.hasNext()) {
+                lambdaNestingLevel++
+                lambda = lambdaSequenceIterator.next()
+            }
+            addInfo(getInfo(lambda, rainbowElement, "$lambdaNestingLevel${rainbowElement.text}", attributesKeyToUse))
+            return
+        }
+
+        val context = target.getStrictParentOfType<KtDeclarationWithBody>() ?: return
+        addInfo(getInfo(context, rainbowElement, rainbowElement.text, attributesKeyToUse))
+    }
+
+    private fun PsiElement?.isRainbowDeclaration() =
+            (this is KtProperty && isLocal) || (this is KtParameter && getStrictParentOfType<KtPrimaryConstructor>() == null)
+}
+
+private fun PsiElement.isAnonymousFunction(): Boolean = this is KtNamedFunction && nameIdentifier == null

idea/tests/org/jetbrains/kotlin/idea/highlighter/KotlinRainbowHighlighterTest.kt

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2010-2017 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.highlighter
+
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
+
+class KotlinRainbowHighlighterTest : KotlinLightCodeInsightFixtureTestCase() {
+    override fun getProjectDescriptor() = KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE
+
+    fun testRainbowSimple() {
+        checkRainbow("""
+             fun main(<rainbow color='ff000003'>args</rainbow>: Array<String>) {
+                  val <rainbow color='ff000004'>local1</rainbow> = ""
+                  println(<rainbow color='ff000004'>local1</rainbow> + <rainbow color='ff000003'>args</rainbow>)
+             }
+        """)
+    }
+
+    fun testRainbowNestedIt() {
+        checkRainbow("""
+             fun main(<rainbow color='ff000003'>args</rainbow>: Array<String>) {
+                  listOf("abc", "def").filter { <rainbow color='ff000002'>it</rainbow>.any { <rainbow color='ff000003'>it</rainbow> == 'a' } }
+             }
+        """)
+    }
+
+    fun testRainbowNestedAnonymousFunction() {
+        checkRainbow("""
+            fun main() {
+                listOf("abc", "def").filter(fun(<rainbow color='ff000002'>it</rainbow>): Boolean {
+                    return <rainbow color='ff000002'>it</rainbow>.any(fun(<rainbow color='ff000003'>it</rainbow>): Boolean {
+                        return <rainbow color='ff000003'>it</rainbow> == 'a'
+                    })
+                })
+            }
+        """)
+    }
+
+    fun testKDoc() {
+        checkRainbow("""
+             /**
+              * @param <rainbow color='ff000003'>args</rainbow> foo
+              */
+             fun main(<rainbow color='ff000003'>args</rainbow>: Array<String>) {
+             }
+        """)
+    }
+
+    fun testQualified() {
+        checkRainbow("""
+             data class Foo(val bar: String)
+
+             fun main(<rainbow color='ff000004'>foo</rainbow>: Foo) {
+                  println(<rainbow color='ff000004'>foo</rainbow>.bar)
+                  System.out.println(<rainbow color='ff000004'>foo</rainbow>)
+             }
+        """)
+
+    }
+
+    private fun checkRainbow(code: String) {
+        myFixture.testRainbow("rainbow.kt", code, true, true)
+    }
+}