Refactor KotlinType::asPsiType for ultra-light classes

- Extract HOF that UltraLightSupport::mapType that allows to call
arbitrary functions of type mapper
- Use it for computing return type of functions
- "declaration: KtDeclaration" parameter became unused since it was only
used for return types

compiler/backend/src/org/jetbrains/kotlin/codegen/state/KotlinTypeMapper.java

@@ -388,7 +388,7 @@ public class KotlinTypeMapper {
     }
 
     @NotNull
-    private Type mapReturnType(@NotNull CallableDescriptor descriptor, @Nullable JvmSignatureWriter sw) {
+    public Type mapReturnType(@NotNull CallableDescriptor descriptor, @Nullable JvmSignatureWriter sw) {
         KotlinType returnType = descriptor.getReturnType();
         assert returnType != null : "Function has no return type: " + descriptor;
 

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightPsi.kt

@@ -81,7 +81,7 @@ class KtUltraLightClass(classOrObject: KtClassOrObject, private val support: Ult
         getDescriptor()?.typeConstructor?.supertypes.orEmpty().asSequence()
 
     private fun mapSupertype(supertype: KotlinType) =
-        supertype.asPsiType(classOrObject, support, TypeMappingMode.SUPER_TYPE, this) as? PsiClassType
+        supertype.asPsiType(support, TypeMappingMode.SUPER_TYPE, this) as? PsiClassType
 
     override fun createExtendsList(): PsiReferenceList? =
         if (tooComplex) super.createExtendsList()
@@ -336,16 +336,19 @@ class KtUltraLightClass(classOrObject: KtClassOrObject, private val support: Ult
         }
     }
 
-    private fun methodReturnType(f: KtDeclaration, wrapper: KtUltraLightMethod): PsiType {
-        val desc = f.resolve()?.let { if (it is PropertyDescriptor) it.getter else it }
-        val kotlinType = (desc as? FunctionDescriptor)?.returnType ?: return PsiType.NULL
-        val mode = when {
-            typeMapper(support).forceBoxedReturnType(desc) -> TypeMappingMode.RETURN_TYPE_BOXED
-            else -> TypeMappingMode.getOptimalModeForReturnType(kotlinType, false)
+    private fun methodReturnType(ktDeclaration: KtDeclaration, wrapper: KtUltraLightMethod): PsiType {
+        val desc =
+            ktDeclaration.resolve()?.getterIfProperty() as? FunctionDescriptor
+                ?: return PsiType.NULL
+
+        return support.mapType(wrapper) { typeMapper, signatureWriter ->
+            typeMapper.mapReturnType(desc, signatureWriter)
         }
-        return kotlinType.asPsiType(f, support, mode, wrapper)
     }
 
+    private fun DeclarationDescriptor.getterIfProperty() =
+        if (this@getterIfProperty is PropertyDescriptor) this@getterIfProperty.getter else this@getterIfProperty
+
     private fun lightMethod(name: String, declaration: KtDeclaration, forceStatic: Boolean): LightMethodBuilder {
         val accessedProperty = if (declaration is KtPropertyAccessor) declaration.property else null
         val outer = accessedProperty ?: declaration
@@ -535,7 +538,7 @@ private class KtUltraLightField(
                         val context = LightClassGenerationSupport.getInstance(project).analyze(declaration)
                         PropertyCodegen.getDelegateTypeForProperty(declaration, it, context)
                     }
-                    ?.let { it.asPsiType(declaration, support, TypeMappingMode.getOptimalModeForValueParameter(it), this) }
+                    ?.let { it.asPsiType(support, TypeMappingMode.getOptimalModeForValueParameter(it), this) }
                     ?.let(TypeConversionUtil::erasure)
                     ?: nonExistent()
             declaration is KtObjectDeclaration ->
@@ -547,7 +550,7 @@ private class KtUltraLightField(
                         (declaration.resolve() as? PropertyDescriptor)?.isVar == true -> TypeMappingMode.getOptimalModeForValueParameter(it)
                         else -> TypeMappingMode.getOptimalModeForReturnType(it, false)
                     }
-                    it.asPsiType(declaration, support, mode, this)
+                    it.asPsiType(support, mode, this)
                 } ?: PsiType.NULL
         }
     }
@@ -654,7 +657,7 @@ internal class KtUltraLightParameter(
         }
     }
     private val _type: PsiType by lazyPub {
-        kotlinType?.let { it.asPsiType(kotlinOrigin, support, TypeMappingMode.getOptimalModeForValueParameter(it), this) } ?: PsiType.NULL
+        kotlinType?.let { it.asPsiType(support, TypeMappingMode.getOptimalModeForValueParameter(it), this) } ?: PsiType.NULL
     }
 
     override fun getType(): PsiType = _type

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightUtils.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.elements.KotlinLightTypeParameterListBuilder
 import org.jetbrains.kotlin.codegen.ClassBuilderMode
 import org.jetbrains.kotlin.codegen.signature.BothSignatureWriter
+import org.jetbrains.kotlin.codegen.signature.JvmSignatureWriter
 import org.jetbrains.kotlin.codegen.state.IncompatibleClassTracker
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.config.JvmTarget
@@ -24,10 +25,8 @@ import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.descriptors.ValueDescriptor
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtTypeParameterListOwner
 import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.types.KotlinType
 import java.text.StringCharacterIterator
 
@@ -44,7 +43,7 @@ internal fun buildTypeParameterList(
                     KotlinLightReferenceListBuilder(manager, PsiReferenceList.Role.EXTENDS_BOUNDS_LIST)
                 if (ktParam.extendsBound != null || declaration.typeConstraints.isNotEmpty()) {
                     val boundTypes = (ktParam.resolve() as? TypeParameterDescriptor)?.upperBounds.orEmpty()
-                        .mapNotNull { it.asPsiType(ktParam, support, TypeMappingMode.DEFAULT, this) as? PsiClassType }
+                        .mapNotNull { it.asPsiType(support, TypeMappingMode.DEFAULT, this) as? PsiClassType }
                     val hasDefaultBound = boundTypes.size == 1 && boundTypes[0].equalsToText(CommonClassNames.JAVA_LANG_OBJECT)
                     if (!hasDefaultBound) {
                         boundTypes.forEach(boundList::addReference)
@@ -75,16 +74,19 @@ internal fun KtDeclaration.resolve() = LightClassGenerationSupport.getInstance(p
 
 // copy-pasted from kotlinInternalUastUtils.kt and post-processed
 internal fun KotlinType.asPsiType(
-    declaration: KtDeclaration,
     support: UltraLightSupport,
     mode: TypeMappingMode,
     psiContext: PsiElement
-): PsiType {
-    val typeFqName = constructor.declarationDescriptor?.fqNameSafe?.asString()
-    if (typeFqName == "kotlin.Unit" && declaration is KtFunction) return PsiType.VOID
+): PsiType = support.mapType(psiContext) { typeMapper, signatureWriter ->
+    typeMapper.mapType(this, signatureWriter, mode)
+}
 
+internal fun UltraLightSupport.mapType(
+    psiContext: PsiElement,
+    mapTypeToSignatureWriter: (KotlinTypeMapper, JvmSignatureWriter) -> Unit
+): PsiType {
     val signatureWriter = BothSignatureWriter(BothSignatureWriter.Mode.TYPE)
-    typeMapper(support).mapType(this, signatureWriter, mode)
+    mapTypeToSignatureWriter(typeMapper(this), signatureWriter)
     val signature = StringCharacterIterator(signatureWriter.toString())
 
     val javaType = SignatureParsing.parseTypeString(signature, StubBuildingVisitor.GUESSING_MAPPER)