FIR: make ArrayMapOwner.arrayMap really protected

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

@@ -6,9 +6,7 @@
 package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.fir.utils.AttributeArrayOwner
-import org.jetbrains.kotlin.fir.utils.Protected
 import org.jetbrains.kotlin.fir.utils.TypeRegistry
-import org.jetbrains.kotlin.fir.utils.isEmpty
 import org.jetbrains.kotlin.types.model.AnnotationMarker
 import org.jetbrains.kotlin.utils.addIfNotNull
 import kotlin.properties.ReadOnlyProperty
@@ -24,7 +22,6 @@ abstract class ConeAttribute<T : ConeAttribute<T>> : AnnotationMarker {
     abstract val key: KClass<out T>
 }
 
-@OptIn(Protected::class)
 class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : AttributeArrayOwner<ConeAttribute<*>, ConeAttribute<*>>(),
     Iterable<ConeAttribute<*>> {
 
@@ -85,16 +82,12 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
     }
 
     fun remove(attribute: ConeAttribute<*>): ConeAttributes {
-        if (arrayMap.isEmpty()) return this
+        if (isEmpty()) return this
         val attributes = arrayMap.filter { it != attribute }
         if (attributes.size == arrayMap.size) return this
         return create(attributes)
     }
 
-    override fun iterator(): Iterator<ConeAttribute<*>> {
-        return arrayMap.iterator()
-    }
-
     private inline fun perform(other: ConeAttributes, op: ConeAttribute<*>.(ConeAttribute<*>?) -> ConeAttribute<*>?): ConeAttributes {
         if (this.isEmpty() && other.isEmpty()) return this
         val attributes = mutableListOf<ConeAttribute<*>>()
@@ -109,8 +102,4 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
 
     override val typeRegistry: TypeRegistry<ConeAttribute<*>, ConeAttribute<*>>
         get() = Companion
-
-    private fun isEmpty(): Boolean {
-        return arrayMap.isEmpty()
-    }
 }

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

@@ -14,9 +14,6 @@ sealed class ArrayMap<T : Any> : Iterable<T> {
     abstract fun copy(): ArrayMap<T>
 }
 
-fun ArrayMap<*>.isEmpty(): Boolean = size == 0
-fun ArrayMap<*>.isNotEmpty(): Boolean = size != 0
-
 internal object EmptyArrayMap : ArrayMap<Nothing>() {
     override val size: Int
         get() = 0

compiler/fir/cones/src/org/jetbrains/kotlin/fir/utils/ArrayMapOwner.kt

@@ -14,10 +14,8 @@ import kotlin.reflect.KProperty
 @RequiresOptIn
 annotation class Protected
 
-@OptIn(Protected::class)
-abstract class AbstractArrayMapOwner<K : Any, V : Any> {
-    @Protected
-    abstract val arrayMap: ArrayMap<V>
+abstract class AbstractArrayMapOwner<K : Any, V : Any> : Iterable<V> {
+    protected abstract val arrayMap: ArrayMap<V>
     protected abstract val typeRegistry: TypeRegistry<K, V>
 
     abstract class AbstractArrayMapAccessor<K : Any, V : Any, T : V>(
@@ -31,6 +29,14 @@ abstract class AbstractArrayMapOwner<K : Any, V : Any> {
     }
 
     protected abstract fun registerComponent(tClass: KClass<out K>, value: V)
+
+    final override fun iterator(): Iterator<V> = arrayMap.iterator()
+
+    fun isEmpty(): Boolean = arrayMap.size == 0
+
+    fun isNotEmpty(): Boolean = arrayMap.size != 0
+
+    operator fun get(index: Int): V? = arrayMap[index]
 }
 
 class ArrayMapAccessor<K : Any, V : Any, T : V>(

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

@@ -16,7 +16,6 @@ import kotlin.reflect.KClass
  * Note that you can remove attributes from [AttributeArrayOwner] despite
  *   from components in [ComponentArrayOwner]
  */
-@OptIn(Protected::class)
 abstract class AttributeArrayOwner<K : Any, T : Any> protected constructor(
     arrayMap: ArrayMap<T>
 ) : AbstractArrayMapOwner<K, T>() {

compiler/fir/cones/src/org/jetbrains/kotlin/fir/utils/ComponentArrayOwner.kt

@@ -11,7 +11,6 @@ import kotlin.reflect.KClass
  * [ComponentArrayOwner] based on [ArrayMap] with flexible size and should be used for
  *   storing services in entities with limited number of instances, like FirSession
  */
-@OptIn(Protected::class)
 abstract class ComponentArrayOwner<K : Any, V : Any> : AbstractArrayMapOwner<K, V>() {
     final override val arrayMap: ArrayMap<V> =
         ArrayMapImpl()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionService.kt

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.NoMutableState
 import org.jetbrains.kotlin.fir.utils.ArrayMapAccessor
 import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner
-import org.jetbrains.kotlin.fir.utils.Protected
 import org.jetbrains.kotlin.fir.utils.TypeRegistry
 import kotlin.reflect.KClass
 
@@ -50,7 +49,6 @@ class FirExtensionService(val session: FirSession) : ComponentArrayOwner<FirExte
     }
 
     @PluginServicesInitialization
-    @OptIn(Protected::class)
     fun getAllExtensions(): List<FirExtension> {
         return arrayMap.flatten()
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -377,7 +377,7 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 
     override fun KotlinTypeMarker.getAnnotations(): List<AnnotationMarker> {
         require(this is ConeKotlinType)
-        return attributes.arrayMap.toList()
+        return attributes.toList()
     }
 
     override fun SimpleTypeMarker.isStubType(): Boolean {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.utils.isNotEmpty
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.types.Variance
@@ -413,7 +412,7 @@ class FirRenderer(builder: StringBuilder, private val mode: RenderMode = RenderM
     }
 
     private fun FirDeclaration.renderDeclarationAttributesIfNeeded() {
-        if (mode.renderDeclarationAttributes && attributes.arrayMap.isNotEmpty()) {
+        if (mode.renderDeclarationAttributes && attributes.isNotEmpty()) {
             val attributes = getAttributesWithValues().mapNotNull { (klass, value) ->
                 value?.let { klass.simpleName to value.renderAsDeclarationAttributeValue() }
             }.joinToString { (name, value) -> "$name=$value" }
@@ -427,10 +426,9 @@ class FirRenderer(builder: StringBuilder, private val mode: RenderMode = RenderM
         }
     }
 
-
     private fun FirDeclaration.getAttributesWithValues(): List<Pair<KClass<out FirDeclarationDataKey>, Any?>> {
         val attributesMap = FirDeclarationDataRegistry.allValuesThreadUnsafeForRendering()
-        return attributesMap.entries.sortedBy { it.key.simpleName }.map { (klass, index) -> klass to attributes.arrayMap[index] }
+        return attributesMap.entries.sortedBy { it.key.simpleName }.map { (klass, index) -> klass to attributes[index] }
     }
 
     private fun Any.renderAsDeclarationAttributeValue() = when (this) {