Support dsl highlighting for object references

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

@@ -17,8 +17,31 @@
 package org.jetbrains.kotlin.idea.highlighter
 
 import com.intellij.lang.annotation.AnnotationHolder
+import com.intellij.openapi.editor.colors.TextAttributesKey
+import com.intellij.openapi.extensions.Extensions
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 abstract class AfterAnalysisHighlightingVisitor protected constructor(
-        holder: AnnotationHolder, protected var bindingContext: BindingContext
-) : HighlightingVisitor(holder)
+    holder: AnnotationHolder, protected var bindingContext: BindingContext
+) : HighlightingVisitor(holder) {
+
+    protected fun attributeKeyForDeclarationFromExtensions(element: PsiElement, descriptor: DeclarationDescriptor) =
+        Extensions.getExtensions(HighlighterExtension.EP_NAME).firstNotNullResult { extension ->
+            extension.highlightDeclaration(element, descriptor)
+        }
+
+    protected fun attributeKeyForCallFromExtensions(
+        expression: KtSimpleNameExpression,
+        resolvedCall: ResolvedCall<out CallableDescriptor>
+    ): TextAttributesKey? {
+        return Extensions.getExtensions(HighlighterExtension.EP_NAME).firstNotNullResult { extension ->
+            extension.highlightCall(expression, resolvedCall)
+        }
+    }
+}

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

@@ -19,12 +19,9 @@ package org.jetbrains.kotlin.idea.highlighter
 import com.intellij.lang.annotation.Annotation
 import com.intellij.lang.annotation.AnnotationHolder
 import com.intellij.openapi.editor.colors.TextAttributesKey
-import com.intellij.openapi.extensions.Extensions
 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
@@ -47,9 +44,4 @@ abstract class HighlightingVisitor protected constructor(
             createInfoAnnotation(textRange, message).textAttributes = attributesKey
         }
     }
-
-    protected fun attributeKeyForDeclarationFromExtensions(element: PsiElement, descriptor: DeclarationDescriptor) =
-        Extensions.getExtensions(HighlighterExtension.EP_NAME).firstNotNullResult { extension ->
-            extension.highlightDeclaration(element, descriptor)
-        }
 }

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

@@ -25,6 +25,8 @@ import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.util.FakeCallableDescriptorForObject
 
 internal class TypeKindHighlightingVisitor(holder: AnnotationHolder, bindingContext: BindingContext) :
     AfterAnalysisHighlightingVisitor(holder, bindingContext) {
@@ -40,7 +42,7 @@ internal class TypeKindHighlightingVisitor(holder: AnnotationHolder, bindingCont
 
         val referenceTarget = computeReferencedDescriptor(expression) ?: return
 
-        val key = when (referenceTarget) {
+        val key = attributeKeyForObjectAccess(expression) ?: when (referenceTarget) {
             is TypeParameterDescriptor -> TYPE_PARAMETER
             is TypeAliasDescriptor -> TYPE_ALIAS
             !is ClassDescriptor -> return
@@ -53,6 +55,13 @@ internal class TypeKindHighlightingVisitor(holder: AnnotationHolder, bindingCont
         highlightName(computeHighlightingRangeForUsage(expression, referenceTarget), key)
     }
 
+    private fun attributeKeyForObjectAccess(expression: KtSimpleNameExpression): TextAttributesKey? {
+        val resolvedCall = expression.getResolvedCall(bindingContext)
+        return if (resolvedCall?.resultingDescriptor is FakeCallableDescriptorForObject)
+            attributeKeyForCallFromExtensions(expression, resolvedCall)
+        else null
+    }
+
     private fun computeReferencedDescriptor(expression: KtSimpleNameExpression): DeclarationDescriptor? {
         val referenceTarget = bindingContext.get(BindingContext.REFERENCE_TARGET, expression)
 

idea/testData/dslHighlighter/objectAccess.kt

@@ -0,0 +1,49 @@
+package p
+
+import p.D.*
+import p.C.*
+
+@DslMarker
+annotation class A
+
+@DslMarker
+annotation class B
+
+@A
+class C {
+    @A
+    object C1
+
+    @A
+    object C2
+
+    class C3 {
+        @A
+        companion object
+    }
+}
+
+@B
+class D {
+    @B
+    object D1
+
+    class D2 {
+        @B
+        companion object
+    }
+}
+
+fun test() {
+    with(C()) {
+        C1 // 4
+        C2 // 4
+        C3 // 4
+        C3.Companion // 4
+    }
+    with(D()) {
+        D1 // 1
+        D2 // 1
+        D2.Companion // 1
+    }
+}

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

@@ -34,6 +34,11 @@ public class DslHighlighterTestGenerated extends AbstractDslHighlighterTest {
         runTest("idea/testData/dslHighlighter/functionCalls.kt");
     }
 
+    @TestMetadata("objectAccess.kt")
+    public void testObjectAccess() throws Exception {
+        runTest("idea/testData/dslHighlighter/objectAccess.kt");
+    }
+
     @TestMetadata("propertyAccess.kt")
     public void testPropertyAccess() throws Exception {
         runTest("idea/testData/dslHighlighter/propertyAccess.kt");