[FIR] Fix enhancement of FlexibleNullability and EnhancedNullability

- Add utilities to add new attribute to ConeAttributes
- Get rid of FlexibleNullability attribute (it can be easily inferred
    for any flexible type at any moment)
- Fix determining of EnhancedNullability attribute

compiler/fir/cones/src/org/jetbrains/kotlin/fir/symbols/StandardClassIds.kt

@@ -108,6 +108,9 @@ object StandardClassIds {
     val MutableMapEntry = MutableMap.createNestedClassId(Name.identifier("MutableEntry"))
 
     val Suppress = "Suppress".baseId()
+
+    val FlexibleNullability = ClassId(FqName("kotlin.internal.ir"), Name.identifier("FlexibleNullability"))
+    val EnhancedNullability = ClassId(FqName("kotlin.jvm.internal"), Name.identifier("EnhancedNullability"))
 }
 
 private fun <K, V> Map<K, V>.inverseMap() = entries.associate { (k, v) -> v to k }

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -36,7 +37,7 @@ object CompilerConeAttributes {
     }
 
     object EnhancedNullability : ConeAttribute<EnhancedNullability>() {
-        val ANNOTATION_CLASS_ID = ClassId(FqName("kotlin.jvm.internal"), Name.identifier("EnhancedNullability"))
+        val ANNOTATION_CLASS_ID = StandardClassIds.EnhancedNullability
 
         override fun union(other: EnhancedNullability?): EnhancedNullability? = other
         override fun intersect(other: EnhancedNullability?): EnhancedNullability = this
@@ -59,18 +60,6 @@ object CompilerConeAttributes {
         override fun toString(): String = "@ExtensionFunctionType"
     }
 
-    object FlexibleNullability : ConeAttribute<FlexibleNullability>() {
-        val ANNOTATION_CLASS_ID = ClassId(FqName("kotlin.internal.ir"), Name.identifier("FlexibleNullability"))
-
-        override fun union(other: FlexibleNullability?): FlexibleNullability? = other
-        override fun intersect(other: FlexibleNullability?): FlexibleNullability = this
-        override fun isSubtypeOf(other: FlexibleNullability?): Boolean = true
-
-        override val key: KClass<out FlexibleNullability> = FlexibleNullability::class
-
-        override fun toString(): String = "@FlexibleNullability"
-    }
-
     object UnsafeVariance : ConeAttribute<UnsafeVariance>() {
         val ANNOTATION_CLASS_ID = ClassId(FqName("kotlin"), Name.identifier("UnsafeVariance"))
 
@@ -88,7 +77,6 @@ object CompilerConeAttributes {
         NoInfer.ANNOTATION_CLASS_ID to NoInfer,
         EnhancedNullability.ANNOTATION_CLASS_ID to EnhancedNullability,
         ExtensionFunctionType.ANNOTATION_CLASS_ID to ExtensionFunctionType,
-        FlexibleNullability.ANNOTATION_CLASS_ID to FlexibleNullability,
         UnsafeVariance.ANNOTATION_CLASS_ID to UnsafeVariance
     )
 
@@ -99,14 +87,17 @@ val ConeAttributes.exact: CompilerConeAttributes.Exact? by ConeAttributes.attrib
 val ConeAttributes.noInfer: CompilerConeAttributes.NoInfer? by ConeAttributes.attributeAccessor<CompilerConeAttributes.NoInfer>()
 val ConeAttributes.enhancedNullability: CompilerConeAttributes.EnhancedNullability? by ConeAttributes.attributeAccessor<CompilerConeAttributes.EnhancedNullability>()
 val ConeAttributes.extensionFunctionType: CompilerConeAttributes.ExtensionFunctionType? by ConeAttributes.attributeAccessor<CompilerConeAttributes.ExtensionFunctionType>()
-val ConeAttributes.flexibleNullability: CompilerConeAttributes.FlexibleNullability? by ConeAttributes.attributeAccessor<CompilerConeAttributes.FlexibleNullability>()
 val ConeAttributes.unsafeVarianceType: CompilerConeAttributes.UnsafeVariance? by ConeAttributes.attributeAccessor<CompilerConeAttributes.UnsafeVariance>()
 
+// ------------------------------------------------------------------
+
+val ConeAttributes.hasEnhancedNullability: Boolean
+    get() = enhancedNullability != null
+
+// ------------------------------------------------------------------
+
 val ConeKotlinType.hasEnhancedNullability: Boolean
     get() = attributes.enhancedNullability != null
 
 val ConeKotlinType.isExtensionFunctionType: Boolean
     get() = attributes.extensionFunctionType != null
-
-val ConeKotlinType.hasFlexibleNullability: Boolean
-    get() = attributes.flexibleNullability != null

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

@@ -36,7 +36,12 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
 
         val Empty: ConeAttributes = ConeAttributes(emptyList())
         val WithExtensionFunctionType: ConeAttributes = ConeAttributes(listOf(CompilerConeAttributes.ExtensionFunctionType))
-        internal val WithFlexibleNullability: ConeAttributes = ConeAttributes(listOf(CompilerConeAttributes.FlexibleNullability))
+
+        private val predefinedAttributes: Map<ConeAttribute<*>, ConeAttributes> = mapOf(
+            CompilerConeAttributes.EnhancedNullability.predefined()
+        )
+
+        private fun ConeAttribute<*>.predefined(): Pair<ConeAttribute<*>, ConeAttributes> = this to ConeAttributes(this)
 
         fun create(attributes: List<ConeAttribute<*>>): ConeAttributes {
             return if (attributes.isEmpty()) {
@@ -47,21 +52,14 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
         }
     }
 
+    private constructor(attribute: ConeAttribute<*>) : this(listOf(attribute))
+
     init {
         for (attribute in attributes) {
             registerComponent(attribute.key, attribute)
         }
-        assert(!hasEnhancedNullability || !hasFlexibleNullability) {
-            "It doesn't make sense to have @EnhancedNullability and @FlexibleNullability at the same time."
-        }
     }
 
-    val hasEnhancedNullability: Boolean
-        get() = enhancedNullability != null
-
-    private val hasFlexibleNullability: Boolean
-        get() = flexibleNullability != null
-
     fun union(other: ConeAttributes): ConeAttributes {
         return perform(other) { this.union(it) }
     }
@@ -70,6 +68,22 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
         return perform(other) { this.intersect(it) }
     }
 
+    operator fun contains(attribute: ConeAttribute<*>): Boolean {
+        val index = getId(attribute.key)
+        return arrayMap[index] != null
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    operator fun plus(attribute: ConeAttribute<*>): ConeAttributes {
+        if (attribute in this) return this
+        if (isEmpty()) return predefinedAttributes[attribute] ?: ConeAttributes(attribute)
+        val newAttributes = buildList {
+            addAll(this)
+            add(attribute)
+        }
+        return ConeAttributes(newAttributes)
+    }
+
     fun remove(attribute: ConeAttribute<*>): ConeAttributes {
         if (arrayMap.isEmpty()) return this
         val attributes = arrayMap.filter { it != attribute }
@@ -100,12 +114,3 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
         return arrayMap.isEmpty()
     }
 }
-
-private fun ConeAttributes.intersectUnless(other: ConeAttributes, predicate: (ConeAttributes) -> Boolean): ConeAttributes =
-    if (predicate.invoke(this)) this else intersect(other)
-
-fun ConeAttributes.withFlexible(): ConeAttributes =
-    intersect(ConeAttributes.WithFlexibleNullability)
-
-fun ConeAttributes.withFlexibleUnless(predicate: (ConeAttributes) -> Boolean): ConeAttributes =
-    intersectUnless(ConeAttributes.WithFlexibleNullability, predicate)