Add K2 KotlinCallHighlighterExtension

This commit adds KotlinCallHighlighterExtension for K2 call highlighter.
If we register an implementation of KotlinCallHighlighterExtension
interface in XML, FunctionCallHighlighter will use it instead of the
default call highlighter (the default call highlighter function is
FunctionCallHighlighter::getTextAttributesForCall(call: KtCall)
function).

GitOrigin-RevId: ba69aff2c819887bd0d3b5cac761def9301d28ab

plugins/kotlin/highlighting/src/org/jetbrains/kotlin/idea/highlighting/KotlinCallHighlighterExtension.kt

@@ -0,0 +1,33 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.kotlin.idea.highlighting
+
+import com.intellij.codeInsight.daemon.impl.HighlightInfoType
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
+import org.jetbrains.kotlin.analysis.api.calls.KtCall
+import org.jetbrains.kotlin.idea.highlighting.highlighters.FunctionCallHighlighter
+
+/**
+ * An interface for an extension point that enables a user-defined call highlighter.
+ *
+ * To add a new custom call highlighter, you have to add
+ *
+ *   <extensions defaultExtensionNs="org.jetbrains.kotlin"><callHighlighterExtension implementation="..."/></extensions>
+ *
+ * in XML. The "implementation" must be a class that implements this interface. See highlighting-k2.xml.
+ */
+interface KotlinCallHighlighterExtension {
+    /**
+     * A function determines how to highlight [call] for PSI [elementToHighlight].
+     *
+     * Note that this function is used by [FunctionCallHighlighter]. When this function returns a non-null [HighlightInfoType],
+     * [FunctionCallHighlighter] uses it instead of the default one to highlight [call].
+     */
+    context(KtAnalysisSession)
+    fun highlightCall(elementToHighlight: PsiElement, call: KtCall): HighlightInfoType?
+
+    companion object {
+        val EP_NAME = ExtensionPointName.create<KotlinCallHighlighterExtension>("org.jetbrains.kotlin.callHighlighterExtension")
+    }
+}

plugins/kotlin/highlighting/src/org/jetbrains/kotlin/idea/highlighting/highlighters/FunctionCallHighlighter.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.idea.base.highlighting.HighlightingFactory
 import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightInfoTypeSemanticNames
