FIR type enhancement: extract indexed Java type qualifiers

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/EnhancementSignatureParts.kt

@@ -46,8 +46,8 @@ internal class EnhancementSignatureParts(
     ): PartEnhancementResult {
         val qualifiers = computeIndexedQualifiersForOverride(jsr305State)
 
-        val qualifiersWithPredefined: ((Int) -> JavaTypeQualifiers)? = predefined?.let {
+        val qualifiersWithPredefined = predefined?.let {
-            { index ->
+            IndexedJavaTypeQualifiers(qualifiers.size) { index ->
                 predefined.map[index] ?: qualifiers(index)
             }
         }
@@ -260,7 +260,7 @@ internal class EnhancementSignatureParts(
         )
     }
 
-    private fun computeIndexedQualifiersForOverride(jsr305State: Jsr305State): (Int) -> JavaTypeQualifiers {
+    private fun computeIndexedQualifiersForOverride(jsr305State: Jsr305State): IndexedJavaTypeQualifiers {
         val indexedFromSupertypes = fromOverridden.map { it.toIndexed(typeQualifierResolver, jsr305State, context) }
         val indexedThisType = current.toIndexed(typeQualifierResolver, jsr305State, context)
 
@@ -283,7 +283,7 @@ internal class EnhancementSignatureParts(
             qualifiers.computeQualifiersForOverride(verticalSlice, defaultQualifiers, isHeadTypeConstructor, jsr305State)
         }
 
-        return { index -> computedResult.getOrElse(index) { JavaTypeQualifiers.NONE } }
+        return IndexedJavaTypeQualifiers(computedResult)
     }
 
     @Suppress("unused")

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/javaTypeUtils.kt

@@ -38,8 +38,16 @@ import org.jetbrains.kotlin.types.AbstractStrictEqualityTypeChecker
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
+internal class IndexedJavaTypeQualifiers(private val data: Array<JavaTypeQualifiers>) {
+    constructor(size: Int, compute: (Int) -> JavaTypeQualifiers) : this(Array(size) { compute(it) })
+
+    operator fun invoke(index: Int) = data.getOrElse(index) { JavaTypeQualifiers.NONE }
+
+    val size: Int get() = data.size
+}
+
 internal fun FirJavaTypeRef.enhance(
-    qualifiers: (Int) -> JavaTypeQualifiers
+    qualifiers: IndexedJavaTypeQualifiers
 ): FirResolvedTypeRef {
     return type.enhancePossiblyFlexible(session, annotations, qualifiers, 0)
 }
@@ -51,7 +59,7 @@ internal fun FirJavaTypeRef.enhance(
 private fun JavaType?.enhancePossiblyFlexible(
     session: FirSession,
     annotations: List<FirAnnotationCall>,
-    qualifiers: (Int) -> JavaTypeQualifiers,
+    qualifiers: IndexedJavaTypeQualifiers,
     index: Int
 ): FirResolvedTypeRef {
     val type = this
@@ -117,7 +125,7 @@ private fun JavaClassifierType.enhanceInflexibleType(
     annotations: List<FirAnnotationCall>,
     arguments: List<JavaType?>,
     position: TypeComponentPosition,
-    qualifiers: (Int) -> JavaTypeQualifiers,
+    qualifiers: IndexedJavaTypeQualifiers,
     index: Int
 ): ConeLookupTagBasedType {
     val originalSymbol = when (val classifier = classifier) {