FIR: Refactor extracting type arguments from FirTypeRef

Extract them all instead of re-iterating for each argument

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt

@@ -629,38 +629,33 @@ internal fun checkCondition(condition: FirExpression, context: CheckerContext, r
     }
 }
 
-fun extractArgumentTypeRefAndSource(typeRef: FirTypeRef?, index: Int): FirTypeRefSource? {
-    if (typeRef is FirResolvedTypeRef) {
-        val delegatedTypeRef = typeRef.delegatedTypeRef
-        if (delegatedTypeRef is FirUserTypeRef) {
-            var currentIndex = index
+fun extractArgumentsTypeRefAndSource(typeRef: FirTypeRef?): List<FirTypeRefSource>? {
+    if (typeRef !is FirResolvedTypeRef) return null
+    val result = mutableListOf<FirTypeRefSource>()
+    when (val delegatedTypeRef = typeRef.delegatedTypeRef) {
+        is FirUserTypeRef -> {
             val qualifier = delegatedTypeRef.qualifier
 
             for (i in qualifier.size - 1 downTo 0) {
-                val typeArguments = qualifier[i].typeArgumentList.typeArguments
-                if (currentIndex < typeArguments.size) {
-                    val typeArgument = typeArguments.elementAtOrNull(currentIndex)
-                    return if (typeArgument is FirTypeProjection)
-                        FirTypeRefSource((typeArgument as? FirTypeProjectionWithVariance)?.typeRef, typeArgument.source)
-                    else null
-                } else {
-                    currentIndex -= typeArguments.size
+                for (typeArgument in qualifier[i].typeArgumentList.typeArguments) {
+                    result.add(FirTypeRefSource((typeArgument as? FirTypeProjectionWithVariance)?.typeRef, typeArgument.source))
                 }
             }
-        } else if (delegatedTypeRef is FirFunctionTypeRef) {
-            val valueParameters = delegatedTypeRef.valueParameters
-            if (index < valueParameters.size) {
-                val valueParamTypeRef = valueParameters.elementAt(index).returnTypeRef
-                return FirTypeRefSource(valueParamTypeRef, valueParamTypeRef.source)
-            }
-            if (index == valueParameters.size) {
-                val returnTypeRef = delegatedTypeRef.returnTypeRef
-                return FirTypeRefSource(returnTypeRef, returnTypeRef.source)
-            }
         }
+        is FirFunctionTypeRef -> {
+            val valueParameters = delegatedTypeRef.valueParameters
+
+            for (valueParameter in valueParameters) {
+                val valueParamTypeRef = valueParameter.returnTypeRef
+                result.add(FirTypeRefSource(valueParamTypeRef, valueParamTypeRef.source))
+            }
+            val returnTypeRef = delegatedTypeRef.returnTypeRef
+            result.add(FirTypeRefSource(returnTypeRef, returnTypeRef.source))
+        }
+        else -> return null
     }
 
-    return null
+    return result
 }
 
 data class FirTypeRefSource(val typeRef: FirTypeRef?, val source: KtSourceElement?)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirUpperBoundViolatedHelpers.kt

@@ -61,10 +61,11 @@ fun checkUpperBoundViolated(
     }
 
     val substitutor = substitutorByMap(substitution, context.session)
+    val typeRefAndSourcesForArguments = extractArgumentsTypeRefAndSource(typeRef) ?: return
     val typeArgumentsWithSourceInfo = type.typeArguments.withIndex().map { (index, projection) ->
         val (argTypeRef, source) =
             if (!isAbbreviatedType)
-                extractArgumentTypeRefAndSource(typeRef, index) ?: return
+                typeRefAndSourcesForArguments.getOrNull(index) ?: return
             else
                 // For abbreviated arguments we use the whole typeRef as a place to report
                 FirTypeRefSource(null, typeRef.source)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirClassVarianceChecker.kt

@@ -8,11 +8,11 @@ package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.descriptors.Visibilities
-import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.extractArgumentTypeRefAndSource
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.extractArgumentsTypeRefAndSource
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -136,6 +136,7 @@ object FirClassVarianceChecker : FirClassChecker() {
             val fullyExpandedType = type.fullyExpandedType(context.session)
             val classSymbol = fullyExpandedType.lookupTag.toSymbol(context.session)
             if (classSymbol is FirClassSymbol<*>) {
+                val typeRefAndSourcesForArguments = extractArgumentsTypeRefAndSource(typeRef)
                 for ((index, typeArgument) in fullyExpandedType.typeArguments.withIndex()) {
                     val paramVariance = classSymbol.typeParameterSymbols.getOrNull(index)?.variance ?: continue
 
@@ -156,7 +157,7 @@ object FirClassVarianceChecker : FirClassChecker() {
                     }
 
                     if (newVariance != null) {
-                        val subTypeRefAndSource = extractArgumentTypeRefAndSource(typeRef, index)
+                        val subTypeRefAndSource = typeRefAndSourcesForArguments?.getOrNull(index)
 
                         checkVarianceConflict(
                             typeArgumentType, newVariance, subTypeRefAndSource?.typeRef, containingType,

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirOuterClassArgumentsRequiredChecker.kt

@@ -5,11 +5,11 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
-import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.extractArgumentTypeRefAndSource
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.extractArgumentsTypeRefAndSource
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.resolve.isValidTypeParameterFromOuterDeclaration
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -56,8 +56,8 @@ object FirOuterClassArgumentsRequiredChecker : FirRegularClassChecker() {
             }
         }
 
-        for (index in type.typeArguments.indices) {
-            val firTypeRefSource = extractArgumentTypeRefAndSource(typeRef, index) ?: continue
+        val typeRefAndSourcesForArguments = extractArgumentsTypeRefAndSource(typeRef) ?: return
+        for (firTypeRefSource in typeRefAndSourcesForArguments) {
             firTypeRefSource.typeRef?.let { checkOuterClassArgumentsRequired(it, declaration, context, reporter) }
         }
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirProjectionRelationChecker.kt

@@ -6,11 +6,11 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
 import org.jetbrains.kotlin.KtFakeSourceElementKind
-import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.extractArgumentTypeRefAndSource
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.extractArgumentsTypeRefAndSource
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
@@ -52,6 +52,7 @@ object FirProjectionRelationChecker : FirBasicDeclarationChecker() {
 
         val size = minOf(typeParameters.size, typeArguments.size)
 
+        val typeRefAndSourcesForArguments = extractArgumentsTypeRefAndSource(typeRef) ?: return
         for (it in 0 until size) {
             val proto = typeParameters[it]
             val actual = typeArguments[it]
@@ -72,7 +73,7 @@ object FirProjectionRelationChecker : FirBasicDeclarationChecker() {
                 ProjectionRelation.None
             }
 
-            val argTypeRefSource = extractArgumentTypeRefAndSource(typeRef, it) ?: continue
+            val argTypeRefSource = typeRefAndSourcesForArguments.getOrNull(it) ?: continue
 
             if (projectionRelation != ProjectionRelation.None && typeRef.source?.kind !is KtFakeSourceElementKind) {
                 reporter.reportOn(

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirSupertypesChecker.kt

@@ -9,13 +9,15 @@ import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.checkers.*
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
-import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.diagnostics.withSuppressedDiagnostics
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirField
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.utils.modality
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.languageVersionSettings
@@ -133,9 +135,10 @@ object FirSupertypesChecker : FirClassChecker() {
         reporter: DiagnosticReporter,
         context: CheckerContext
     ) {
+        val typeRefAndSourcesForArguments = extractArgumentsTypeRefAndSource(superTypeRef) ?: return
         for ((index, typeArgument) in coneType.typeArguments.withIndex()) {
             if (typeArgument.isConflictingOrNotInvariant) {
-                val (_, argSource) = extractArgumentTypeRefAndSource(superTypeRef, index) ?: continue
+                val (_, argSource) = typeRefAndSourcesForArguments.getOrNull(index) ?: continue
                 reporter.reportOn(
                     argSource ?: superTypeRef.source,
                     FirErrors.PROJECTION_IN_IMMEDIATE_ARGUMENT_TO_SUPERTYPE,