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
     }