KTIJ-26211: IDEA: On selecting it highlight the call accepting it's in lambda

plugins/kotlin/code-insight/utils/src/org/jetbrains/kotlin/idea/codeinsight/utils/implicitLambdaParameterUtils.kt

@@ -12,7 +12,9 @@ import org.jetbrains.kotlin.analysis.api.symbols.KaAnonymousFunctionSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaValueParameterSymbol
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtLambdaArgument
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 
 fun KtNameReferenceExpression.isReferenceToImplicitLambdaParameter(): Boolean {
@@ -51,3 +53,22 @@ fun KaValueParameterSymbol.getFunctionLiteralByImplicitLambdaParameterSymbol():
     val lambda = containingDeclaration as? KaAnonymousFunctionSymbol ?: return null
     return lambda.psi as? KtFunctionLiteral
 }
+
+
+fun KtNameReferenceExpression.getLambdaFunctionLiteralByNameReference(): KtFunctionLiteral? {
+    val lambda = (this.parent as? KtCallExpression)?.lambdaArguments?.singleOrNull()?.getLambdaExpression() ?: return null
+    return lambda.functionLiteral
+}
+
+fun KtFunctionLiteral.getImplicitItLambdaParameterSymbol(): KaValueParameterSymbol? {
+    @OptIn(KaAllowAnalysisOnEdt::class)
+    allowAnalysisOnEdt {
+        analyze(this) {
+            if (valueParameters.isNotEmpty() || symbol.valueParameters.size != 1 || !symbol.valueParameters.first().isImplicitLambdaParameter) {
+                return null
+            } else {
+                return symbol.valueParameters.first()
+            }
+        }
+    }
+}

plugins/kotlin/highlighting/highlighting-k2/resources/intellij.kotlin.highlighting.xml

@@ -23,6 +23,7 @@
     <highlightErrorFilter implementation="org.jetbrains.kotlin.idea.highlighting.visitor.KotlinHighlightInjectionErrorFilter"/>
 
     <highlightUsagesHandlerFactory implementation="org.jetbrains.kotlin.idea.highlighting.KotlinHighlightExitPointsHandlerFactory"/>
+    <highlightUsagesHandlerFactory implementation="org.jetbrains.kotlin.idea.highlighter.KotlinHighlightImplicitItHandlerFactory"/>
 
     <editorNotificationProvider implementation="org.jetbrains.kotlin.idea.highlighting.SupportAvailabilityNotificationProvider"/>
 

plugins/kotlin/highlighting/highlighting-k2/test/org/jetbrains/kotlin/idea/k2/highlighting/HighlightImplicitItHandlerTestGenerated.java

