FIR LC: use optimal type mapping mode for return type

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10PsiTypeProvider.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.withValidityAssertion
 import org.jetbrains.kotlin.codegen.signature.BothSignatureWriter
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
+import org.jetbrains.kotlin.load.kotlin.getOptimalModeForReturnType
 import org.jetbrains.kotlin.load.kotlin.getOptimalModeForValueParameter
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.types.*
@@ -41,6 +42,11 @@ internal class KtFe10PsiTypeProvider(override val analysisSession: KtFe10Analysi
         return asPsiType(simplifyType(kotlinType), useSitePosition, mode)
     }
 
+    override fun getOptimalModeForReturnType(type: KtType, isAnnotationMethod: Boolean): TypeMappingMode = withValidityAssertion {
+        require(type is KtFe10Type)
+        typeMapper.typeContext.getOptimalModeForReturnType(type.type, isAnnotationMethod)
+    }
+
     override fun getOptimalModeForValueParameter(type: KtType): TypeMappingMode = withValidityAssertion {
         require(type is KtFe10Type)
         typeMapper.typeContext.getOptimalModeForValueParameter(type.type)

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirPsiTypeProvider.kt

@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
+import org.jetbrains.kotlin.load.kotlin.getOptimalModeForReturnType
 import org.jetbrains.kotlin.load.kotlin.getOptimalModeForValueParameter
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.psi
@@ -59,6 +60,11 @@ internal class KtFirPsiTypeProvider(
         type.coneType.asPsiType(rootModuleSession, analysisSession.firResolveState, mode, useSitePosition)
     }
 
+    override fun getOptimalModeForReturnType(type: KtType, isAnnotationMethod: Boolean): TypeMappingMode = withValidityAssertion {
+        require(type is KtFirType)
+        rootModuleSession.jvmTypeMapper.typeContext.getOptimalModeForReturnType(type.coneType, isAnnotationMethod)
+    }
+
     override fun getOptimalModeForValueParameter(type: KtType): TypeMappingMode = withValidityAssertion {
         require(type is KtFirType)
         rootModuleSession.jvmTypeMapper.typeContext.getOptimalModeForValueParameter(type.coneType)

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtPsiTypeProvider.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 
 public abstract class KtPsiTypeProvider : KtAnalysisSessionComponent() {
     public abstract fun asPsiType(type: KtType, useSitePosition: PsiElement, mode: TypeMappingMode): PsiType?
+    public abstract fun getOptimalModeForReturnType(type: KtType, isAnnotationMethod: Boolean): TypeMappingMode
     public abstract fun getOptimalModeForValueParameter(type: KtType): TypeMappingMode
 }
 
@@ -38,6 +39,9 @@ public interface KtPsiTypeProviderMixIn : KtAnalysisSessionMixIn {
     ): PsiType? =
         analysisSession.psiTypeProvider.asPsiType(this, useSitePosition, mode)
 
+    public fun KtType.getOptimalModeForReturnType(isAnnotationMethod: Boolean): TypeMappingMode =
+        analysisSession.psiTypeProvider.getOptimalModeForReturnType(this, isAnnotationMethod)
+
     public fun KtType.getOptimalModeForValueParameter(): TypeMappingMode =
         analysisSession.psiTypeProvider.getOptimalModeForValueParameter(this)
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/fields/FirLightFieldForPropertySymbol.kt

@@ -34,10 +34,19 @@ internal class FirLightFieldForPropertySymbol(
             when {
                 isDelegated ->
                     (kotlinOrigin as? KtProperty)?.delegateExpression?.let {
-                        it.getKtType()?.asPsiType(this@FirLightFieldForPropertySymbol)
+                        val ktType = it.getKtType()
+                        ktType?.asPsiType(
+                            this@FirLightFieldForPropertySymbol,
+                            ktType.getOptimalModeForReturnType(containingClass.isAnnotationType)
+                        )
                     }
-                else ->
-                    propertySymbol.annotatedType.type.asPsiType(this@FirLightFieldForPropertySymbol)
+                else -> {
+                    val ktType = propertySymbol.annotatedType.type
+                    ktType.asPsiType(
+                        this@FirLightFieldForPropertySymbol,
+                        ktType.getOptimalModeForReturnType(containingClass.isAnnotationType)
+                    )
+                }
             }
         } ?: nonExistentType()
     }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/methods/FirLightAccessorMethodForSymbol.kt

@@ -150,9 +150,12 @@ internal class FirLightAccessorMethodForSymbol(
     private val _returnedType: PsiType by lazyPub {
         if (!isGetter) return@lazyPub PsiType.VOID
         analyzeWithSymbolAsContext(containingPropertySymbol) {
-            containingPropertySymbol.annotatedType.type.asPsiType(this@FirLightAccessorMethodForSymbol)
-                ?: this@FirLightAccessorMethodForSymbol.nonExistentType()
-        }
+            val ktType = containingPropertySymbol.annotatedType.type
+            ktType.asPsiType(
+                this@FirLightAccessorMethodForSymbol,
+                ktType.getOptimalModeForReturnType(containingClass.isAnnotationType)
+            )
+        } ?: nonExistentType()
     }
 
     override fun getReturnType(): PsiType = _returnedType

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/methods/FirLightSimpleMethodForSymbol.kt

@@ -145,7 +145,10 @@ internal class FirLightSimpleMethodForSymbol(
                     else ->
                         functionSymbol.annotatedType.type
                 }
-            ktType.asPsiType(this@FirLightSimpleMethodForSymbol)
+            ktType.asPsiType(
+                this@FirLightSimpleMethodForSymbol,
+                ktType.getOptimalModeForReturnType(containingClass.isAnnotationType)
+            )
         } ?: nonExistentType()
     }
 

compiler/testData/asJava/lightClasses/nullabilityAnnotations/PlatformTypes.kt

@@ -6,4 +6,6 @@ class PlatformTypes {
     fun simplyPlatform() = Collections.singletonList("")[0]
     fun bothNullable() = Collections.emptyList<String>() ?: null
     fun bothNotNull() = Collections.emptyList<String>()!!
-}
+}
+
+// FIR_COMPARISON

compiler/testData/asJava/ultraLightClasses/wildcardOptimization.kt

@@ -51,3 +51,5 @@ class Container {
 
     fun <Q : Final> typeParameter(x: Out<Q>, y: In<Q>) {}
 }
+
+// FIR_COMPARISON