+import org.jetbrains.kotlin.idea.highlighting.KotlinCallHighlighterExtension
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.Name
@@ -38,8 +39,8 @@ internal class FunctionCallHighlighter(
         if (operationReference.isAssignment()) return emptyList()
         val call = expression.resolveCall()?.successfulCallOrNull<KtCall>() ?: return emptyList()
         if (call is KtSimpleFunctionCall && (call.symbol as? KtFunctionSymbol)?.isOperator == true) return emptyList()
-        val h = getTextAttributesForCall(call)?.let { attributes ->
-            HighlightingFactory.highlightName(operationReference, attributes)
+        val h = getDefaultHighlightInfoTypeForCall(call)?.let { highlightInfoType ->
+            HighlightingFactory.highlightName(operationReference, highlightInfoType)
         }
         return listOfNotNull(h)
     }
@@ -49,20 +50,21 @@ internal class FunctionCallHighlighter(
 
     context(KtAnalysisSession)
     private fun highlightCallExpression(expression: KtCallExpression): List<HighlightInfo.Builder> {
-        return listOfNotNull(expression.calleeExpression
-            ?.takeUnless { it is KtLambdaExpression }
-            ?.takeUnless { it is KtCallExpression /* KT-16159 */ }
-            ?.let { callee ->
-                expression.resolveCall()?.singleCallOrNull<KtCall>()?.let { call ->
-                    getTextAttributesForCall(call)?.let { attributes ->
-                        HighlightingFactory.highlightName (callee, attributes)
-                    }
-                }
-            })
+        val callee = expression.calleeExpression ?: return emptyList()
+        val call = expression.resolveCall()?.singleCallOrNull<KtCall>() ?: return emptyList()
+        if (callee is KtLambdaExpression || callee is KtCallExpression /* KT-16159 */) return emptyList()
+        val highlightInfoType = getHighlightInfoTypeForCallFromExtension(callee, call)
+            ?: getDefaultHighlightInfoTypeForCall(call)
+            ?: return emptyList()
+        return listOfNotNull(HighlightingFactory.highlightName(callee, highlightInfoType))
     }
 
     context(KtAnalysisSession)
-    private fun getTextAttributesForCall(call: KtCall): HighlightInfoType? {
+    private fun getHighlightInfoTypeForCallFromExtension(callee: KtExpression, call: KtCall): HighlightInfoType? =
+        KotlinCallHighlighterExtension.EP_NAME.getExtensions(project).firstNotNullOfOrNull { it.highlightCall(callee, call) }
+
+    context(KtAnalysisSession)
+    private fun getDefaultHighlightInfoTypeForCall(call: KtCall): HighlightInfoType? {
         if (call !is KtSimpleFunctionCall) return null
         return when (val function = call.symbol) {
             is KtConstructorSymbol -> KotlinHighlightInfoTypeSemanticNames.CONSTRUCTOR_CALL
@@ -89,4 +91,4 @@ internal class FunctionCallHighlighter(
         private val KOTLIN_SUSPEND_BUILT_IN_FUNCTION_FQ_NAME_CALLABLE_ID =
             CallableId(StandardNames.BUILT_INS_PACKAGE_FQ_NAME, Name.identifier("suspend"))
     }
-}
+}

plugins/kotlin/highlighting/test/org/jetbrains/kotlin/idea/k2/highlighting/AbstractK2HighlightingMetaInfoWithExtensionTest.kt

@@ -0,0 +1,48 @@
+// Copyright 2000-2023 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.codeInsight.daemon.impl.HighlightInfoType
+import com.intellij.openapi.extensions.ExtensionPoint
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
+import org.jetbrains.kotlin.analysis.api.annotations.annotations
+import org.jetbrains.kotlin.analysis.api.calls.KtCall
+import org.jetbrains.kotlin.analysis.api.calls.KtCallableMemberCall
+import org.jetbrains.kotlin.analysis.api.calls.symbol
+import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightInfoTypeSemanticNames
+import org.jetbrains.kotlin.idea.highlighting.KotlinCallHighlighterExtension
+
+object KotlinCallHighlighterExtensionForTest : KotlinCallHighlighterExtension {
+    private enum class HighlightType(val annotationName: String) {
+        SUSPEND("MySuspend"),
+        DYNAMIC("MyDynamic"),
+        EXTENSION("MyExtension"),
+    }
+
+    context(KtAnalysisSession)
+    override fun highlightCall(elementToHighlight: PsiElement, call: KtCall): HighlightInfoType? {
+        if (call !is KtCallableMemberCall<*, *>) return null
+        val highlightType = call.partiallyAppliedSymbol.symbol.annotations.firstNotNullOfOrNull { annotation ->
+            HighlightType.values().singleOrNull { it.annotationName == annotation.classId?.shortClassName?.asString() }
+        } ?: return null
+        return when (highlightType) {
+            HighlightType.SUSPEND -> KotlinHighlightInfoTypeSemanticNames.SUSPEND_FUNCTION_CALL
+            HighlightType.DYNAMIC -> KotlinHighlightInfoTypeSemanticNames.DYNAMIC_FUNCTION_CALL
+            HighlightType.EXTENSION -> KotlinHighlightInfoTypeSemanticNames.EXTENSION_FUNCTION_CALL
+        }
+    }
+}
+
+abstract class AbstractK2HighlightingMetaInfoWithExtensionTest : AbstractK2HighlightingMetaInfoTest() {
+    override fun setUp() {
+        super.setUp()
+        project.extensionArea.registerExtensionPoint(
+            KotlinCallHighlighterExtension.EP_NAME.name,
+            KotlinCallHighlighterExtension::class.java.name,
+            ExtensionPoint.Kind.INTERFACE,
+            true
+        )
+        project.extensionArea.getExtensionPoint(KotlinCallHighlighterExtension.EP_NAME)
+            .registerExtension(KotlinCallHighlighterExtensionForTest, testRootDisposable)
+    }
+}

plugins/kotlin/highlighting/test/org/jetbrains/kotlin/idea/k2/highlighting/K2HighlightingMetaInfoWithExtensionTestGenerated.java

@@ -0,0 +1,30 @@
+// Copyright 2000-2023 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.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.idea.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.jetbrains.kotlin.idea.base.test.TestRoot;
+import org.junit.runner.RunWith;
+
+/**
+ * This class is generated by {@link org.jetbrains.kotlin.testGenerator.generator.TestGenerator}.
+ * DO NOT MODIFY MANUALLY.
+ */
+@SuppressWarnings("all")
+@TestRoot("highlighting")
+@TestDataPath("$CONTENT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+@TestMetadata("../idea/tests/testData/highlighterMetaInfoWithExtension")
+public class K2HighlightingMetaInfoWithExtensionTestGenerated extends AbstractK2HighlightingMetaInfoWithExtensionTest {
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    @TestMetadata("Functions.kt")
+    public void testFunctions() throws Exception {
+        runTest("../idea/tests/testData/highlighterMetaInfoWithExtension/Functions.kt");
+    }
+}

plugins/kotlin/idea/tests/testData/highlighterMetaInfoWithExtension/Functions.kt

@@ -0,0 +1,52 @@
+// FIR_IDENTICAL
+// WITH_STDLIB
+// CHECK_SYMBOL_NAMES
+// HIGHLIGHTER_ATTRIBUTES_KEY
+fun global() {
+    fun inner() {
+
+    }
+    inner()
+}
+
+fun Int.ext() {
+}
+
+infix fun Int.fif(y: Int) {
+    this * y
+}
+
+open class Container {
+    open fun member() {
+        global()
+        5.ext()
+        member()
+        5 fif 6
+    }
+}
+
+fun foo() {
+    suspend {
+
+    }
+}
+
+annotation class MySuspend
+annotation class MyDynamic
+annotation class MyExtension
+
+@MySuspend
+fun suspendFunction() {
+}
+@MyDynamic
+fun dynamicFunction() {
+}
+@MyExtension
+fun extensionFunction() {
+}
+
+fun testFunctionsWithSpecialAnnotations() {
+    suspendFunction()
+    dynamicFunction()
+    extensionFunction()
+}

plugins/kotlin/idea/tests/testData/highlighterMetaInfoWithExtension/Functions.kt.highlighting

@@ -0,0 +1,52 @@
+// FIR_IDENTICAL
+// WITH_STDLIB
+// CHECK_SYMBOL_NAMES
+// HIGHLIGHTER_ATTRIBUTES_KEY
+fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>global<!>() {
+    fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>inner<!>() {
+
+    }
+    <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_CALL'")!>inner<!>()
+}
+
+fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_CLASS'")!>Int<!>.<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>ext<!>() {
+}
+
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_BUILTIN_ANNOTATION'")!>infix<!> fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_CLASS'")!>Int<!>.<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>fif<!>(<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_PARAMETER'")!>y<!>: <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_CLASS'")!>Int<!>) {
+    this * <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_PARAMETER'")!>y<!>
+}
+
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_BUILTIN_ANNOTATION'")!>open<!> class <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_CLASS'")!>Container<!> {
+    <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_BUILTIN_ANNOTATION'")!>open<!> fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>member<!>() {
+        <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_PACKAGE_FUNCTION_CALL'")!>global<!>()
+        5.<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_EXTENSION_FUNCTION_CALL'")!>ext<!>()
+        <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_CALL'")!>member<!>()
+        5 <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_EXTENSION_FUNCTION_CALL'")!>fif<!> 6
+    }
+}
+
+fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>foo<!>() {
+    <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_KEYWORD'")!>suspend<!> {
+
+    }
+}
+
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_BUILTIN_ANNOTATION'")!>annotation<!> class <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_ANNOTATION'")!>MySuspend<!>
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_BUILTIN_ANNOTATION'")!>annotation<!> class <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_ANNOTATION'")!>MyDynamic<!>
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_BUILTIN_ANNOTATION'")!>annotation<!> class <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_ANNOTATION'")!>MyExtension<!>
+
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_ANNOTATION'")!>@MySuspend<!>
+fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>suspendFunction<!>() {
+}
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_ANNOTATION'")!>@MyDynamic<!>
+fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>dynamicFunction<!>() {
+}
+<!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_ANNOTATION'")!>@MyExtension<!>
+fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>extensionFunction<!>() {
+}
+
+fun <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_FUNCTION_DECLARATION'")!>testFunctionsWithSpecialAnnotations<!>() {
+    <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_SUSPEND_FUNCTION_CALL'")!>suspendFunction<!>()
+    <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_DYNAMIC_FUNCTION_CALL'")!>dynamicFunction<!>()
+    <!HIGHLIGHTING("severity='SYMBOL_TYPE_SEVERITY'; highlightingTextAttributesKey='KOTLIN_EXTENSION_FUNCTION_CALL'")!>extensionFunction<!>()
+}