@@ -0,0 +1,43 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+
+package org.jetbrains.kotlin.idea.k2.highlighting;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.idea.base.plugin.KotlinPluginMode;
+import org.jetbrains.kotlin.idea.base.test.TestRoot;
+import org.jetbrains.kotlin.idea.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.idea.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+import org.jetbrains.kotlin.idea.highlighter.AbstractHighlightImplicitItHandlerTest;
+
+/**
+ * This class is generated by {@link org.jetbrains.kotlin.testGenerator.generator.TestGenerator}.
+ * DO NOT MODIFY MANUALLY.
+ */
+@SuppressWarnings("all")
+@TestRoot("highlighting/highlighting-k2")
+@TestDataPath("$CONTENT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+@TestMetadata("../../idea/tests/testData/implicitIt")
+public class HighlightImplicitItHandlerTestGenerated extends AbstractHighlightImplicitItHandlerTest {
+    @java.lang.Override
+    @org.jetbrains.annotations.NotNull
+    public final KotlinPluginMode getPluginMode() {
+        return KotlinPluginMode.K2;
+    }
+
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    @TestMetadata("HighlightCallAcceptingLambdaImplicitIt.kt")
+    public void testHighlightCallAcceptingLambdaImplicitIt() throws Exception {
+        runTest("../../idea/tests/testData/implicitIt/HighlightCallAcceptingLambdaImplicitIt.kt");
+    }
+
+    @TestMetadata("HighlightImplicitIt.kt")
+    public void testHighlightImplicitIt() throws Exception {
+        runTest("../../idea/tests/testData/implicitIt/HighlightImplicitIt.kt");
+    }
+}

plugins/kotlin/highlighting/highlighting-k1/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightImplicitItHandlerFactory.kt → plugins/kotlin/highlighting/highlighting-shared/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightImplicitItHandlerFactory.kt

@@ -7,27 +7,33 @@ 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.psi.util.elementType
 import com.intellij.util.Consumer
 import org.jetbrains.kotlin.K1Deprecation
+import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.idea.codeinsight.utils.getFunctionLiteralByImplicitLambdaParameter
+import org.jetbrains.kotlin.idea.codeinsight.utils.getImplicitItLambdaParameterSymbol
+import org.jetbrains.kotlin.idea.codeinsight.utils.getLambdaFunctionLiteralByNameReference
+import org.jetbrains.kotlin.idea.references.getCalleeByLambdaArgument
 import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
 
 @K1Deprecation
 class KotlinHighlightImplicitItHandlerFactory : HighlightUsagesHandlerFactoryBase() {
     override fun createHighlightUsagesHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
-        if (target !is LeafPsiElement
-            || target.elementType !is KtToken // do not trigger loading of KtTokens in Java
+        if (target.elementType !is KtToken // do not trigger loading of KtTokens in Java
             || target.elementType != KtTokens.IDENTIFIER) {
             return null
         }
 
         val refExpr = target.parent as? KtNameReferenceExpression ?: return null
-        val lambda = refExpr.getFunctionLiteralByImplicitLambdaParameter() ?: return null
+        val lambda = refExpr.getFunctionLiteralByImplicitLambdaParameter()
+            ?: refExpr.getLambdaFunctionLiteralByNameReference()?.takeIf { it.getImplicitItLambdaParameterSymbol() != null }
+            ?: return null
         return object : HighlightUsagesHandlerBase<KtNameReferenceExpression>(editor, file) {
             override fun getTargets() = listOf(refExpr)
 
@@ -37,15 +43,31 @@ class KotlinHighlightImplicitItHandlerFactory : HighlightUsagesHandlerFactoryBas
             ) = selectionConsumer.consume(targets)
 
             override fun computeUsages(targets: List<KtNameReferenceExpression>) {
+                var addLambdaOccurrence = false
                 lambda.accept(
                     object : KtTreeVisitorVoid() {
                         override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
                             if (expression is KtNameReferenceExpression && expression.getFunctionLiteralByImplicitLambdaParameter() == lambda) {
                                 addOccurrence(expression)
+                                addLambdaOccurrence = true
+                            }
+                        }
+
+                        //implicit lambda parameter is not available inside another lambda with implicit parameter or name `it`
+                        override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression) {
+                            if (lambdaExpression.functionLiteral.getImplicitItLambdaParameterSymbol() == null &&
+                                lambdaExpression.functionLiteral.valueParameters.find { it.name == StandardNames.IMPLICIT_LAMBDA_PARAMETER_NAME.asString() } == null
+                            ) {
+                                super.visitLambdaExpression(lambdaExpression)
                             }
                         }
                     }
                 )
+                if (addLambdaOccurrence) {
+                    lambda.getCalleeByLambdaArgument()?.also {
+                        addOccurrence(it)
+                    }
+                }
             }
         }
     }

plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/highlighter/HighlightImplicitItHandlerTestGenerated.java

