From cf494ada0b495e280365550e904d35bd168cb768 Mon Sep 17 00:00:00 2001
From: Kirill Rakhman <kirill.rakhman@jetbrains.com>
Date: Tue, 16 Jan 2024 14:03:06 +0100
Subject: [PATCH] [FIR] Flatten ConeAttributeWithConeType when transforming

This is necessary to prevent exponential growth of the attribute
structure.
Since we usually care for the most inner value of a
ConeAttributeWithConeType like EnhancedTypeForWarningAttribute, this
shouldn't alter any behavior.
---
 .../jetbrains/kotlin/fir/types/ConeAttributes.kt    | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt
index d358fc9f423..51d259bbba4 100644
--- a/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt
+++ b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt
@@ -63,7 +63,11 @@ inline fun <T : ConeAttributeWithConeType<T>> ConeAttributeWithConeType<T>.trans
 ): ConeAttributeWithConeType<T>? {
     val transformedType = transform(coneType) ?: return null
     if (transformedType == coneType) return this
-    return copyWith(transformedType)
+
+    // If the type contains the attribute itself, use the nested type from the attribute to prevent exponential growth.
+    // As an example, consider a substitution {T -> Attr(Foo) Bar} applied to a type `Attr(T) T`.
+    // If we don't flatten the attribute chain, we would get `Attr(Attr(Foo) Bar) Bar`.
+    return copyWith(transformedType.attributes[key]?.coneType ?: transformedType)
 }
 
 typealias ConeAttributeKey = KClass<out ConeAttribute<*>>
@@ -120,8 +124,13 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
     }
 
     operator fun contains(attributeKey: KClass<out ConeAttribute<*>>): Boolean {
+        return get(attributeKey) != null
+    }
+
+    operator fun <T : ConeAttribute<*>> get(attributeKey: KClass<T>) : T? {
         val index = getId(attributeKey)
-        return arrayMap[index] != null
+        @Suppress("UNCHECKED_CAST")
+        return arrayMap[index] as T?
     }
 
     operator fun plus(attribute: ConeAttribute<*>): ConeAttributes {