FIR IDE: Move some declarations highlighting to before resolve highlighting pass

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -28,21 +28,6 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 internal class FunctionsHighlightingVisitor(holder: AnnotationHolder, bindingContext: BindingContext) :
     AfterAnalysisHighlightingVisitor(holder, bindingContext) {
 
-    override fun visitNamedFunction(function: KtNamedFunction) {
-        function.nameIdentifier?.let { highlightName(it, FUNCTION_DECLARATION) }
-
-        super.visitNamedFunction(function)
-    }
-
-    override fun visitSuperTypeCallEntry(call: KtSuperTypeCallEntry) {
-        val calleeExpression = call.calleeExpression
-        val typeElement = calleeExpression.typeReference?.typeElement
-        if (typeElement is KtUserType) {
-            typeElement.referenceExpression?.let { highlightName(it, CONSTRUCTOR_CALL) }
-        }
-        super.visitSuperTypeCallEntry(call)
-    }
-
     override fun visitBinaryExpression(expression: KtBinaryExpression) {
         if (expression.operationReference.getIdentifier() != null) {
             val resolvedCall = expression.getResolvedCall(bindingContext)

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

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2016 JetBrains s.r.o.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- *
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
  */
 
 package org.jetbrains.kotlin.idea.highlighter
@@ -92,11 +81,6 @@ internal class TypeKindHighlightingVisitor(holder: AnnotationHolder, bindingCont
         return TextRange(atSymbol.textRange.startOffset, expression.textRange.endOffset)
     }
 
-    override fun visitTypeParameter(parameter: KtTypeParameter) {
-        parameter.nameIdentifier?.let { highlightName(it, TYPE_PARAMETER) }
-        super.visitTypeParameter(parameter)
-    }
-
     override fun visitClassOrObject(classOrObject: KtClassOrObject) {
         val identifier = classOrObject.nameIdentifier
         val classDescriptor = bindingContext.get(BindingContext.CLASS, classOrObject)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/highlighter/HighlightingUtils.kt

@@ -8,33 +8,35 @@ package org.jetbrains.kotlin.idea.highlighter
 import com.intellij.lang.jvm.JvmModifier
 import com.intellij.openapi.editor.colors.TextAttributesKey
 import com.intellij.psi.*
-import com.intellij.psi.impl.source.PsiFieldImpl
-import com.intellij.psi.util.elementType
-import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.isAbstract
 import org.jetbrains.kotlin.psi.psiUtil.isExtensionDeclaration
 import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors as Colors
 
-internal fun textAttributesKeyForPropertyDeclaration(declaration: PsiElement): TextAttributesKey? = when {
+internal fun textAttributesKeyForPropertyDeclaration(declaration: PsiElement): TextAttributesKey? = when (declaration) {
-    declaration is KtProperty && declaration.isExtensionDeclaration() -> Colors.EXTENSION_PROPERTY
+    is KtProperty -> textAttributesForKtPropertyDeclaration(declaration)
-    declaration is KtProperty && declaration.isLocal || declaration is PsiLocalVariable ->
+    is KtParameter -> textAttributesForKtParameterDeclaration(declaration)
-        Colors.LOCAL_VARIABLE
+    is PsiLocalVariable -> Colors.LOCAL_VARIABLE
-    declaration is KtParameter -> {
+    is PsiParameter -> Colors.PARAMETER
-        if (declaration.valOrVarKeyword != null) Colors.INSTANCE_PROPERTY
+    is PsiField -> Colors.INSTANCE_PROPERTY
-        else Colors.PARAMETER
+    else -> null
-    }
+}
-    declaration is PsiParameter -> Colors.PARAMETER
+
-    declaration is KtProperty && declaration.isTopLevel -> {
+internal fun textAttributesForKtParameterDeclaration(parameter: KtParameter) =
-        if (declaration.isCustomPropertyDeclaration()) Colors.PACKAGE_PROPERTY_CUSTOM_PROPERTY_DECLARATION
+    if (parameter.valOrVarKeyword != null) Colors.INSTANCE_PROPERTY
+    else Colors.PARAMETER
+
+internal fun textAttributesForKtPropertyDeclaration(property: KtProperty): TextAttributesKey? = when {
+    property.isExtensionDeclaration() -> Colors.EXTENSION_PROPERTY
+    property.isLocal -> Colors.LOCAL_VARIABLE
+    property.isTopLevel -> {
+        if (property.isCustomPropertyDeclaration()) Colors.PACKAGE_PROPERTY_CUSTOM_PROPERTY_DECLARATION
         else Colors.PACKAGE_PROPERTY
     }
-    declaration is KtProperty -> {
+    else -> {
-        if (declaration.isCustomPropertyDeclaration()) Colors.INSTANCE_PROPERTY_CUSTOM_PROPERTY_DECLARATION
+        if (property.isCustomPropertyDeclaration()) Colors.INSTANCE_PROPERTY_CUSTOM_PROPERTY_DECLARATION
         else Colors.INSTANCE_PROPERTY
     }
-    declaration is PsiField -> Colors.INSTANCE_PROPERTY
-    else -> null
 }
 
 private fun KtProperty.isCustomPropertyDeclaration() =
@@ -44,16 +46,23 @@ private fun KtProperty.isCustomPropertyDeclaration() =
 internal fun textAttributesKeyForTypeDeclaration(declaration: PsiElement): TextAttributesKey? = when {
     declaration is KtTypeParameter || declaration is PsiTypeParameter -> Colors.TYPE_PARAMETER
     declaration is KtTypeAlias -> Colors.TYPE_ALIAS
-    declaration is KtClass && declaration.isInterface()
+    declaration is KtClass -> textAttributesForClass(declaration)
-            || declaration is PsiClass && declaration.isInterface && !declaration.isAnnotationType -> Colors.TRAIT
+    declaration is PsiClass && declaration.isInterface && !declaration.isAnnotationType -> Colors.TRAIT
     declaration.isAnnotationClass() -> Colors.ANNOTATION
     declaration is KtObjectDeclaration -> Colors.OBJECT
-    declaration is KtEnumEntry || declaration is PsiEnumConstant -> Colors.ENUM_ENTRY
+    declaration is PsiEnumConstant -> Colors.ENUM_ENTRY
-    declaration is KtClass && declaration.isAbstract()
+    declaration is PsiClass && declaration.hasModifier(JvmModifier.ABSTRACT) -> Colors.ABSTRACT_CLASS
-            || declaration is PsiClass && declaration.hasModifier(JvmModifier.ABSTRACT) -> Colors.ABSTRACT_CLASS
+    declaration is PsiClass -> Colors.CLASS
-    declaration is KtClass || declaration is PsiClass -> Colors.CLASS
     else -> null
 }
 
+fun textAttributesForClass(klass: KtClass): TextAttributesKey = when {
+    klass.isInterface() -> Colors.TRAIT
+    klass.isAnnotation() -> Colors.ANNOTATION
+    klass is KtEnumEntry -> Colors.ENUM_ENTRY
+    klass.isAbstract() -> Colors.ABSTRACT_CLASS
+    else -> Colors.CLASS
+}
+
 internal fun PsiElement.isAnnotationClass() =
     this is KtClass && isAnnotation() || this is PsiClass && isAnnotationType

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/highlighter/visitors/DeclarationHighlightingVisitor.kt

@@ -6,57 +6,53 @@
 package org.jetbrains.kotlin.idea.highlighter.visitors
 
 import com.intellij.lang.annotation.AnnotationHolder
-import com.intellij.psi.util.elementType
+import com.intellij.psi.util.PsiUtilCore
-import org.jetbrains.kotlin.idea.frontend.api.FrontendAnalysisSession
+import org.jetbrains.kotlin.idea.highlighter.*
-import org.jetbrains.kotlin.idea.highlighter.textAttributesKeyForPropertyDeclaration
-import org.jetbrains.kotlin.idea.highlighter.textAttributesKeyForTypeDeclaration
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingColors as Colors
 
-internal class DeclarationHighlightingVisitor(
+internal class DeclarationHighlightingVisitor(holder: AnnotationHolder) : HighlightingVisitor(holder) {
-    analysisSession: FrontendAnalysisSession,
-    holder: AnnotationHolder
-) : FirAfterResolveHighlightingVisitor(analysisSession, holder) {
-    override fun visitNamedFunction(function: KtNamedFunction) {
-        highlightNamedDeclaration(function, Colors.FUNCTION_DECLARATION)
-        super.visitNamedFunction(function)
-    }
-
     override fun visitTypeAlias(typeAlias: KtTypeAlias) {
         highlightNamedDeclaration(typeAlias, Colors.TYPE_ALIAS)
         super.visitTypeAlias(typeAlias)
     }
 
-    override fun visitClassOrObject(classOrObject: KtClassOrObject) {
+    override fun visitObjectDeclaration(declaration: KtObjectDeclaration) {
-        textAttributesKeyForTypeDeclaration(classOrObject)?.let { attributes -> highlightNamedDeclaration(classOrObject, attributes) }
+        highlightNamedDeclaration(declaration, Colors.OBJECT)
-        super.visitClassOrObject(classOrObject)
+        super.visitObjectDeclaration(declaration)
     }
 
-    override fun visitTypeParameter(parameter: KtTypeParameter) {
+    override fun visitClass(klass: KtClass) {
-        highlightNamedDeclaration(parameter, Colors.TYPE_PARAMETER)
+        highlightNamedDeclaration(klass, textAttributesForClass(klass))
-        super.visitTypeParameter(parameter)
+        super.visitClass(klass)
     }
 
-    override fun visitNamedDeclaration(declaration: KtNamedDeclaration) {
+    override fun visitProperty(property: KtProperty) {
-        if (declaration is KtValVarKeywordOwner) {
+        textAttributesForKtPropertyDeclaration(property)?.let { attributes ->
-            textAttributesKeyForPropertyDeclaration(declaration)?.let { attributes ->
+            highlightNamedDeclaration(property, attributes)
-                highlightNamedDeclaration(declaration, attributes)
-            }
-            if (declaration.valOrVarKeyword?.elementType == KtTokens.VAR_KEYWORD) {
-                highlightNamedDeclaration(declaration, Colors.MUTABLE_VARIABLE)
-            }
         }
-        super.visitNamedDeclaration(declaration)
+        highlightMutability(property)
+        super.visitProperty(property)
     }
 
-
+    override fun visitParameter(parameter: KtParameter) {
-    override fun visitSuperTypeCallEntry(call: KtSuperTypeCallEntry) {
+        textAttributesForKtParameterDeclaration(parameter)?.let { attributes ->
-        val calleeExpression = call.calleeExpression
+            highlightNamedDeclaration(parameter, attributes)
-        val typeElement = calleeExpression.typeReference?.typeElement
-        if (typeElement is KtUserType) {
-            typeElement.referenceExpression?.let { highlightName(it, Colors.CONSTRUCTOR_CALL) }
         }
-        super.visitSuperTypeCallEntry(call)
+        highlightMutability(parameter)
+        super.visitParameter(parameter)
     }
+
+
+    private fun <D> highlightMutability(declaration: D) where D : KtValVarKeywordOwner, D : KtNamedDeclaration {
+        if (PsiUtilCore.getElementType(declaration.valOrVarKeyword) == KtTokens.VAR_KEYWORD) {
+            highlightNamedDeclaration(declaration, Colors.MUTABLE_VARIABLE)
+        }
+    }
+}
+
+class DeclarationHighlightingExtension : BeforeResolveHighlightingExtension {
+    override fun createVisitor(holder: AnnotationHolder): HighlightingVisitor =
+        DeclarationHighlightingVisitor(holder)
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/highlighter/visitors/FirAfterResolveHighlightingVisitor.kt

@@ -14,17 +14,12 @@ abstract class FirAfterResolveHighlightingVisitor(
     protected val holder: AnnotationHolder
 ) : HighlightingVisitor(holder) {
 
-    protected fun highlightNamedDeclaration(declaration: KtNamedDeclaration, key: TextAttributesKey) {
-        declaration.nameIdentifier?.let { highlightName(it, key) }
-    }
-
     companion object {
         fun createListOfVisitors(
             analysisSession: FrontendAnalysisSession,
             holder: AnnotationHolder
         ): List<FirAfterResolveHighlightingVisitor> = listOf(
             TypeHighlightingVisitor(analysisSession, holder),
-            DeclarationHighlightingVisitor(analysisSession, holder),
             FunctionCallHighlightingVisitor(analysisSession, holder),
             ExpressionsSmartcastHighlightingVisitor(analysisSession, holder),
             VariableReferenceHighlightingVisitor(analysisSession, holder),

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/highlighter/BeforeResolveHighlightingVisitor.kt

@@ -16,10 +16,7 @@ import org.jetbrains.kotlin.kdoc.parser.KDocKnownTag
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocLink
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocTag
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtExpressionWithLabel
-import org.jetbrains.kotlin.psi.KtLambdaExpression
-import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
@@ -73,7 +70,7 @@ internal class BeforeResolveHighlightingVisitor(holder: AnnotationHolder) : High
         val argumentName = argument.getArgumentName() ?: return
         val eq = argument.equalsToken ?: return
         createInfoAnnotation(TextRange(argumentName.startOffset, eq.endOffset), null).textAttributes =
-            if(argument.parent.parent is KtAnnotationEntry)
+            if (argument.parent.parent is KtAnnotationEntry)
                 KotlinHighlightingColors.ANNOTATION_ATTRIBUTE_NAME_ATTRIBUTES
             else
                 KotlinHighlightingColors.NAMED_ARGUMENT
@@ -85,4 +82,24 @@ internal class BeforeResolveHighlightingVisitor(holder: AnnotationHolder) : High
             highlightName(targetLabel, KotlinHighlightingColors.LABEL)
         }
     }
+
+    override fun visitSuperTypeCallEntry(call: KtSuperTypeCallEntry) {
+        val calleeExpression = call.calleeExpression
+        val typeElement = calleeExpression.typeReference?.typeElement
+        if (typeElement is KtUserType) {
+            typeElement.referenceExpression?.let { highlightName(it, KotlinHighlightingColors.CONSTRUCTOR_CALL) }
+        }
+        super.visitSuperTypeCallEntry(call)
+    }
+
+
+    override fun visitTypeParameter(parameter: KtTypeParameter) {
+        parameter.nameIdentifier?.let { highlightName(it, KotlinHighlightingColors.TYPE_PARAMETER) }
+        super.visitTypeParameter(parameter)
+    }
+
+    override fun visitNamedFunction(function: KtNamedFunction) {
+        highlightNamedDeclaration(function, KotlinHighlightingColors.FUNCTION_DECLARATION)
+        super.visitNamedFunction(function)
+    }
 }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/highlighter/HighlightingVisitor.kt

@@ -10,6 +10,7 @@ import com.intellij.lang.annotation.AnnotationHolder
 import com.intellij.openapi.editor.colors.TextAttributesKey
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.KtNamedDeclaration
 import org.jetbrains.kotlin.psi.KtVisitorVoid
 
 abstract class HighlightingVisitor protected constructor(
@@ -41,4 +42,8 @@ abstract class HighlightingVisitor protected constructor(
             createInfoAnnotation(textRange, attributesKey, message)
         }
     }
+
+    protected fun highlightNamedDeclaration(declaration: KtNamedDeclaration, attributesKey: TextAttributesKey) {
+        declaration.nameIdentifier?.let { highlightName(it, attributesKey) }
+    }
 }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/highlighter/KotlinBeforeResolveHighlightingPass.kt

@@ -13,7 +13,6 @@ import com.intellij.lang.annotation.AnnotationSession
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.extensions.ExtensionPointName
-import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.project.DumbAware
 import com.intellij.openapi.project.Project
@@ -21,7 +20,6 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiRecursiveElementVisitor
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 class KotlinBeforeResolveHighlightingPass(
     private val file: KtFile,
@@ -34,10 +32,12 @@ class KotlinBeforeResolveHighlightingPass(
     override fun doCollectInformation(progress: ProgressIndicator) {
         val annotationHolder = AnnotationHolderImpl(AnnotationSession(file))
         val visitor = BeforeResolveHighlightingVisitor(annotationHolder)
+        val extensions = EP_NAME.extensionList
         file.accept(object : PsiRecursiveElementVisitor() {
             override fun visitElement(element: PsiElement) {
                 super.visitElement(element)
                 element.accept(visitor)
+                extensions.forEach(element::accept)
             }
         })
         this.annotationHolder = annotationHolder
@@ -71,7 +71,7 @@ class KotlinBeforeResolveHighlightingPass(
         }
     }
 
-//    companion object {
+    companion object {
-//        val EP_NAME = ExtensionPointName.create<HighlightingVisitor>("org.jetbrains.kotlin.beforeResolveHighlightingVisitor")
+        val EP_NAME = ExtensionPointName.create<HighlightingVisitor>("org.jetbrains.kotlin.beforeResolveHighlightingVisitor")
-//    }
+    }
 }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/highlighter/KotlinBeforeResolveHighlightingPass.kt.192

@@ -13,6 +13,7 @@ import com.intellij.codeInsight.daemon.impl.AnnotationHolderImpl
 import com.intellij.codeInsight.daemon.impl.HighlightInfo
 import com.intellij.codeInsight.daemon.impl.UpdateHighlightersUtil
 import com.intellij.lang.annotation.AnnotationSession
+import com.intellij.lang.annotation.AnnotationHolder
 import com.intellij.openapi.components.ProjectComponent
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
@@ -34,10 +35,12 @@ class KotlinBeforeResolveHighlightingPass(
     override fun doCollectInformation(progress: ProgressIndicator) {
         val annotationHolder = AnnotationHolderImpl(AnnotationSession(file))
         val visitor = BeforeResolveHighlightingVisitor(annotationHolder)
+        val extensions = EP_NAME.extensionList
         file.accept(object : PsiRecursiveElementVisitor() {
             override fun visitElement(element: PsiElement) {
                 super.visitElement(element)
                 element.accept(visitor)
+                extensions.forEach(element::accept)
             }
         })
         this.annotationHolder = annotationHolder
@@ -68,4 +71,8 @@ class KotlinBeforeResolveHighlightingPass(
             return KotlinBeforeResolveHighlightingPass(file, editor.document)
         }
     }
+
+     companion object {
+        val EP_NAME = ExtensionPointName.create<HighlightingVisitor>("org.jetbrains.kotlin.beforeResolveHighlightingVisitor")
+    }
 }