FIR/UAST: commonize UImportStatement

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

@@ -6,6 +6,7 @@
 package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.uast.UElement
 
 interface BaseKotlinUastResolveProviderService {
@@ -13,4 +14,5 @@ interface BaseKotlinUastResolveProviderService {
 
     fun convertParent(uElement: UElement): UElement?
 
+    fun resolveToDeclaration(ktExpression: KtExpression): PsiElement?
 }

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/declarations/KotlinUImportStatement.kt

@@ -6,8 +6,6 @@
 package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.idea.frontend.api.analyse
-import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtImportDirective
 import org.jetbrains.kotlin.psi.KtReferenceExpression
@@ -15,20 +13,22 @@ import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UImportStatement
 import org.jetbrains.uast.USimpleNameReferenceExpression
+import org.jetbrains.uast.kotlin.internal.DelegatedMultiResolve
 
-class FirKotlinUImportStatement(
+class KotlinUImportStatement(
     override val psi: KtImportDirective,
     givenParent: UElement?
-) : KotlinAbstractUElement(givenParent), UImportStatement {
+) : KotlinAbstractUElement(givenParent), UImportStatement, DelegatedMultiResolve {
+
     override val javaPsi: PsiElement? = null
 
     override val sourcePsi: KtImportDirective = psi
 
     override val isOnDemand: Boolean = sourcePsi.isAllUnder
 
-    private val importRef: FirImportReference? by lz {
+    private val importRef: ImportReference? by lz {
         sourcePsi.importedReference?.let {
-            FirImportReference(it, sourcePsi.name ?: sourcePsi.text, this, sourcePsi)
+            ImportReference(it, sourcePsi.name ?: sourcePsi.text, this, sourcePsi)
         }
     }
 
@@ -36,7 +36,7 @@ class FirKotlinUImportStatement(
 
     override fun resolve(): PsiElement? = importRef?.resolve()
 
-    private class FirImportReference(
+    private class ImportReference(
         override val psi: KtExpression,
         override val identifier: String,
         givenParent: UElement?,
@@ -50,9 +50,7 @@ class FirKotlinUImportStatement(
 
         override fun resolve(): PsiElement? {
             val reference = sourcePsi.getQualifiedElementSelector() as? KtReferenceExpression ?: return null
-            analyse(reference) {
+            return baseResolveProviderService?.resolveToDeclaration(reference)
-                return reference.mainReference.resolve()
-            }
         }
     }
 }

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/internal/DelegatedMultiResolve.kt

@@ -0,0 +1,16 @@
+/*
+ * 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.internal
+
+import com.intellij.psi.PsiSubstitutor
+import com.intellij.psi.ResolveResult
+import com.intellij.psi.infos.CandidateInfo
+import org.jetbrains.uast.UMultiResolvable
+import org.jetbrains.uast.UResolvable
+
+interface DelegatedMultiResolve : UMultiResolvable, UResolvable {
+    override fun multiResolve(): Iterable<ResolveResult> = listOfNotNull(resolve()?.let { CandidateInfo(it, PsiSubstitutor.EMPTY) })
+}

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

@@ -190,7 +190,7 @@ internal object FirKotlinConverter {
                     convertExpression(element, givenParent, requiredTypes)
                 }
                 is KtImportDirective -> {
-                    el<UImportStatement>(build(::FirKotlinUImportStatement))
+                    el<UImportStatement>(build(::KotlinUImportStatement))
                 }
                 else -> null
             }

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

@@ -9,22 +9,13 @@ import com.intellij.lang.Language
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.KotlinLanguage
-import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtParameter
-import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.uast.DEFAULT_TYPES_LIST
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UExpression
 import org.jetbrains.uast.UastLanguagePlugin
 import org.jetbrains.uast.kotlin.FirKotlinConverter.convertDeclarationOrElement
 
-interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderService {
-    override fun convertParent(uElement: UElement): UElement? {
-        TODO("Not yet implemented")
-    }
-}
-
 class FirKotlinUastLanguagePlugin : UastLanguagePlugin {
     override val priority: Int = 10
 

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

@@ -0,0 +1,31 @@
+/*
+ * 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.PsiElement
+import org.jetbrains.kotlin.idea.frontend.api.analyse
+import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtReferenceExpression
+import org.jetbrains.uast.UElement
+
+interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderService {
+    override fun convertParent(uElement: UElement): UElement? {
+        TODO("Not yet implemented")
+    }
+
+    override fun resolveToDeclaration(ktExpression: KtExpression): PsiElement? {
+        when (ktExpression) {
+            is KtReferenceExpression -> {
+                analyse(ktExpression) {
+                    return ktExpression.mainReference.resolve()
+                }
+            }
+            else ->
+                return null
+        }
+    }
+}

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

@@ -62,6 +62,10 @@ interface KotlinUastResolveProviderService : BaseKotlinUastResolveProviderServic
     override fun convertParent(uElement: UElement): UElement? {
         return convertParentImpl(uElement)
     }
+
+    override fun resolveToDeclaration(ktExpression: KtExpression): PsiElement? {
+        return resolveToDeclarationImpl(ktExpression)
+    }
 }
 
 var PsiElement.destructuringDeclarationInitializer: Boolean? by UserDataProperty(Key.create("kotlin.uast.destructuringDeclarationInitializer"))

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

@@ -64,7 +64,7 @@ abstract class KotlinUAnnotationBase<T : KtCallElement>(
     protected abstract fun computeClassDescriptor(): ClassDescriptor?
 
     override fun resolve(): PsiClass? = computeClassDescriptor()?.let {
-        sourcePsi.calleeExpression?.let { ktExpression -> resolveToDeclaration(ktExpression, it) }
+        sourcePsi.calleeExpression?.let { ktExpression -> resolveToDeclarationImpl(ktExpression, it) }
     } as? PsiClass
 
     override fun findAttributeValue(name: String?): UExpression? =

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

@@ -1,68 +0,0 @@
-/*
- * Copyright 2010-2016 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.jetbrains.uast.kotlin
-
-import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtImportDirective
-import org.jetbrains.kotlin.psi.KtReferenceExpression
-import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
-import org.jetbrains.uast.UElement
-import org.jetbrains.uast.UImportStatement
-import org.jetbrains.uast.USimpleNameReferenceExpression
-import org.jetbrains.uast.kotlin.internal.DelegatedMultiResolve
-
-class KotlinUImportStatement(
-        override val psi: KtImportDirective,
-        givenParent: UElement?
-) : KotlinAbstractUElement(givenParent), UImportStatement, DelegatedMultiResolve {
-
-    override val javaPsi = null
-
-    override val sourcePsi = psi
-
-    override val isOnDemand: Boolean
-        get() = psi.isAllUnder
-
-    private val importRef by lz {
-        psi.importedReference?.let {
-            ImportReference(it, psi.name ?: psi.text, this, psi)
-        }
-    }
-    
-    override val importReference: UElement?
-        get() = importRef
-
-    override fun resolve() = importRef?.resolve()
-
-    private class ImportReference(
-            override val psi: KtExpression,
-            override val identifier: String,
-            givenParent: UElement?,
-            private val importDirective: KtImportDirective
-    ) : KotlinAbstractUExpression(givenParent), USimpleNameReferenceExpression {
-        override val resolvedName: String?
-            get() = identifier
-
-        override fun asRenderString(): String = importDirective.importedFqName?.asString() ?: psi.text
-
-        override fun resolve(): PsiElement? {
-            val reference = psi.getQualifiedElementSelector() as? KtReferenceExpression ?: return null
-            return resolveToDeclaration(reference)
-        }
-    }
-}

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

@@ -46,8 +46,8 @@ class KotlinUCallableReferenceExpression(
     override val resolvedName: String?
         get() = (resolve() as? PsiNamedElement)?.name
 
-    override fun resolve(): PsiElement? = resolveToDeclaration(sourcePsi.callableReference)
+    override fun resolve(): PsiElement? = baseResolveProviderService?.resolveToDeclaration(sourcePsi.callableReference)
 
     override fun multiResolve(): Iterable<ResolveResult> = getResolveResultVariants(sourcePsi.callableReference)
 
-}
+}

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

@@ -144,7 +144,8 @@ class KotlinUFunctionCallExpression(
             }
 
             val ktNameReferenceExpression = sourcePsi.calleeExpression as? KtNameReferenceExpression ?: return null
-            val localCallableDeclaration = resolveToDeclaration(ktNameReferenceExpression) as? PsiVariable ?: return null
+            val localCallableDeclaration =
+                baseResolveProviderService?.resolveToDeclaration(ktNameReferenceExpression) as? PsiVariable ?: return null
             if (localCallableDeclaration !is PsiLocalVariable && localCallableDeclaration !is PsiParameter) return null
 
             // an implicit receiver for variables calls (KT-25524)

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

@@ -19,7 +19,6 @@ package org.jetbrains.uast.kotlin
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNamedElement
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
-import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.uast.UCallExpression
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UQualifiedReferenceExpression
@@ -35,7 +34,7 @@ class KotlinUQualifiedReferenceExpression(
     override val selector by lz { KotlinConverter.convertOrEmpty(sourcePsi.selectorExpression, this) }
     override val accessType = UastQualifiedExpressionAccessType.SIMPLE
 
-    override fun resolve(): PsiElement? = sourcePsi.selectorExpression?.let { resolveToDeclaration(it) }
+    override fun resolve(): PsiElement? = sourcePsi.selectorExpression?.let { baseResolveProviderService?.resolveToDeclaration(it) }
 
     override val resolvedName: String?
         get() = (resolve() as? PsiNamedElement)?.name

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

@@ -37,6 +37,6 @@ class KotlinUSafeQualifiedExpression(
     override val resolvedName: String?
         get() = (resolve() as? PsiNamedElement)?.name
 
-    override fun resolve(): PsiElement? = sourcePsi.selectorExpression?.let { resolveToDeclaration(it) }
+    override fun resolve(): PsiElement? = sourcePsi.selectorExpression?.let { baseResolveProviderService?.resolveToDeclaration(it) }
     override fun multiResolve(): Iterable<ResolveResult> = getResolveResultVariants(sourcePsi.selectorExpression)
-}
+}

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

@@ -38,7 +38,7 @@ open class KotlinUSimpleReferenceExpression(
     override val sourcePsi: KtSimpleNameExpression,
     givenParent: UElement?
 ) : KotlinAbstractUExpression(givenParent), USimpleNameReferenceExpression, KotlinUElementWithType, KotlinEvaluatableUElement {
-    private val resolvedDeclaration: PsiElement? by lz { resolveToDeclaration(sourcePsi) }
+    private val resolvedDeclaration: PsiElement? by lz { baseResolveProviderService?.resolveToDeclaration(sourcePsi) }
 
     override val identifier get() = sourcePsi.getReferencedName()
 
@@ -174,7 +174,7 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
     private val resolved by lazy {
         when (val resultingDescriptor = sourcePsi.getResolvedCall(sourcePsi.analyze())?.descriptorForResolveViaConstructor()) {
             is ConstructorDescriptor -> {
-                sourcePsi.calleeExpression?.let { resolveToDeclaration(it, resultingDescriptor.constructedClass) }
+                sourcePsi.calleeExpression?.let { resolveToDeclarationImpl(it, resultingDescriptor.constructedClass) }
             }
             is SamConstructorDescriptor ->
                 (resultingDescriptor.returnType?.getFunctionalInterfaceType(this, sourcePsi) as? PsiClassType)?.resolve()
@@ -213,4 +213,4 @@ class KotlinStringUSimpleReferenceExpression(
     override fun resolve() = null
     override val resolvedName: String?
         get() = identifier
-}
+}

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinInternalUastUtils.kt

@@ -311,17 +311,17 @@ internal fun resolveToPsiMethod(
     }
 }
 
-internal fun resolveToDeclaration(sourcePsi: KtExpression): PsiElement? =
+internal fun resolveToDeclarationImpl(sourcePsi: KtExpression): PsiElement? =
     when (sourcePsi) {
         is KtSimpleNameExpression ->
             sourcePsi.analyze()[BindingContext.REFERENCE_TARGET, sourcePsi]
-                ?.let { resolveToDeclaration(sourcePsi, it) }
+                ?.let { resolveToDeclarationImpl(sourcePsi, it) }
         else ->
             sourcePsi.getResolvedCall(sourcePsi.analyze())?.resultingDescriptor
-                ?.let { descriptor -> resolveToDeclaration(sourcePsi, descriptor) }
+                ?.let { descriptor -> resolveToDeclarationImpl(sourcePsi, descriptor) }
     }
 
-internal fun resolveToDeclaration(sourcePsi: KtExpression, declarationDescriptor: DeclarationDescriptor): PsiElement? {
+internal fun resolveToDeclarationImpl(sourcePsi: KtExpression, declarationDescriptor: DeclarationDescriptor): PsiElement? {
     declarationDescriptor.toSource()?.getMaybeLightElement(sourcePsi)?.let { return it }
 
     var declarationDescriptor = declarationDescriptor
@@ -349,7 +349,7 @@ internal fun resolveToDeclaration(sourcePsi: KtExpression, declarationDescriptor
     resolveDeserialized(sourcePsi, declarationDescriptor, sourcePsi.readWriteAccess())?.let { return it }
 
     if (declarationDescriptor is ValueParameterDescriptor) {
-        val parentDeclaration = resolveToDeclaration(sourcePsi, declarationDescriptor.containingDeclaration)
+        val parentDeclaration = resolveToDeclarationImpl(sourcePsi, declarationDescriptor.containingDeclaration)
         if (parentDeclaration is PsiClass && parentDeclaration.isAnnotationType) {
             parentDeclaration.findMethodsByName(declarationDescriptor.name.asString(), false).firstOrNull()?.let { return it }
         }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinUastMultiresolve.kt

@@ -13,16 +13,14 @@ import com.intellij.psi.ResolveResult
 import com.intellij.psi.infos.CandidateInfo
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.uast.UMultiResolvable
-import org.jetbrains.uast.UResolvable
 import org.jetbrains.uast.kotlin.KotlinUastResolveProviderService
-import org.jetbrains.uast.kotlin.resolveToDeclaration
+import org.jetbrains.uast.kotlin.resolveToDeclarationImpl
 
 
 internal fun getReferenceVariants(ktElement: KtExpression, nameHint: String): Sequence<PsiElement> =
     ServiceManager.getService(ktElement.project, KotlinUastResolveProviderService::class.java)
         .getReferenceVariants(ktElement, nameHint)
-        .mapNotNull { resolveToDeclaration(ktElement, it) }
+        .mapNotNull { resolveToDeclarationImpl(ktElement, it) }
 
 internal fun getResolveResultVariants(ktExpression: KtExpression?): Iterable<ResolveResult> {
     ktExpression ?: return emptyList()
@@ -41,10 +39,6 @@ internal fun KtElement.multiResolveResults(): Sequence<ResolveResult> =
         }
     }
 
-interface DelegatedMultiResolve : UMultiResolvable, UResolvable {
-    override fun multiResolve(): Iterable<ResolveResult> = listOfNotNull(resolve()?.let { CandidateInfo(it, PsiSubstitutor.EMPTY) })
-}
-
 class TypedResolveResult<T : PsiElement>(element: T) : CandidateInfo(element, PsiSubstitutor.EMPTY) {
     @Suppress("UNCHECKED_CAST")
     override fun getElement(): T = super.getElement() as T