[FIR] Add ability to rewrite attributes in FirDeclarationAttributes

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt

@@ -18,8 +18,12 @@ class FirDeclarationAttributes : AttributeArrayOwner<FirDeclarationDataKey, Any>
     override val typeRegistry: TypeRegistry<FirDeclarationDataKey, Any>
         get() = FirDeclarationDataRegistry
 
-    internal operator fun set(key: KClass<out FirDeclarationDataKey>, value: Any) {
-        registerComponent(key, value)
+    internal operator fun set(key: KClass<out FirDeclarationDataKey>, value: Any?) {
+        if (value == null) {
+            removeComponent(key)
+        } else {
+            registerComponent(key, value)
+        }
     }
 }
 
@@ -44,9 +48,7 @@ object FirDeclarationDataRegistry : TypeRegistry<FirDeclarationDataKey, Any>() {
         }
 
         override fun setValue(thisRef: FirDeclaration, property: KProperty<*>, value: V?) {
-            if (value != null) {
-                thisRef.attributes[key] = value
-            }
+            thisRef.attributes[key] = value
         }
     }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ArrayMap.kt

@@ -38,7 +38,7 @@ internal class OneElementArrayMap<T : Any>(val value: T, val index: Int) : Array
     }
 }
 
-internal  class ArrayMapImpl<T : Any> : ArrayMap<T>() {
+internal class ArrayMapImpl<T : Any> : ArrayMap<T>() {
     companion object {
         private const val DEFAULT_SIZE = 20
         private const val INCREASE_K = 2
@@ -56,13 +56,28 @@ internal  class ArrayMapImpl<T : Any> : ArrayMap<T>() {
 
     override operator fun set(index: Int, value: T) {
         ensureCapacity(index)
+        if (data[index] == null) {
+            size++
+        }
         data[index] = value
-
-        size++
     }
 
     override operator fun get(index: Int): T? {
         @Suppress("UNCHECKED_CAST")
         return data.getOrNull(index) as T?
     }
+
+    fun remove(index: Int) {
+        if (data[index] != null) {
+            size--
+        }
+        data[index] = null
+    }
+
+    fun entries(): List<Entry<T>> {
+        @Suppress("UNCHECKED_CAST")
+        return data.mapIndexedNotNull { index, value -> if (value != null) Entry(index, value as T) else null }
+    }
+
+    data class Entry<T>(override val key: Int, override val value: T) : Map.Entry<Int, T>
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/AttributeArrayOwner.kt

@@ -12,6 +12,9 @@ import kotlin.reflect.KClass
  *   depending on array map fullness
  * [AttributeArrayOwner] can be used in classes with many instances,
  *   like user data for Fir elements or attributes for cone types
+ *
+ * Note that you can remove attributes from [AttributeArrayOwner] despite
+ *   from components in [ComponentArrayOwner]
  */
 abstract class AttributeArrayOwner<K : Any, T : Any> : AbstractArrayMapOwner<K, T>() {
     @Suppress("UNCHECKED_CAST")
@@ -36,4 +39,21 @@ abstract class AttributeArrayOwner<K : Any, T : Any> : AbstractArrayMapOwner<K,
 
         arrayMap[id] = value
     }
+
+    protected fun removeComponent(tClass: KClass<out K>) {
+        val id = typeRegistry.getId(tClass)
+        if (arrayMap[id] == null) return
+        @Suppress("UNCHECKED_CAST")
+        when (arrayMap.size) {
+            1 -> arrayMap = EmptyArrayMap as ArrayMap<T>
+            else -> {
+                val map = arrayMap as ArrayMapImpl<T>
+                map.remove(id)
+                if (map.size == 1) {
+                    val (index, value) = map.entries().first()
+                    arrayMap = OneElementArrayMap(value, index)
+                }
+            }
+        }
+    }
 }