Implement custom highlighting for dsls

A declaration is considered a part of dsl if it's marked with
    a marker annotation (which is in turn must be marked as @DslMarker)
Highlighting attributes are then chosen randomly from a preset list
    based on marker annotation fqname

idea/src/META-INF/plugin.xml

@@ -2729,6 +2729,8 @@
     <binaryExtension implementation="org.jetbrains.kotlin.idea.util.KotlinBuiltInBinary"/>
     <binaryExtension implementation="org.jetbrains.kotlin.idea.util.KotlinModuleBinary"/>
     <binaryExtension implementation="org.jetbrains.kotlin.idea.util.KotlinJsMetaBinary"/>
+
+    <highlighterExtension implementation="org.jetbrains.kotlin.idea.highlighter.dsl.DslHighlighterExtension"/>
   </extensions>
 
 </idea-plugin>

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

@@ -26,6 +26,7 @@ import com.intellij.openapi.options.colors.RainbowColorSettingsPage
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.KotlinIcons
 import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.idea.highlighter.dsl.DslHighlighterExtension
 import java.lang.reflect.Modifier
 import java.util.*
 
@@ -109,6 +110,9 @@ var <PACKAGE_PROPERTY><MUTABLE_VARIABLE>globalCounter</MUTABLE_VARIABLE></PACKAG
 
             }
         }
+
+        map.putAll(DslHighlighterExtension.descriptionsToStyles)
+
         return map
     }
 
@@ -172,7 +176,8 @@ var <PACKAGE_PROPERTY><MUTABLE_VARIABLE>globalCounter</MUTABLE_VARIABLE></PACKAG
                        KotlinBundle.message("options.kotlin.attribute.descriptor.smart.constant") to KotlinHighlightingColors.SMART_CONSTANT,
                        KotlinBundle.message("options.kotlin.attribute.descriptor.smart.cast.receiver") to KotlinHighlightingColors.SMART_CAST_RECEIVER,
                        KotlinBundle.message("options.kotlin.attribute.descriptor.label") to KotlinHighlightingColors.LABEL,
-                       "Named argument" to KotlinHighlightingColors.NAMED_ARGUMENT)
+                       "Named argument" to KotlinHighlightingColors.NAMED_ARGUMENT) +
+                        DslHighlighterExtension.descriptionsToStyles.map { (description, key) -> description to key }.toTypedArray()
     }
 
     override fun getColorDescriptors(): Array<ColorDescriptor> = ColorDescriptor.EMPTY_ARRAY

idea/src/org/jetbrains/kotlin/idea/highlighter/dsl/DslHighlighterExtension.kt

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.highlighter.dsl
+
+import com.intellij.openapi.editor.colors.TextAttributesKey
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.idea.highlighter.HighlighterExtension
+import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.descriptorUtil.annotationClass
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import kotlin.math.absoluteValue
+
+class DslHighlighterExtension : HighlighterExtension() {
+    override fun highlightDeclaration(elementToHighlight: PsiElement, descriptor: DeclarationDescriptor): TextAttributesKey? {
+        return null
+    }
+
+    override fun highlightCall(elementToHighlight: PsiElement, resolvedCall: ResolvedCall<*>): TextAttributesKey? {
+        val markerAnnotationFqName = markerAnnotationFqName(resolvedCall) ?: return null
+        return styleForDsl(markerAnnotationFqName)
+    }
+
+    private fun markerAnnotationFqName(resolvedCall: ResolvedCall<*>): FqName? {
+        val markerAnnotation = resolvedCall.resultingDescriptor.annotations.find { annotation ->
+            annotation.annotationClass?.isDslHighlightingMarker() ?: false
+        }
+
+        return markerAnnotation?.fqName
+    }
+
+    companion object {
+        private const val numStyles = 5
+
+        private val styles = (1..numStyles).map { index ->
+            TextAttributesKey.createTextAttributesKey("KOTLIN_DSL_STYLE$index", KotlinHighlightingColors.KEYWORD)
+        }
+
+        val descriptionsToStyles = (1..numStyles).associate { index ->
+            "Dsl//Style$index" to styles[index - 1]
+        }
+
+        private fun styleForDsl(markerAnnotationFqName: FqName) =
+            styles[(markerAnnotationFqName.asString().hashCode() % numStyles).absoluteValue]
+    }
+}
+
+internal fun ClassDescriptor.isDslHighlightingMarker(): Boolean {
+    return annotations.any {
+        it.annotationClass?.fqNameSafe?.asString() == "kotlin.DslMarker"
+    }
+}