FIR IDE: add ability to specify type argument in KtClassTypeBuilder

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtTypeCreator.kt

@@ -5,12 +5,14 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.components
 
+import org.jetbrains.kotlin.idea.frontend.api.KtTypeArgument
+import org.jetbrains.kotlin.idea.frontend.api.KtTypeArgumentWithVariance
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtClassType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtTypeNullability
 import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.types.Variance
 
 abstract class KtTypeCreator : KtAnalysisSessionComponent() {
     abstract fun buildClassType(builder: KtClassTypeBuilder): KtClassType
@@ -35,16 +37,20 @@ inline fun KtTypeCreatorMixIn.buildClassType(
 sealed class KtTypeBuilder
 
 sealed class KtClassTypeBuilder : KtTypeBuilder() {
-    private val _arguments = mutableListOf<KtType>()
+    private val _arguments = mutableListOf<KtTypeArgument>()
 
     var nullability: KtTypeNullability = KtTypeNullability.NON_NULLABLE
 
-    val arguments: List<KtType> get() = _arguments
+    val arguments: List<KtTypeArgument> get() = _arguments
 
-    fun argument(argument: KtType) {
+    fun argument(argument: KtTypeArgument) {
         _arguments += argument
     }
 
+    fun argument(type: KtType, variance: Variance = Variance.INVARIANT) {
+        _arguments += KtTypeArgumentWithVariance(type, variance, type.token)
+    }
+
     class ByClassId(val classId: ClassId) : KtClassTypeBuilder()
     class BySymbol(val symbol: KtClassOrObjectSymbol) : KtClassTypeBuilder()
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirAnalysisSessionComponent.kt

@@ -13,7 +13,10 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.toFirDiagnostics
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.typeContext
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.idea.frontend.api.KtStarProjectionTypeArgument
+import org.jetbrains.kotlin.idea.frontend.api.KtTypeArgument
+import org.jetbrains.kotlin.idea.frontend.api.KtTypeArgumentWithVariance
 import org.jetbrains.kotlin.fir.types.ConeTypeCheckerContext
 import org.jetbrains.kotlin.idea.asJava.asPsiType
 import org.jetbrains.kotlin.idea.frontend.api.diagnostics.KtDiagnosticWithPsi
@@ -22,11 +25,13 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KT_DIAGNOSTIC_CONV
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
+import org.jetbrains.kotlin.types.model.convertVariance
 
 internal interface KtFirAnalysisSessionComponent {
     val analysisSession: KtFirAnalysisSession
 
     val rootModuleSession: FirSession get() = analysisSession.firResolveState.rootModuleSession
+    val typeContext: ConeInferenceContext get() = rootModuleSession.typeContext
     val firSymbolBuilder get() = analysisSession.firSymbolBuilder
     val firResolveState get() = analysisSession.firResolveState
 
@@ -50,6 +55,14 @@ internal interface KtFirAnalysisSessionComponent {
             return coneType
         }
 
+    val KtTypeArgument.coneTypeProjection: ConeTypeProjection
+        get() = when (this) {
+            is KtStarProjectionTypeArgument -> ConeStarProjection
+            is KtTypeArgumentWithVariance -> {
+                typeContext.createTypeArgument(type.coneType, variance.convertVariance()) as ConeTypeProjection
+            }
+        }
+
     fun createTypeCheckerContext() = ConeTypeCheckerContext(
         isErrorTypeEqualsToAnything = true,
         isStubTypeEqualsToAnything = true,

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirTypeCreator.kt

@@ -40,9 +40,10 @@ internal class KtFirTypeCreator(
             }
         }
 
-        val coneType = rootModuleSession.typeContext.createSimpleType(
+        val typeContext = rootModuleSession.typeContext
+        val coneType = typeContext.createSimpleType(
             lookupTag,
-            builder.arguments.map { it.coneType },
+            builder.arguments.map { it.coneTypeProjection },
             builder.nullability.isNullable
         ) as ConeClassLikeType
 

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/TypeCreatorTest.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescrip
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.types.Variance
 import org.junit.Test
 
 internal class TypeCreatorTest : KotlinLightCodeInsightFixtureTestCase() {
@@ -73,6 +74,15 @@ internal class TypeCreatorTest : KotlinLightCodeInsightFixtureTestCase() {
         }
     }
 
+    @Test
+    fun testTypeArgumentVariance() {
+        doTest("KClass<out Number>") {
+            buildClassType(StandardClassIds.KClass) {
+                argument(buildClassType(StandardClassIds.Number), variance = Variance.OUT_VARIANCE)
+            }
+        }
+    }
+
 
     companion object {
         private val RENDERING_OPTIONS = KtTypeRendererOptions.SHORT_NAMES