FIR/UAST: commonize KotlinUElementWithType

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

@@ -6,8 +6,10 @@
 package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UExpression
 
 interface BaseKotlinUastResolveProviderService {
     fun isJvmElement(psiElement: PsiElement): Boolean
@@ -15,4 +17,6 @@ interface BaseKotlinUastResolveProviderService {
     fun convertParent(uElement: UElement): UElement?
 
     fun resolveToDeclaration(ktExpression: KtExpression): PsiElement?
+
+    fun getExpressionType(uExpression: UExpression): PsiType?
 }

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

@@ -20,7 +20,7 @@ abstract class KotlinAbstractUElement(
         psi?.let { UastFacade.findPlugin(it) }
     }
 
-    protected val baseResolveProviderService: BaseKotlinUastResolveProviderService by lz {
+    open val baseResolveProviderService: BaseKotlinUastResolveProviderService by lz {
         baseResolveProviderServiceSupplier?.get()
             ?: psi?.project?.let { ServiceManager.getService(it, BaseKotlinUastResolveProviderService::class.java) }
             ?: error("${BaseKotlinUastResolveProviderService::class.java.name} is not available")

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

@@ -0,0 +1,17 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.uast.kotlin
+
+import com.intellij.psi.PsiType
+import org.jetbrains.uast.UExpression
+
+interface KotlinUElementWithType : UExpression {
+    val baseResolveProviderService: BaseKotlinUastResolveProviderService
+
+    override fun getExpressionType(): PsiType? {
+        return baseResolveProviderService.getExpressionType(this)
+    }
+}

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

@@ -6,12 +6,14 @@
 package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.idea.frontend.api.analyseWithCustomToken
 import org.jetbrains.kotlin.idea.frontend.api.tokens.AlwaysAccessibleValidityTokenFactory
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtReferenceExpression
 import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UExpression
 
 interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderService {
     override fun convertParent(uElement: UElement): UElement? {
@@ -29,4 +31,8 @@ interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderSer
                 return null
         }
     }
+
+    override fun getExpressionType(uExpression: UExpression): PsiType? {
+        TODO("Not yet implemented")
+    }
 }

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

@@ -6,6 +6,7 @@
 package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
@@ -13,6 +14,7 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UExpression
 
 interface KotlinUastResolveProviderService : BaseKotlinUastResolveProviderService {
     fun getBindingContext(element: KtElement): BindingContext
@@ -27,4 +29,10 @@ interface KotlinUastResolveProviderService : BaseKotlinUastResolveProviderServic
     override fun resolveToDeclaration(ktExpression: KtExpression): PsiElement? {
         return resolveToDeclarationImpl(ktExpression)
     }
+
+    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
+        return ktType.toPsiType(uExpression, ktElement, boxed = false)
+    }
 }

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

@@ -16,21 +16,12 @@
 
 package org.jetbrains.uast.kotlin
 
-import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.constants.UnsignedErrorValueTypeConstant
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.uast.UExpression
 
-interface KotlinUElementWithType : UExpression {
-    override fun getExpressionType(): PsiType? {
-        val ktElement = sourcePsi as? KtExpression ?: return null
-        val ktType = ktElement.analyze()[BindingContext.EXPRESSION_TYPE_INFO, ktElement]?.type ?: return null
-        return ktType.toPsiType(this, ktElement, boxed = false)
-    }
-}
-
 interface KotlinEvaluatableUElement : UExpression {
     override fun evaluate(): Any? {
         val ktElement = sourcePsi as? KtExpression ?: return null
@@ -40,4 +31,4 @@ interface KotlinEvaluatableUElement : UExpression {
 
         return compileTimeConst?.getValue(TypeUtils.NO_EXPECTED_TYPE)
     }
-}
+}

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

@@ -25,8 +25,8 @@ import org.jetbrains.uast.kotlin.internal.KotlinFakeUElement
 import org.jetbrains.uast.kotlin.internal.toSourcePsiFakeAware
 
 class KotlinUReturnExpression(
-        override val sourcePsi: KtReturnExpression,
+    override val sourcePsi: KtReturnExpression,
-        givenParent: UElement?
+    givenParent: UElement?
 ) : KotlinAbstractUExpression(givenParent), UReturnExpression, KotlinUElementWithType {
     override val returnExpression by lz { KotlinConverter.convertOrNull(sourcePsi.returnedExpression, this) }
 }
@@ -40,8 +40,13 @@ class KotlinUImplicitReturnExpression(
     override lateinit var returnExpression: UExpression
         internal set
 
+    // Due to the lack of [psi], (lazily) delegate to the one in [returnExpression]
+    override val baseResolveProviderService: BaseKotlinUastResolveProviderService by lz {
+        (returnExpression as KotlinAbstractUElement).baseResolveProviderService
+    }
+
     override fun unwrapToSourcePsi(): List<PsiElement> {
         return returnExpression.toSourcePsiFakeAware()
     }
 
-}
+}