FIR: Prepare ConeAttribute infrastructure for non-singleton attrs

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/CompilerConeAttributes.kt

@@ -81,20 +81,21 @@ object CompilerConeAttributes {
         override fun toString(): String = "@UnsafeVariance"
     }
 
-    val compilerAttributeByClassId: Map<ClassId, ConeAttribute<*>> = mapOf(
-        Exact.ANNOTATION_CLASS_ID to Exact,
-        NoInfer.ANNOTATION_CLASS_ID to NoInfer,
-        EnhancedNullability.ANNOTATION_CLASS_ID to EnhancedNullability,
-        ExtensionFunctionType.ANNOTATION_CLASS_ID to ExtensionFunctionType,
-        UnsafeVariance.ANNOTATION_CLASS_ID to UnsafeVariance
+    private val compilerAttributeByClassId: Map<ClassId, ConeAttributeKey> = mapOf(
+        Exact.ANNOTATION_CLASS_ID to Exact.key,
+        NoInfer.ANNOTATION_CLASS_ID to NoInfer.key,
+        EnhancedNullability.ANNOTATION_CLASS_ID to EnhancedNullability.key,
+        ExtensionFunctionType.ANNOTATION_CLASS_ID to ExtensionFunctionType.key,
+        UnsafeVariance.ANNOTATION_CLASS_ID to UnsafeVariance.key
     )
 
-    val classIdByCompilerAttribute: Map<ConeAttribute<*>, ClassId> = compilerAttributeByClassId.entries.associateBy(
+    val classIdByCompilerAttributeKey: Map<ConeAttributeKey, ClassId> = compilerAttributeByClassId.entries.associateBy(
         keySelector = { it.value },
         valueTransform = { it.key }
     )
 
-    val compilerAttributeByFqName: Map<FqName, ConeAttribute<*>> = compilerAttributeByClassId.mapKeys { it.key.asSingleFqName() }
+    val compilerAttributeKeyByFqName: Map<FqName, ConeAttributeKey> =
+        compilerAttributeByClassId.mapKeys { it.key.asSingleFqName() }
 }
 
 val ConeAttributes.exact: CompilerConeAttributes.Exact? by ConeAttributes.attributeAccessor<CompilerConeAttributes.Exact>()

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt

@@ -6,9 +6,9 @@
 package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.fir.util.ConeTypeRegistry
+import org.jetbrains.kotlin.types.model.AnnotationMarker
 import org.jetbrains.kotlin.util.AttributeArrayOwner
 import org.jetbrains.kotlin.util.TypeRegistry
-import org.jetbrains.kotlin.types.model.AnnotationMarker
 import org.jetbrains.kotlin.utils.addIfNotNull
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KClass
@@ -34,6 +34,8 @@ abstract class ConeAttribute<T : ConeAttribute<T>> : AnnotationMarker {
     abstract val key: KClass<out T>
 }
 
+typealias ConeAttributeKey = KClass<out ConeAttribute<*>>
+
 class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : AttributeArrayOwner<ConeAttribute<*>, ConeAttribute<*>>(),
     Iterable<ConeAttribute<*>> {
 
@@ -82,7 +84,11 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
     }
 
     operator fun contains(attribute: ConeAttribute<*>): Boolean {
-        val index = getId(attribute.key)
+        return contains(attribute.key)
+    }
+
+    operator fun contains(attributeKey: KClass<out ConeAttribute<*>>): Boolean {
+        val index = getId(attributeKey)
         return arrayMap[index] != null
     }
 

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirElementSerializer.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.serialization
 
+import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
@@ -17,11 +18,10 @@ import org.jetbrains.kotlin.fir.declarations.comparators.FirCallableDeclarationC
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.utils.*
 import org.jetbrains.kotlin.fir.deserialization.projection
-import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirArgumentList
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirNamedArgumentExpression
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
+import org.jetbrains.kotlin.fir.expressions.builder.buildConstExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyAnnotationArgumentMapping
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.typeAttributeExtensions
@@ -51,6 +51,7 @@ import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.resolve.RequireKotlinConstants
 import org.jetbrains.kotlin.serialization.deserialization.ProtoEnumFlags
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
+import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 
 class FirElementSerializer private constructor(
@@ -597,9 +598,9 @@ class FirElementSerializer private constructor(
             val compilerAttributes = mutableListOf<ConeAttribute<*>>()
             val extensionAttributes = mutableListOf<ConeAttribute<*>>()
             for (attribute in coneType.attributes) {
-                when (attribute) {
-                    is CustomAnnotationTypeAttribute -> continue
-                    in CompilerConeAttributes.classIdByCompilerAttribute -> compilerAttributes += attribute
+                when {
+                    attribute is CustomAnnotationTypeAttribute -> continue
+                    attribute.key in CompilerConeAttributes.classIdByCompilerAttributeKey -> compilerAttributes += attribute
                     else -> extensionAttributes += attribute
                 }
             }
@@ -608,7 +609,7 @@ class FirElementSerializer private constructor(
                 val annotation = buildAnnotation {
                     annotationTypeRef = buildResolvedTypeRef {
                         type = ConeClassLikeTypeImpl(
-                            ConeClassLikeLookupTagImpl(CompilerConeAttributes.classIdByCompilerAttribute.getValue(attribute)),
+                            ConeClassLikeLookupTagImpl(CompilerConeAttributes.classIdByCompilerAttributeKey.getValue(attribute.key)),
                             emptyArray(),
                             isNullable = false
                         )
@@ -661,11 +662,6 @@ class FirElementSerializer private constructor(
                     return functionType
                 }
                 fillFromPossiblyInnerType(builder, type)
-                if (type.isExtensionFunctionType) {
-                    serializeAnnotationFromAttribute(
-                        correspondingTypeRef?.annotations, CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID, builder
-                    )
-                }
             }
             is ConeTypeParameterType -> {
                 val typeParameter = type.lookupTag.typeParameterSymbol.fir

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/ConeInferenceContext.kt

@@ -408,8 +408,8 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
     }
 
     private fun AnnotationMarker.isCustomAttribute(): Boolean {
-        val compilerAttributes = CompilerConeAttributes.classIdByCompilerAttribute
-        return this !in compilerAttributes && this !is CustomAnnotationTypeAttribute
+        val compilerAttributes = CompilerConeAttributes.classIdByCompilerAttributeKey
+        return (this as? ConeAttribute<*>)?.key !in compilerAttributes && this !is CustomAnnotationTypeAttribute
     }
 
     override fun KotlinTypeMarker.replaceCustomAttributes(newAttributes: List<AnnotationMarker>): KotlinTypeMarker {

compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -429,14 +429,14 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 
     override fun KotlinTypeMarker.hasCustomAttributes(): Boolean {
         require(this is ConeKotlinType)
-        val compilerAttributes = CompilerConeAttributes.classIdByCompilerAttribute
-        return this.attributes.any { it !in compilerAttributes && it !is CustomAnnotationTypeAttribute }
+        val compilerAttributes = CompilerConeAttributes.classIdByCompilerAttributeKey
+        return this.attributes.any { it.key !in compilerAttributes && it !is CustomAnnotationTypeAttribute }
     }
 
     override fun KotlinTypeMarker.getCustomAttributes(): List<AnnotationMarker> {
         require(this is ConeKotlinType)
-        val compilerAttributes = CompilerConeAttributes.classIdByCompilerAttribute
-        return attributes.filterNot { it in compilerAttributes || it is CustomAnnotationTypeAttribute }
+        val compilerAttributes = CompilerConeAttributes.classIdByCompilerAttributeKey
+        return attributes.filterNot { it.key in compilerAttributes || it is CustomAnnotationTypeAttribute }
     }
 
     override fun SimpleTypeMarker.isStubType(): Boolean {
@@ -520,7 +520,7 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 
     override fun KotlinTypeMarker.hasAnnotation(fqName: FqName): Boolean {
         require(this is ConeKotlinType)
-        val compilerAttribute = CompilerConeAttributes.compilerAttributeByFqName[fqName]
+        val compilerAttribute = CompilerConeAttributes.compilerAttributeKeyByFqName[fqName]
         if (compilerAttribute != null) {
             return compilerAttribute in attributes
         }