@@ -0,0 +1,42 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+
+package org.jetbrains.kotlin.idea.highlighter;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.idea.base.plugin.KotlinPluginMode;
+import org.jetbrains.kotlin.idea.base.test.TestRoot;
+import org.jetbrains.kotlin.idea.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.idea.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+
+/**
+ * This class is generated by {@link org.jetbrains.kotlin.testGenerator.generator.TestGenerator}.
+ * DO NOT MODIFY MANUALLY.
+ */
+@SuppressWarnings("all")
+@TestRoot("idea/tests")
+@TestDataPath("$CONTENT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+@TestMetadata("testData/implicitIt")
+public class HighlightImplicitItHandlerTestGenerated extends AbstractHighlightImplicitItHandlerTest {
+    @java.lang.Override
+    @org.jetbrains.annotations.NotNull
+    public final KotlinPluginMode getPluginMode() {
+        return KotlinPluginMode.K1;
+    }
+
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    @TestMetadata("HighlightCallAcceptingLambdaImplicitIt.kt")
+    public void testHighlightCallAcceptingLambdaImplicitIt() throws Exception {
+        runTest("testData/implicitIt/HighlightCallAcceptingLambdaImplicitIt.kt");
+    }
+
+    @TestMetadata("HighlightImplicitIt.kt")
+    public void testHighlightImplicitIt() throws Exception {
+        runTest("testData/implicitIt/HighlightImplicitIt.kt");
+    }
+}

plugins/kotlin/idea/tests/testData/implicitIt/HighlightCallAcceptingLambdaImplicitIt.kt

@@ -0,0 +1,19 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+fun acceptL1(fn: (Int) -> Unit) {}
+fun acceptL2(fn: (Int, Int) -> Unit) {}
+
+fun nestLambdas() {
+    acceptL1 {
+        // More code.
+        <info descr="null">~acceptL1</info> {
+            // And more code.
+            acceptL2 { x, y ->
+                // And more code.
+                println(<info descr="null">it</info>)
+                acceptL1 {
+                    it
+                }
+            }
+        }
+    }
+}

plugins/kotlin/idea/tests/testData/implicitIt/HighlightImplicitIt.kt

@@ -0,0 +1,19 @@
+// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+fun acceptL1(fn: (Int) -> Unit) {}
+fun acceptL2(fn: (Int, Int) -> Unit) {}
+
+fun nestLambdas() {
+    acceptL1 {
+        // More code.
+        <info descr="null">acceptL1</info> {
+            // And more code.
+            acceptL2 { x, y ->
+                // And more code.
+                println(<info descr="null">~it</info>)
+                acceptL1 {
+                    it
+                }
+            }
+        }
+    }
+}

plugins/kotlin/intellij.kotlin.base.test/test/org/jetbrains/kotlin/idea/highlighter/AbstractHighlightImplicitItHandlerTest.kt

@@ -0,0 +1,79 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+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.HighlightUsagesHandler
+import com.intellij.codeInsight.highlighting.highlightUsages
+import com.intellij.openapi.application.ReadAction
+import com.intellij.openapi.editor.asTextRange
+import com.intellij.openapi.editor.colors.EditorColors
+import com.intellij.openapi.project.DumbService
+import com.intellij.psi.PsiElement
+import com.intellij.testFramework.ExpectedHighlightingData
+import com.intellij.util.concurrency.AppExecutorUtil
+import org.jetbrains.kotlin.idea.base.test.IgnoreTests
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.idea.test.extractMarkerOffset
+import java.util.concurrent.Callable
+
+abstract class AbstractHighlightImplicitItHandlerTest : KotlinLightCodeInsightFixtureTestCase() {
+
+    companion object {
+        // Not standard <caret> to leave it in text after configureByFile and remove manually after collecting highlighting information
+        const val CARET_TAG = "~"
+    }
+
+    open fun doTest(unused: String) {
+        val disableDirective = IgnoreTests.DIRECTIVES.of(pluginMode)
+
+        IgnoreTests.runTestIfNotDisabledByFileDirective(
+            dataFile().toPath(),
+            disableDirective
+        ) {
+        myFixture.configureByFile(fileName())
+
+        val editor = myFixture.editor
+
+        val document = myFixture.editor.document
+        val data = ExpectedHighlightingData(document, false, false, true, false)
+        data.init()
+
+        val caret = document.extractMarkerOffset(project, CARET_TAG)
+        assert(caret != -1) { "Caret marker '${CARET_TAG}' expected" }
+        editor.caretModel.moveToOffset(caret)
+
+        val customHandler =
+            HighlightUsagesHandler.createCustomHandler<PsiElement>(editor, myFixture.file)
+
+        if (customHandler != null) {
+            ReadAction.nonBlocking(Callable {
+                customHandler.highlightUsages()
+            }).submit(AppExecutorUtil.getAppExecutorService())
+                .get()
+        } else {
+            DumbService.getInstance(project).withAlternativeResolveEnabled(Runnable {
+                highlightUsages(project, editor, file)
+            })
+        }
+
+        val ranges = editor.markupModel.allHighlighters
+            .filter { it.textAttributesKey == EditorColors.SEARCH_RESULT_ATTRIBUTES || it.textAttributesKey == EditorColors.WRITE_SEARCH_RESULT_ATTRIBUTES }
+            .mapNotNull { it.asTextRange }
+
+        val infos = ranges.toHashSet()
+            .map {
+                var startOffset = it.startOffset
+                var endOffset = it.endOffset
+
+                if (startOffset > caret) startOffset += CARET_TAG.length
+                if (endOffset > caret) endOffset += CARET_TAG.length
+
+                HighlightInfo.newHighlightInfo(HighlightInfoType.INFORMATION)
+                    .range(startOffset, endOffset)
+                    .create()
+            }
+
+        data.checkResult(myFixture.file, infos, StringBuilder(document.text).insert(caret, CARET_TAG).toString())
+    }}
+}

plugins/kotlin/util/test-generator-fe10/test/org/jetbrains/kotlin/fe10/testGenerator/Fe10GenerateTests.kt

@@ -117,6 +117,7 @@ import org.jetbrains.kotlin.idea.hierarchy.AbstractHierarchyWithLibTest
 import org.jetbrains.kotlin.idea.highlighter.AbstractCustomHighlightUsageHandlerTest
 import org.jetbrains.kotlin.idea.highlighter.AbstractDiagnosticMessageTest
 import org.jetbrains.kotlin.idea.highlighter.AbstractDslHighlighterTest
+import org.jetbrains.kotlin.idea.highlighter.AbstractHighlightImplicitItHandlerTest
 import org.jetbrains.kotlin.idea.highlighter.AbstractK1HighlightingMetaInfoTest
 import org.jetbrains.kotlin.idea.highlighter.AbstractK1HighlightingTest
 import org.jetbrains.kotlin.idea.highlighter.AbstractKotlinReceiverUsageHighlightingTest
@@ -1016,6 +1017,10 @@ private fun assembleWorkspace(): TWorkspace = workspace(KotlinPluginMode.K1) {
         testClass<AbstractMultiPlatformHighlightingTest> {
             model("multiModuleHighlighting/multiplatform/", isRecursive = false, pattern = DIRECTORY)
         }
+
+        testClass<AbstractHighlightImplicitItHandlerTest>("org.jetbrains.kotlin.idea.highlighter.HighlightImplicitItHandlerTestGenerated") {
+            model("implicitIt")
+        }
     }
 
     testGroup("idea/tests", category = CODE_INSIGHT) {

plugins/kotlin/util/test-generator-fir/test/org/jetbrains/kotlin/fir/testGenerator/GenerateK2HighlighterTests.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.idea.k2.highlighting.AbstractK2HighlightUsagesTest
 import org.jetbrains.kotlin.idea.k2.highlighting.AbstractK2HighlightingMetaInfoTest
 import org.jetbrains.kotlin.idea.k2.highlighting.AbstractK2HighlightingMetaInfoWithExtensionTest
 import org.jetbrains.kotlin.idea.core.script.k2.definitions.AbstractScriptHighlightingMetaInfoTest
+import org.jetbrains.kotlin.idea.highlighter.AbstractHighlightImplicitItHandlerTest
 import org.jetbrains.kotlin.idea.k2.highlighting.AbstractOutsiderHighlightingTest
 import org.jetbrains.kotlin.idea.test.kmp.KMPTestPlatform
 import org.jetbrains.kotlin.testGenerator.model.GroupCategory.HIGHLIGHTING
@@ -78,4 +79,10 @@ internal fun MutableTWorkspace.generateK2HighlighterTests() {
             model("outsider", pattern = Patterns.DIRECTORY, isRecursive = false, passTestDataPath = false)
         }
     }
+
+    testGroup("highlighting/highlighting-k2", category = HIGHLIGHTING, testDataPath = "../../idea/tests/testData") {
+        testClass<AbstractHighlightImplicitItHandlerTest>("org.jetbrains.kotlin.idea.k2.highlighting.HighlightImplicitItHandlerTestGenerated") {
+            model("implicitIt")
+        }
+    }
 }