[FIR] Implement extension for providing additional type attributes

compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirTypeAttributeExtension.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2021 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.
+ */
+
+package org.jetbrains.kotlin.fir.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.types.ConeAttribute
+import kotlin.reflect.KClass
+
+/**
+ * Note that when you declare your own attribute you HAVE TO declare accessor to it using following syntax:
+ *
+ * class MyAttribute : ConeAttribute<MyAttribute>() {...}
+ *
+ * val ConeAttributes.myAttribute: MyAttribute? by ConeAttributes.attributeAccessor<MyAttribute>()
+ */
+abstract class FirTypeAttributeExtension(session: FirSession) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("AdditionalTypeAttributeExtension")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension>
+        get() = FirTypeAttributeExtension::class
+
+    abstract fun extractAttributeFromAnnotation(annotation: FirAnnotation): ConeAttribute<*>?
+
+    /**
+     * Please don't convert attributes which you didn't create
+     * If [attribute] came from compiler or another plugin just return null
+     */
+    abstract fun convertAttributeToAnnotation(attribute: ConeAttribute<*>): FirAnnotation?
+
+    fun interface Factory : FirExtension.Factory<FirTypeAttributeExtension>
+}
+
+val FirExtensionService.typeAttributeExtensions: List<FirTypeAttributeExtension> by FirExtensionService.registeredExtensions()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirTypeUtils.kt

@@ -6,10 +6,13 @@
 package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirConstExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.extensions.typeAttributeExtensions
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.name.ClassId
@@ -109,26 +112,25 @@ fun ConeClassLikeType.toConstKind(): ConstantValueKind<*>? = when (lookupTag.cla
     else -> null
 }
 
-fun List<FirAnnotation>.computeTypeAttributes(
-    additionalProcessor: MutableList<ConeAttribute<*>>.(ClassId) -> Unit = {}
-): ConeAttributes {
+fun List<FirAnnotation>.computeTypeAttributes(session: FirSession): ConeAttributes {
     if (this.isEmpty()) return ConeAttributes.Empty
     val attributes = mutableListOf<ConeAttribute<*>>()
     val customAnnotations = mutableListOf<FirAnnotation>()
     for (annotation in this) {
         val type = annotation.annotationTypeRef.coneTypeSafe<ConeClassLikeType>() ?: continue
-        when (val classId = type.lookupTag.classId) {
+        when (type.lookupTag.classId) {
             CompilerConeAttributes.Exact.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.Exact
             CompilerConeAttributes.NoInfer.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.NoInfer
             CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.ExtensionFunctionType
             CompilerConeAttributes.UnsafeVariance.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.UnsafeVariance
             else -> {
-                val annotationAttributes = mutableListOf<ConeAttribute<*>>()
-                additionalProcessor.invoke(annotationAttributes, classId)
-                if (annotationAttributes.isEmpty()) {
-                    customAnnotations += annotation
+                val attributeFromPlugin = session.extensionService.typeAttributeExtensions.firstNotNullOfOrNull {
+                    it.extractAttributeFromAnnotation(annotation)
+                }
+                if (attributeFromPlugin != null) {
+                    attributes += attributeFromPlugin
                 } else {
-                    attributes += annotationAttributes
+                    customAnnotations += annotation
                 }
             }
         }