FIR/UAST: commonize type reference expression

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/BaseKotlinConverter.kt

@@ -5,14 +5,17 @@
 
 package org.jetbrains.uast.kotlin
 
-import org.jetbrains.kotlin.psi.KtEscapeStringTemplateEntry
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtStringTemplateEntry
-import org.jetbrains.kotlin.psi.KtStringTemplateEntryWithExpression
 import org.jetbrains.uast.*
 
 interface BaseKotlinConverter {
 
+    fun convertReceiverParameter(receiver: KtTypeReference): UParameter? {
+        val call = (receiver.parent as? KtCallableDeclaration) ?: return null
+        if (call.receiverTypeReference != receiver) return null
+        return call.toUElementOfType<UMethod>()?.uastParameters?.firstOrNull()
+    }
+
     fun convertExpression(
         expression: KtExpression,
         givenParent: UElement?,

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/BaseKotlinUastResolveProviderService.kt

@@ -8,6 +8,7 @@ package org.jetbrains.uast.kotlin
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtTypeReference
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UExpression
 
@@ -20,6 +21,8 @@ interface BaseKotlinUastResolveProviderService {
 
     fun resolveToDeclaration(ktExpression: KtExpression): PsiElement?
 
+    fun resolveToType(ktTypeReference: KtTypeReference, source: UElement): PsiType?
+
     fun getExpressionType(uExpression: UExpression): PsiType?
 
     fun evaluate(uExpression: UExpression): Any?

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/expressions/KotlinUTypeReferenceExpression.kt

@@ -11,20 +11,14 @@ import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UTypeReferenceExpression
 import org.jetbrains.uast.UastErrorType
 
-class FirKotlinUTypeReferenceExpression(
+class KotlinUTypeReferenceExpression(
-    override val sourcePsi: KtTypeReference,
+    override val sourcePsi: KtTypeReference?,
     givenParent: UElement?,
-    private val typeSupplier: (() -> PsiType)? = null,
+    private val typeSupplier: (() -> PsiType)? = null
-) : KotlinAbstractUExpression(givenParent), UTypeReferenceExpression {
+) : KotlinAbstractUExpression(givenParent), UTypeReferenceExpression, KotlinUElementWithType {
     override val type: PsiType by lz {
-        typeSupplier?.invoke() ?: sourcePsi.toPsiType(uastParent ?: this)
+        typeSupplier?.invoke()
+            ?: sourcePsi?.let { baseResolveProviderService.resolveToType(it, uastParent ?: this) }
+            ?: UastErrorType
     }
 }
-
-private fun KtTypeReference?.toPsiType(
-    source: UElement,
-): PsiType {
-    if (this == null) return UastErrorType
-    // TODO: use type conversions in firLightUtils.kt
-    return PsiType.NULL
-}

plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/FirKotlinConverter.kt

@@ -200,6 +200,11 @@ internal object FirKotlinConverter : BaseKotlinConverter {
                     element.expression?.let { convertExpression(it, givenParent, requiredTypes) }
                         ?: expr<UExpression> { UastEmptyExpression(givenParent) }
                 }
+                is KtTypeReference ->
+                    requiredTypes.accommodate(
+                        alternative { KotlinUTypeReferenceExpression(element, givenParent) },
+                        alternative { convertReceiverParameter(element) }
+                    ).firstOrNull()
                 is KtImportDirective -> {
                     el<UImportStatement>(build(::KotlinUImportStatement))
                 }

plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/FirKotlinUastResolveProviderService.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtReferenceExpression
+import org.jetbrains.kotlin.psi.KtTypeReference
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UExpression
 
@@ -36,6 +37,11 @@ interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderSer
         }
     }
 
+    override fun resolveToType(ktTypeReference: KtTypeReference, source: UElement): PsiType? {
+        // TODO: use type conversions in firLightUtils.kt
+        return PsiType.NULL
+    }
+
     override fun getExpressionType(uExpression: UExpression): PsiType? {
         val ktExpression = uExpression.sourcePsi as? KtExpression ?: return null
         analyseWithCustomToken(ktExpression, AlwaysAccessibleValidityTokenFactory) {

plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUClass.kt

@@ -38,7 +38,7 @@ sealed class AbstractFirKotlinUClass(
 
     override val uastSuperTypes: List<UTypeReferenceExpression>
         get() = ktClass?.superTypeListEntries.orEmpty().mapNotNull { it.typeReference }.map {
-            FirKotlinUTypeReferenceExpression(it, this)
+            KotlinUTypeReferenceExpression(it, this)
         }
 
     // TODO: delegateExpressions

plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUMethod.kt

@@ -84,7 +84,7 @@ open class FirKotlinUMethod(
 
     override val returnTypeReference: UTypeReferenceExpression? by lz {
         (sourcePsi as? KtCallableDeclaration)?.typeReference?.let {
-            FirKotlinUTypeReferenceExpression(it, this) { javaPsi.returnType ?: UastErrorType }
+            KotlinUTypeReferenceExpression(it, this) { javaPsi.returnType ?: UastErrorType }
         }
     }
 

plugins/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/declarations/FirKotlinUVariable.kt

@@ -46,7 +46,7 @@ sealed class AbstractFirKotlinUVariable(
 
     override val typeReference: UTypeReferenceExpression? by lz {
         (sourcePsi as? KtCallableDeclaration)?.typeReference?.let {
-            FirKotlinUTypeReferenceExpression(it, this) { type }
+            KotlinUTypeReferenceExpression(it, this) { type }
         }
     }
 

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinConverter.kt

@@ -133,12 +133,6 @@ internal object KotlinConverter : BaseKotlinConverter {
         }}
     }
 
-    internal fun convertReceiverParameter(receiver: KtTypeReference): UParameter? {
-        val call = (receiver.parent as? KtCallableDeclaration) ?: return null
-        if (call.receiverTypeReference != receiver) return null
-        return call.toUElementOfType<UMethod>()?.uastParameters?.firstOrNull()
-    }
-
     var forceUInjectionHost = Registry.`is`("kotlin.uast.force.uinjectionhost", false)
         @TestOnly
         set(value) {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastResolveProviderService.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtTypeReference
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.constants.UnsignedErrorValueTypeConstant
 import org.jetbrains.kotlin.types.TypeUtils
@@ -35,6 +36,10 @@ interface KotlinUastResolveProviderService : BaseKotlinUastResolveProviderServic
         return resolveToDeclarationImpl(ktExpression)
     }
 
+    override fun resolveToType(ktTypeReference: KtTypeReference, source: UElement): PsiType? {
+        return ktTypeReference.toPsiType(source)
+    }
+
     override fun getExpressionType(uExpression: UExpression): PsiType? {
         val ktElement = uExpression.sourcePsi as? KtExpression ?: return null
         val ktType = ktElement.analyze()[BindingContext.EXPRESSION_TYPE_INFO, ktElement]?.type ?: return null

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt

@@ -91,7 +91,6 @@ abstract class AbstractKotlinUVariable(givenParent: UElement?) : KotlinAbstractU
         annotations
     }
 
-
     protected abstract fun acceptsAnnotationTarget(target: AnnotationUseSiteTarget?): Boolean
 
     override val typeReference: UTypeReferenceExpression? by lz {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUTypeReferenceExpression.kt

@@ -1,16 +0,0 @@
-package org.jetbrains.uast.kotlin
-
-import com.intellij.psi.PsiType
-import org.jetbrains.kotlin.psi.KtTypeReference
-import org.jetbrains.uast.UElement
-import org.jetbrains.uast.UTypeReferenceExpression
-
-class KotlinUTypeReferenceExpression(
-    override val sourcePsi: KtTypeReference?,
-    givenParent: UElement?,
-    private val typeSupplier: (() -> PsiType)? = null
-) : KotlinAbstractUExpression(givenParent), UTypeReferenceExpression, KotlinUElementWithType {
-    override val type: PsiType by lz {
-        typeSupplier?.invoke() ?: sourcePsi.toPsiType(uastParent ?: this)
-    }
-}