Highlighting: Support implicit 'it' references

#KT-21002 Fixed

idea/src/META-INF/plugin.xml

@@ -662,6 +662,7 @@
     <elementDescriptionProvider
         implementation="org.jetbrains.kotlin.idea.findUsages.KotlinNonCodeSearchElementDescriptionProvider"/>
     <highlightUsagesHandlerFactory implementation="org.jetbrains.kotlin.idea.highlighter.KotlinHighlightExitPointsHandlerFactory"/>
+    <highlightUsagesHandlerFactory implementation="org.jetbrains.kotlin.idea.highlighter.KotlinHighlightImplicitItHandlerFactory"/>
     <findUsagesHandlerFactory implementation="org.jetbrains.kotlin.idea.findUsages.KotlinFindUsagesHandlerFactory"/>
     <usageTypeProvider implementation="org.jetbrains.kotlin.idea.findUsages.KotlinUsageTypeProvider"/>
 

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

@@ -0,0 +1,58 @@
+/*
+ * 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.highlighting.HighlightUsagesHandlerBase
+import com.intellij.codeInsight.highlighting.HighlightUsagesHandlerFactoryBase
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.impl.source.tree.LeafPsiElement
+import com.intellij.util.Consumer
+import org.jetbrains.kotlin.idea.intentions.getLambdaByImplicitItReference
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
+
+class KotlinHighlightImplicitItHandlerFactory : HighlightUsagesHandlerFactoryBase() {
+    override fun createHighlightUsagesHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
+        if (!(target is LeafPsiElement && target.elementType == KtTokens.IDENTIFIER)) return null
+        val refExpr = target.parent as? KtNameReferenceExpression ?: return null
+        val lambda = getLambdaByImplicitItReference(refExpr) ?: return null
+        return object : HighlightUsagesHandlerBase<KtNameReferenceExpression>(editor, file) {
+            override fun getTargets() = listOf(refExpr)
+
+            override fun selectTargets(
+                    targets: MutableList<KtNameReferenceExpression>,
+                    selectionConsumer: Consumer<MutableList<KtNameReferenceExpression>>
+            ) = selectionConsumer.consume(targets)
+
+            override fun computeUsages(targets: MutableList<KtNameReferenceExpression>?) {
+                lambda.accept(
+                        object : KtTreeVisitorVoid() {
+                            override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
+                                if (expression is KtNameReferenceExpression && getLambdaByImplicitItReference(expression) == lambda) {
+                                    addOccurrence(expression)
+                                }
+                            }
+                        }
+                )
+            }
+        }
+    }
+}

idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt

@@ -18,7 +18,6 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.KtNodeTypes
-import org.jetbrains.kotlin.builtins.BuiltInsPackageFragment
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
@@ -28,16 +27,15 @@ import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.core.setType
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.CollectionLiteralResolver
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
-import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameOrNull
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.isFlexible
 import java.lang.IllegalArgumentException
@@ -62,13 +60,18 @@ fun KtCallExpression.isMethodCall(fqMethodName: String): Boolean {
     return resolvedCall.resultingDescriptor.fqNameUnsafe.asString() == fqMethodName
 }
 
-fun isAutoCreatedItUsage(expression: KtNameReferenceExpression): Boolean {
-    if (expression.getReferencedName() != "it") return false
+fun isAutoCreatedItUsage(expression: KtNameReferenceExpression) = resolveToAutoCreatedItDescriptor(expression) != null
+
+fun resolveToAutoCreatedItDescriptor(expression: KtNameReferenceExpression): ValueParameterDescriptor? {
+    if (expression.getReferencedName() != "it") return null
     val context = expression.analyze(BodyResolveMode.PARTIAL)
-    val target = expression.mainReference.resolveToDescriptors(context).singleOrNull() as? ValueParameterDescriptor? ?: return false
-    return context[BindingContext.AUTO_CREATED_IT, target]!!
+    val target = expression.mainReference.resolveToDescriptors(context).singleOrNull() as? ValueParameterDescriptor ?: return null
+    return if (context[BindingContext.AUTO_CREATED_IT, target] == true) target else null
 }
 
+fun getLambdaByImplicitItReference(expression: KtNameReferenceExpression) =
+        resolveToAutoCreatedItDescriptor(expression)?.containingDeclaration?.source?.getPsi() as? KtFunctionLiteral
+
 // returns assignment which replaces initializer
 fun splitPropertyDeclaration(property: KtProperty): KtBinaryExpression {
     val parent = property.parent

idea/testData/usageHighlighter/implicitIt.kt

@@ -0,0 +1,13 @@
+fun foo(body: (Int) -> Unit) = body(1)
+
+fun test() {
+    foo {
+        val x = <info descr="null">~it</info> + 1
+        val xx = <info descr="null">it</info> + 2
+        foo {
+            val y = it - 1
+            val yy = it - 2
+        }
+        val xxx = <info descr="null">it</info> + 3
+    }
+}

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

@@ -19,10 +19,10 @@ package org.jetbrains.kotlin.idea.highlighter
 import com.intellij.codeInsight.daemon.impl.HighlightInfo
 import com.intellij.codeInsight.daemon.impl.HighlightInfoType
 import com.intellij.codeInsight.highlighting.actions.HighlightUsagesAction
-import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.testFramework.ExpectedHighlightingData
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
+import org.jetbrains.kotlin.idea.test.extractMarkerOffset
 
 abstract class AbstractUsageHighlightingTest: KotlinLightCodeInsightFixtureTestCase() {
     companion object {
@@ -36,13 +36,8 @@ abstract class AbstractUsageHighlightingTest: KotlinLightCodeInsightFixtureTestC
         val data = ExpectedHighlightingData(document, false, false, true, false, myFixture.file)
         data.init()
 
-        val caret = document.text.indexOf(CARET_TAG)
+        val caret = document.extractMarkerOffset(project, CARET_TAG)
         assert(caret != -1) { "Caret marker '$CARET_TAG' expected" }
-
-        WriteCommandAction.runWriteCommandAction(myFixture.project) {
-            document.deleteString(caret, caret + CARET_TAG.length)
-        }
-
         editor.caretModel.moveToOffset(caret)
 
         myFixture.testAction(HighlightUsagesAction())

idea/tests/org/jetbrains/kotlin/idea/highlighter/UsageHighlightingTestGenerated.java

@@ -36,6 +36,12 @@ public class UsageHighlightingTestGenerated extends AbstractUsageHighlightingTes
         KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/usageHighlighter"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
     }
 
+    @TestMetadata("implicitIt.kt")
+    public void testImplicitIt() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/usageHighlighter/implicitIt.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("importAlias.kt")
     public void testImportAlias() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/usageHighlighter/importAlias.kt");