Search same place for it usages in rainbow highlighter (KT-22242)

Reuse KotlinTargetElementEvaluator for getting same declaration
context for implicit 'it' parameter.

 #KT-22242 Fixed

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

@@ -16,25 +16,26 @@
 
 package org.jetbrains.kotlin.idea.highlighter
 
+import com.intellij.codeInsight.TargetElementUtil
 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.idea.references.mainReference
+import org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinTargetElementEvaluator
 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() {
+    companion object {
+        val KOTLIN_TARGET_ELEMENT_EVALUATOR = KotlinTargetElementEvaluator()
+    }
+
     override fun suitableForFile(file: PsiFile) = file is KtFile
 
     override fun clone() = KotlinRainbowVisitor()
@@ -46,20 +47,24 @@ class KotlinRainbowVisitor : RainbowVisitor() {
                 addRainbowHighlight(element, rainbowElement)
             }
             element is KtSimpleNameExpression -> {
-                val qualifiedExpression = PsiTreeUtil.getParentOfType(element, KtQualifiedExpression::class.java, true,
+                val qualifiedExpression = PsiTreeUtil.getParentOfType(
-                                                                      KtLambdaExpression::class.java, KtValueArgumentList::class.java)
+                    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 reference = element.mainReference
-                val targets = element.getReferenceTargets(bindingContext)
+                val targetElement = reference.resolve()
-                val targetVariable = targets.firstIsInstanceOrNull<VariableDescriptor>()
+                if (targetElement != null) {
-                if (targetVariable != null) {
-                    val targetElement = DescriptorToSourceUtils.getSourceFromDescriptor(targetVariable)
                     if (targetElement.isRainbowDeclaration()) {
-                        addRainbowHighlight(targetElement!!, element)
+                        addRainbowHighlight(targetElement, element)
                     }
-                    else if (targetElement == null && element.getReferencedName() == "it") {
+                } else if (element.getReferencedName() == "it") {
-                        addRainbowHighlight(element, element)
+                    val itTargetElement =
+                        KOTLIN_TARGET_ELEMENT_EVALUATOR.getElementByReference(reference, TargetElementUtil.REFERENCED_ELEMENT_ACCEPTED)
+
+                    if (itTargetElement != null) {
+                        addRainbowHighlight(itTargetElement, element)
                     }
                 }
             }
@@ -72,14 +77,13 @@ class KotlinRainbowVisitor : RainbowVisitor() {
         }
     }
 
-    private fun addRainbowHighlight(target: PsiElement, rainbowElement: PsiElement,
+    private fun addRainbowHighlight(target: PsiElement, rainbowElement: PsiElement, attributesKey: TextAttributesKey? = null) {
-                                    attributesKey: TextAttributesKey? = null) {
         val lambdaSequenceIterator = target.parents
-                .takeWhile { it !is KtDeclaration || it.isAnonymousFunction() || it is KtFunctionLiteral }
+            .takeWhile { it !is KtDeclaration || it.isAnonymousFunction() || it is KtFunctionLiteral }
-                .filter { it is KtLambdaExpression || it.isAnonymousFunction() }
+            .filter { it is KtLambdaExpression || it.isAnonymousFunction() }
-                .iterator()
+            .iterator()
 
-       val attributesKeyToUse = attributesKey ?: (if (target is KtParameter) PARAMETER else LOCAL_VARIABLE)
+        val attributesKeyToUse = attributesKey ?: (if (target is KtParameter) PARAMETER else LOCAL_VARIABLE)
         if (lambdaSequenceIterator.hasNext()) {
             var lambda = lambdaSequenceIterator.next()
             var lambdaNestingLevel = 0
@@ -97,10 +101,10 @@ class KotlinRainbowVisitor : RainbowVisitor() {
         addInfo(getInfo(context, rainbowElement, rainbowElement.text, attributesKeyToUse))
     }
 
-    private fun PsiElement?.isRainbowDeclaration(): Boolean =
+    private fun PsiElement.isRainbowDeclaration(): Boolean =
-            (this is KtProperty && isLocal) ||
+        (this is KtProperty && isLocal) ||
-            (this is KtParameter && getStrictParentOfType<KtPrimaryConstructor>() == null) ||
+                (this is KtParameter && getStrictParentOfType<KtPrimaryConstructor>() == null) ||
-            this is KtDestructuringDeclarationEntry
+                this is KtDestructuringDeclarationEntry
 }
 
 private fun PsiElement.isAnonymousFunction(): Boolean = this is KtNamedFunction && nameIdentifier == null

idea/src/org/jetbrains/kotlin/idea/search/ideaExtensions/KotlinTargetElementEvaluator.kt

@@ -39,8 +39,8 @@ import org.jetbrains.kotlin.resolve.source.getPsi
 
 class KotlinTargetElementEvaluator : TargetElementEvaluatorEx, TargetElementUtilExtender {
     companion object {
-        val DO_NOT_UNWRAP_LABELED_EXPRESSION = 0x100
+        const val DO_NOT_UNWRAP_LABELED_EXPRESSION = 0x100
-        val BYPASS_IMPORT_ALIAS = 0x200
+        const val BYPASS_IMPORT_ALIAS = 0x200
 
         // Place caret after the open curly brace in lambda for generated 'it'
         fun findLambdaOpenLBraceForGeneratedIt(ref: PsiReference): PsiElement? {
@@ -105,8 +105,8 @@ class KotlinTargetElementEvaluator : TargetElementEvaluatorEx, TargetElementUtil
         }
 
         if (BitUtil.isSet(flags, TargetElementUtil.REFERENCED_ELEMENT_ACCEPTED)) {
-            return findLambdaOpenLBraceForGeneratedIt(ref) ?:
+            return findLambdaOpenLBraceForGeneratedIt(ref)
-                   findReceiverForThisInExtensionFunction(ref)
+                    ?: findReceiverForThisInExtensionFunction(ref)
         }
 
         return null

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

@@ -39,6 +39,32 @@ class KotlinRainbowHighlighterTest : KotlinLightCodeInsightFixtureTestCase() {
         """)
     }
 
+    fun testRainbowNestedVal() {
+        checkRainbow(
+            """
+             fun some() {
+                 run {
+                     val <rainbow color='ff000001'>name</rainbow> = 1
+                     <rainbow color='ff000001'>name</rainbow>
+                     run {
+                         val <rainbow color='ff000004'>name</rainbow> = 2
+                         <rainbow color='ff000004'>name</rainbow>
+                     }
+                 }
+             }
+            """
+        )
+    }
+
+    fun testAssignmentIt() {
+        checkRainbow("""
+            val f : (Int) -> Unit = {
+                val <rainbow color='ff000004'>t</rainbow> = <rainbow color='ff000002'>it</rainbow>
+                <rainbow color='ff000002'>it</rainbow>
+            }
+        """)
+    }
+
     fun testRainbowNestedAnonymousFunction() {
         checkRainbow("""
             fun main() {