[FIR] Implement a special type of attribute that contains a cone type that can be transformed

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

@@ -8,13 +8,14 @@ package org.jetbrains.kotlin.fir.types
 import kotlin.reflect.KClass
 
 class AbbreviatedTypeAttribute(
-    val coneType: ConeKotlinType,
-) : ConeAttribute<AbbreviatedTypeAttribute>() {
+    override val coneType: ConeKotlinType,
+) : ConeAttributeWithConeType<AbbreviatedTypeAttribute>() {
     override fun union(other: AbbreviatedTypeAttribute?): AbbreviatedTypeAttribute? = null
     override fun intersect(other: AbbreviatedTypeAttribute?): AbbreviatedTypeAttribute? = null
     override fun add(other: AbbreviatedTypeAttribute?): AbbreviatedTypeAttribute = other ?: this
     override fun isSubtypeOf(other: AbbreviatedTypeAttribute?): Boolean = true
     override fun toString(): String = "{${coneType.renderForDebugging()}=}"
+    override fun copyWith(newType: ConeKotlinType): AbbreviatedTypeAttribute = AbbreviatedTypeAttribute(newType)
 
     override val key: KClass<out AbbreviatedTypeAttribute>
         get() = AbbreviatedTypeAttribute::class

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

@@ -40,6 +40,25 @@ abstract class ConeAttribute<out T : ConeAttribute<T>> : AnnotationMarker {
     abstract val keepInInferredDeclarationType: Boolean
 }
 
+/**
+ * An attribute that contains a [ConeKotlinType].
+ * It is assumed that the [coneType] in this attribute is somehow related to the type it is attached to.
+ * Therefore, when the type is transformed (e.g., substitution, making not-null), the same transformation is applied to the attribute.
+ */
+abstract class ConeAttributeWithConeType<out T : ConeAttributeWithConeType<T>> : ConeAttribute<T>() {
+    abstract val coneType: ConeKotlinType
+
+    abstract fun copyWith(newType: ConeKotlinType): T
+}
+
+inline fun <T : ConeAttributeWithConeType<T>> ConeAttributeWithConeType<T>.transformOrNull(
+    transform: (ConeKotlinType) -> ConeKotlinType?,
+): ConeAttributeWithConeType<T>? {
+    val transformedType = transform(coneType) ?: return null
+    if (transformedType == coneType) return this
+    return copyWith(transformedType)
+}
+
 typealias ConeAttributeKey = KClass<out ConeAttribute<*>>
 
 class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : AttributeArrayOwner<ConeAttribute<*>, ConeAttribute<*>>(),
@@ -116,15 +135,6 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
         return create(attributes)
     }
 
-    fun replace(oldAttribute: ConeAttribute<*>, newAttribute: ConeAttribute<*>): ConeAttributes {
-        return create(buildList {
-            arrayMap.mapNotNullTo(this) { attr ->
-                attr.takeUnless { it == oldAttribute }
-            }
-            add(newAttribute)
-        })
-    }
-
     fun filterNecessaryToKeep(): ConeAttributes {
         return if (all { it.keepInInferredDeclarationType }) this
         else create(filter { it.keepInInferredDeclarationType })
@@ -142,6 +152,34 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
         return create(attributes)
     }
 
+    /**
+     * Applies the [transform] to all attributes that are subtypes of [ConeAttributeWithConeType] and returns a [ConeAttributes]
+     * with the results of transforms that were not-`null` or `null` if no attributes were transformed.
+     */
+    inline fun transformTypesWith(transform: (ConeKotlinType) -> ConeKotlinType?): ConeAttributes? {
+        if (isEmpty()) return null
+
+        // List will be allocated on demand
+        var newList: MutableList<ConeAttribute<*>>? = null
+        var hasDifference = false
+
+        for ((i, attr) in this.withIndex()) {
+            if (attr !is ConeAttributeWithConeType) continue
+            val substitutedAttribute = attr.transformOrNull(transform) ?: continue
+            if (newList == null) {
+                newList = this.toMutableList()
+            }
+            newList[i] = substitutedAttribute
+            hasDifference = hasDifference || substitutedAttribute != attr
+        }
+
+        if (newList != null && !hasDifference) {
+            return this
+        }
+
+        return newList?.let(Companion::create)
+    }
+
     override val typeRegistry: TypeRegistry<ConeAttribute<*>, ConeAttribute<*>>
         get() = Companion
 }

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/substitution/Substitutors.kt

@@ -59,7 +59,7 @@ abstract class AbstractConeSubstitutor(protected val typeContext: ConeTypeContex
         }
 
         val substitutedType = newType ?: type.substituteRecursive()
-        val substitutedAttributes = (substitutedType ?: type).attributes.substituteAbbreviationOrNull()
+        val substitutedAttributes = (substitutedType ?: type).attributes.transformTypesWith(this::substituteOrNull)
 
         return if (substitutedType != null || substitutedAttributes != null) {
             var result = substitutedType ?: type
@@ -74,14 +74,6 @@ abstract class AbstractConeSubstitutor(protected val typeContext: ConeTypeContex
         }
     }
 
-    private fun ConeAttributes.substituteAbbreviationOrNull(): ConeAttributes? {
-        val abbreviatedTypeAttribute = abbreviatedType ?: return null
-        substituteOrNull(abbreviatedTypeAttribute.coneType)?.let {
-            return replace(abbreviatedTypeAttribute, AbbreviatedTypeAttribute(it))
-        }
-        return null
-    }
-
     private fun ConeKotlinType.substituteRecursive(): ConeKotlinType? {
         return when (this) {
             is ConeClassLikeType -> this.substituteArguments()