plugins/kotlin/plugin/k2/resources/META-INF/highlighting-k2.xml

@@ -1,4 +1,11 @@
 <idea-plugin>
+  <extensionPoints>
+    <extensionPoint
+            qualifiedName="org.jetbrains.kotlin.callHighlighterExtension"
+            interface="org.jetbrains.kotlin.idea.highlighting.KotlinCallHighlighterExtension"
+            dynamic="true"/>
+  </extensionPoints>
+
   <extensions defaultExtensionNs="com.intellij">
     <problemHighlightFilter implementation="org.jetbrains.kotlin.idea.highlighter.KotlinProblemHighlightFilter"/>
     <problemFileHighlightFilter implementation="org.jetbrains.kotlin.idea.highlighting.KotlinProblemFileHighlightFilter"/>

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

@@ -3,6 +3,7 @@ package org.jetbrains.kotlin.fir.testGenerator
 
 import org.jetbrains.kotlin.idea.k2.highlighting.AbstractK2BundledCompilerPluginsHighlightingMetaInfoTest
 import org.jetbrains.kotlin.idea.k2.highlighting.AbstractK2HighlightingMetaInfoTest
+import org.jetbrains.kotlin.idea.k2.highlighting.AbstractK2HighlightingMetaInfoWithExtensionTest
 import org.jetbrains.kotlin.idea.k2.highlighting.AbstractOutsiderHighlightingTest
 import org.jetbrains.kotlin.testGenerator.model.*
 
@@ -15,6 +16,10 @@ internal fun MutableTWorkspace.generateK2HighlighterTests() {
         testClass<AbstractK2BundledCompilerPluginsHighlightingMetaInfoTest> {
             model("highlighterMetaInfoWithBundledCompilerPlugins")
         }
+
+        testClass<AbstractK2HighlightingMetaInfoWithExtensionTest> {
+            model("highlighterMetaInfoWithExtension")
+        }
     }
 
     testGroup("highlighting") {