HighlighterExtension can provide custom highlighting for calls

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.highlighter
 
 import com.intellij.lang.annotation.AnnotationHolder
+import com.intellij.openapi.extensions.Extensions
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.builtins.isFunctionTypeOrSubtype
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
@@ -31,6 +32,7 @@ import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
 import org.jetbrains.kotlin.resolve.calls.tasks.isDynamic
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 internal class FunctionsHighlightingVisitor(holder: AnnotationHolder, bindingContext: BindingContext) :
         AfterAnalysisHighlightingVisitor(holder, bindingContext) {
@@ -73,37 +75,26 @@ internal class FunctionsHighlightingVisitor(holder: AnnotationHolder, bindingCon
     private fun highlightCall(callee: PsiElement, resolvedCall: ResolvedCall<out CallableDescriptor>) {
         val calleeDescriptor = resolvedCall.resultingDescriptor
 
-        if (applyHighlighterExtensions(callee, calleeDescriptor)) return
-
-        if (calleeDescriptor.isDynamic()) {
-            highlightName(callee, DYNAMIC_FUNCTION_CALL)
-        }
-        else if (resolvedCall is VariableAsFunctionResolvedCall) {
-            val container = calleeDescriptor.containingDeclaration
-            val containedInFunctionClassOrSubclass = container is ClassDescriptor && container.defaultType.isFunctionTypeOrSubtype
-            highlightName(callee, if (containedInFunctionClassOrSubclass)
-                VARIABLE_AS_FUNCTION_CALL
-            else
-                VARIABLE_AS_FUNCTION_LIKE_CALL)
-        }
-        else {
-            if (calleeDescriptor is ConstructorDescriptor) {
-                highlightName(callee, CONSTRUCTOR_CALL)
-            }
-            else if (calleeDescriptor is FunctionDescriptor) {
-                val attributesKey = when {
-                    calleeDescriptor.extensionReceiverParameter != null ->
-                        EXTENSION_FUNCTION_CALL
-
-                    DescriptorUtils.isTopLevelDeclaration(calleeDescriptor) ->
-                        PACKAGE_FUNCTION_CALL
-
-                    else ->
-                        FUNCTION_CALL
-                }
-
-                highlightName(callee, attributesKey)
+        val key = Extensions.getExtensions(HighlighterExtension.EP_NAME).firstNotNullResult { extension ->
+            extension.highlightCall(callee, resolvedCall)
+        } ?: when {
+            calleeDescriptor.isDynamic() -> DYNAMIC_FUNCTION_CALL
+            resolvedCall is VariableAsFunctionResolvedCall -> {
+                val container = calleeDescriptor.containingDeclaration
+                val containedInFunctionClassOrSubclass = container is ClassDescriptor && container.defaultType.isFunctionTypeOrSubtype
+                if (containedInFunctionClassOrSubclass)
+                    VARIABLE_AS_FUNCTION_CALL
+                else
+                    VARIABLE_AS_FUNCTION_LIKE_CALL
             }
+            calleeDescriptor is ConstructorDescriptor -> CONSTRUCTOR_CALL
+            calleeDescriptor !is FunctionDescriptor -> null
+            calleeDescriptor.extensionReceiverParameter != null -> EXTENSION_FUNCTION_CALL
+            DescriptorUtils.isTopLevelDeclaration(calleeDescriptor) -> PACKAGE_FUNCTION_CALL
+            else -> FUNCTION_CALL
+        }
+        if (key != null) {
+            highlightName(callee, key)
         }
     }
 }

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

@@ -20,9 +20,13 @@ import com.intellij.openapi.editor.colors.TextAttributesKey
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 
 abstract class HighlighterExtension {
-    abstract fun highlightReference(elementToHighlight: PsiElement, descriptor: DeclarationDescriptor): TextAttributesKey?
+    abstract fun highlightDeclaration(elementToHighlight: PsiElement, descriptor: DeclarationDescriptor): TextAttributesKey?
+
+    open fun highlightCall(elementToHighlight: PsiElement, resolvedCall: ResolvedCall<*>): TextAttributesKey?
+            = highlightDeclaration(elementToHighlight, resolvedCall.resultingDescriptor)
 
     companion object {
         val EP_NAME = ExtensionPointName.create<HighlighterExtension>("org.jetbrains.kotlin.highlighterExtension")

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

@@ -24,16 +24,17 @@ import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.psi.KtVisitorVoid
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 abstract class HighlightingVisitor protected constructor(
-        private val holder: AnnotationHolder
+    private val holder: AnnotationHolder
 ) : KtVisitorVoid() {
 
     protected fun createInfoAnnotation(element: PsiElement, message: String? = null): Annotation =
-            createInfoAnnotation(element.textRange, message)
+        createInfoAnnotation(element.textRange, message)
 
     protected fun createInfoAnnotation(textRange: TextRange, message: String? = null): Annotation =
-            holder.createInfoAnnotation(textRange, message)
+        holder.createInfoAnnotation(textRange, message)
 
     protected fun highlightName(element: PsiElement, attributesKey: TextAttributesKey, message: String? = null) {
         if (NameHighlighter.namesHighlightingEnabled && !element.textRange.isEmpty) {
@@ -49,13 +50,14 @@ abstract class HighlightingVisitor protected constructor(
 
     protected fun applyHighlighterExtensions(element: PsiElement, descriptor: DeclarationDescriptor): Boolean {
         if (!NameHighlighter.namesHighlightingEnabled) return false
-        for (extension in Extensions.getExtensions(HighlighterExtension.EP_NAME)) {
-            val key = extension.highlightReference(element, descriptor)
-            if (key != null) {
-                highlightName(element, key)
-                return true
-            }
+
+        Extensions.getExtensions(HighlighterExtension.EP_NAME).firstNotNullResult { extension ->
+            extension.highlightDeclaration(element, descriptor)
+        }?.let { key ->
+            highlightName(element, key)
+            return true
         }
+
         return false
     }
 }

idea/idea-android/src/org/jetbrains/kotlin/android/AndroidHighlighterExtension.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.android
 
-import com.intellij.openapi.editor.colors.TextAttributesKey
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.android.synthetic.res.AndroidSyntheticProperty
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
@@ -24,9 +23,11 @@ import org.jetbrains.kotlin.idea.highlighter.HighlighterExtension
 import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors
 
 class AndroidHighlighterExtension : HighlighterExtension() {
-    override fun highlightReference(elementToHighlight: PsiElement, descriptor: DeclarationDescriptor): TextAttributesKey? =
-            if (descriptor is AndroidSyntheticProperty)
-                KotlinHighlightingColors.ANDROID_EXTENSIONS_PROPERTY_CALL
-            else
-                null
+    override fun highlightDeclaration(
+        elementToHighlight: PsiElement,
+        descriptor: DeclarationDescriptor
+    ) = when (descriptor) {
+        is AndroidSyntheticProperty -> KotlinHighlightingColors.ANDROID_EXTENSIONS_PROPERTY_CALL
+        else -> null
+    }
 }