From 49f8f0af92fa48ababa0a414fcf1737283c73e60 Mon Sep 17 00:00:00 2001
From: Alexander Udalov <Alexander.Udalov@jetbrains.com>
Date: Thu, 21 Nov 2013 20:58:25 +0400
Subject: [PATCH] Fix highlighting for objects and enum entries

JetObjectDeclarationName is now highlighted with KOTLIN_CLASS
---
 .../TypeKindHighlightingVisitor.java          | 48 +++++++++++++------
 idea/testData/highlighter/Object.kt           | 11 +++++
 .../HighlightingTestGenerated.java            |  5 ++
 3 files changed, 49 insertions(+), 15 deletions(-)
 create mode 100644 idea/testData/highlighter/Object.kt

diff --git a/idea/src/org/jetbrains/jet/plugin/highlighter/TypeKindHighlightingVisitor.java b/idea/src/org/jetbrains/jet/plugin/highlighter/TypeKindHighlightingVisitor.java
index 1def1ac26cf..825cc057ff2 100644
--- a/idea/src/org/jetbrains/jet/plugin/highlighter/TypeKindHighlightingVisitor.java
+++ b/idea/src/org/jetbrains/jet/plugin/highlighter/TypeKindHighlightingVisitor.java
@@ -41,47 +41,65 @@ class TypeKindHighlightingVisitor extends AfterAnalysisHighlightingVisitor {
             }
 
             if (referenceTarget instanceof ClassDescriptor) {
-                highlightClassByKind((ClassDescriptor) referenceTarget, expression);
+                highlightName(expression, textAttributesKeyForClass((ClassDescriptor) referenceTarget));
             }
             else if (referenceTarget instanceof TypeParameterDescriptor) {
-                JetPsiChecker.highlightName(holder, expression, JetHighlightingColors.TYPE_PARAMETER);
+                highlightName(expression, JetHighlightingColors.TYPE_PARAMETER);
             }
         }
     }
 
+    @Override
+    public void visitObjectDeclarationName(@NotNull JetObjectDeclarationName declaration) {
+        PsiElement nameIdentifier = declaration.getNameIdentifier();
+        if (nameIdentifier != null) {
+            highlightName(nameIdentifier, JetHighlightingColors.CLASS);
+        }
+    }
+
     @Override
     public void visitTypeParameter(@NotNull JetTypeParameter parameter) {
         PsiElement identifier = parameter.getNameIdentifier();
         if (identifier != null) {
-            JetPsiChecker.highlightName(holder, identifier, JetHighlightingColors.TYPE_PARAMETER);
+            highlightName(identifier, JetHighlightingColors.TYPE_PARAMETER);
         }
         super.visitTypeParameter(parameter);
     }
 
+    @Override
+    public void visitEnumEntry(@NotNull JetEnumEntry enumEntry) {
+        // Do nothing, the name was already highlighted in visitObjectDeclarationName
+    }
+
     @Override
     public void visitClass(@NotNull JetClass klass) {
         PsiElement identifier = klass.getNameIdentifier();
         ClassDescriptor classDescriptor = bindingContext.get(BindingContext.CLASS, klass);
         if (identifier != null && classDescriptor != null) {
-            highlightClassByKind(classDescriptor, identifier);
+            highlightName(identifier, textAttributesKeyForClass(classDescriptor));
         }
         super.visitClass(klass);
     }
 
-    private void highlightClassByKind(@NotNull ClassDescriptor classDescriptor, @NotNull PsiElement whatToHighlight) {
-        TextAttributesKey textAttributes;
-        switch (classDescriptor.getKind()) {
+    private void highlightName(@NotNull PsiElement whatToHighlight, @NotNull TextAttributesKey textAttributesKey) {
+        JetPsiChecker.highlightName(holder, whatToHighlight, textAttributesKey);
+    }
+
+    @NotNull
+    private static TextAttributesKey textAttributesKeyForClass(@NotNull ClassDescriptor descriptor) {
+        switch (descriptor.getKind()) {
             case TRAIT:
-                textAttributes = JetHighlightingColors.TRAIT;
-                break;
+                return JetHighlightingColors.TRAIT;
             case ANNOTATION_CLASS:
-                textAttributes = JetHighlightingColors.ANNOTATION;
-                break;
+                return JetHighlightingColors.ANNOTATION;
+            case OBJECT:
+            case ENUM_ENTRY:
+                // Treat object accesses as variables to simulate the old behaviour (when variables were created for objects)
+                return JetHighlightingColors.INSTANCE_PROPERTY;
             default:
-                textAttributes = classDescriptor.getModality() == Modality.ABSTRACT
-                                 ? JetHighlightingColors.ABSTRACT_CLASS
-                                 : JetHighlightingColors.CLASS;
+                return descriptor.getModality() == Modality.ABSTRACT
+                       ? JetHighlightingColors.ABSTRACT_CLASS
+                       : JetHighlightingColors.CLASS;
         }
-        JetPsiChecker.highlightName(holder, whatToHighlight, textAttributes);
     }
 }
diff --git a/idea/testData/highlighter/Object.kt b/idea/testData/highlighter/Object.kt
new file mode 100644
index 00000000000..dde53708c51
--- /dev/null
+++ b/idea/testData/highlighter/Object.kt
@@ -0,0 +1,11 @@
+package testing
+
+object <info textAttributesKey="KOTLIN_CLASS">O</info> {
+    fun <info textAttributesKey="KOTLIN_FUNCTION_DECLARATION">foo</info>() = 42
+}
+
+fun <info textAttributesKey="KOTLIN_FUNCTION_DECLARATION">testing</info>() {
+    <info textAttributesKey="KOTLIN_INSTANCE_PROPERTY">O</info>.<info textAttributesKey="KOTLIN_FUNCTION_CALL">foo</info>()
+    val <info textAttributesKey="KOTLIN_LOCAL_VARIABLE">o</info> = <info textAttributesKey="KOTLIN_INSTANCE_PROPERTY">O</info>
+    <info textAttributesKey="KOTLIN_LOCAL_VARIABLE">o</info>.<info textAttributesKey="KOTLIN_FUNCTION_CALL">foo</info>()
+}
diff --git a/idea/tests/org/jetbrains/jet/plugin/highlighter/HighlightingTestGenerated.java b/idea/tests/org/jetbrains/jet/plugin/highlighter/HighlightingTestGenerated.java
index f32019177c0..9540015cbcd 100644
--- a/idea/tests/org/jetbrains/jet/plugin/highlighter/HighlightingTestGenerated.java
+++ b/idea/tests/org/jetbrains/jet/plugin/highlighter/HighlightingTestGenerated.java
@@ -52,6 +52,11 @@ public class HighlightingTestGenerated extends AbstractHighlightingTest {
         doTest("idea/testData/highlighter/JavaTypes.kt");
     }
     
+    @TestMetadata("Object.kt")
+    public void testObject() throws Exception {
+        doTest("idea/testData/highlighter/Object.kt");
+    }
+    
     @TestMetadata("TypesAndAnnotations.kt")
     public void testTypesAndAnnotations() throws Exception {
         doTest("idea/testData/highlighter/TypesAndAnnotations.kt");