Uast: resolving local calls (KT-36275)

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

@@ -39,7 +39,6 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
@@ -52,10 +51,7 @@ import org.jetbrains.uast.kotlin.declarations.KotlinUIdentifier
 import org.jetbrains.uast.kotlin.declarations.KotlinUMethod
 import org.jetbrains.uast.kotlin.declarations.KotlinUMethodWithFakeLightDelegate
 import org.jetbrains.uast.kotlin.expressions.*
-import org.jetbrains.uast.kotlin.psi.UastFakeLightMethod
+import org.jetbrains.uast.kotlin.psi.*
-import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
-import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameterBase
-import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
 interface KotlinUastResolveProviderService {
     fun getBindingContext(element: KtElement): BindingContext
@@ -167,6 +163,8 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
                 KotlinConverter.convertNonLocalProperty(element, null, requiredTypes) as Sequence<T>
             element is KtParameter -> KotlinConverter.convertParameter(element, null, requiredTypes) as Sequence<T>
             element is KtClassOrObject -> KotlinConverter.convertClassOrObject(element, null, requiredTypes) as Sequence<T>
+            element is UastFakeLightPrimaryConstructor ->
+                KotlinConverter.convertFakeLightConstructorAlternatices(element, null, requiredTypes) as Sequence<T>
             else -> sequenceOf(convertElementWithParent(element, requiredTypes.nonEmptyOr(DEFAULT_TYPES_LIST)) as? T).filterNotNull()
         }
 }
@@ -507,7 +505,15 @@ internal object KotlinConverter {
         return with(expectedTypes) {
             when (original) {
                 is KtLightMethod -> el<UMethod>(build(KotlinUMethod.Companion::create))   // .Companion is needed because of KT-13934
-                is UastFakeLightMethod -> el<UMethod> { KotlinUMethodWithFakeLightDelegate(original.original, original, givenParent) }
+                is UastFakeLightMethod -> el<UMethod> {
+                    val ktFunction = original.original
+                    if (ktFunction.isLocal)
+                        convertDeclaration(ktFunction, givenParent, expectedTypes)
+                    else
+                        KotlinUMethodWithFakeLightDelegate(ktFunction, original, givenParent)
+                }
+                is UastFakeLightPrimaryConstructor ->
+                    convertFakeLightConstructorAlternatices(original, givenParent, expectedTypes).firstOrNull()
                 is KtLightClass -> when (original.kotlinOrigin) {
                     is KtEnumEntry -> el<UEnumConstant> {
                         convertEnumEntry(original.kotlinOrigin as KtEnumEntry, givenParent)
@@ -586,9 +592,20 @@ internal object KotlinConverter {
         }
     }
 
+    internal fun convertFakeLightConstructorAlternatices(
+        original: UastFakeLightPrimaryConstructor,
+        givenParent: UElement?,
+        expectedTypes: Array<out Class<out UElement>>
+    ): Sequence<UElement> {
+        return expectedTypes.accommodate(
+            alternative { convertDeclaration(original.original, givenParent, expectedTypes) as? UClass },
+            alternative { KotlinConstructorUMethod(original.original, original, original.original, givenParent) }
+        )
+    }
+
     private fun getLightClassForFakeMethod(original: KtFunction): KtLightClass? {
         if (original.isLocal) return null
-        return (original.containingClassOrObject?.toLightClass() ?: original.containingKtFile.findFacadeClass())
+        return getContainingLightClass(original)
     }
 
     fun convertDeclarationOrElement(

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

@@ -200,15 +200,22 @@ open class KotlinUClass private constructor(
 
 open class KotlinConstructorUMethod(
     private val ktClass: KtClassOrObject?,
-    override val psi: KtLightMethod,
+    override val psi: PsiMethod,
+    kotlinOrigin: KtDeclaration?,
     givenParent: UElement?
-) : KotlinUMethod(psi, psi.kotlinOrigin, givenParent) {
+) : KotlinUMethod(psi, kotlinOrigin, givenParent) {
+
+    constructor(
+        ktClass: KtClassOrObject?,
+        psi: KtLightMethod,
+        givenParent: UElement?
+    ):this(ktClass, psi, psi.kotlinOrigin, givenParent)
 
     val isPrimary: Boolean
-        get() = psi.kotlinOrigin.let { it is KtPrimaryConstructor || it is KtClassOrObject }
+        get() = sourcePsi.let { it is KtPrimaryConstructor || it is KtClassOrObject }
 
     override val uastBody: UExpression? by lz {
-        val delegationCall: KtCallElement? = psi.kotlinOrigin.let {
+        val delegationCall: KtCallElement? = sourcePsi.let {
             when {
                 isPrimary -> ktClass?.superTypeListEntries?.firstIsInstanceOrNull<KtSuperTypeCallEntry>()
                 it is KtSecondaryConstructor -> it.getDelegationCall()
@@ -232,7 +239,7 @@ open class KotlinConstructorUMethod(
     override val uastAnchor: KotlinUIdentifier by lazy {
         KotlinUIdentifier(
             psi.nameIdentifier,
-            if (isPrimary) ktClass?.nameIdentifier else (psi.kotlinOrigin as? KtSecondaryConstructor)?.getConstructorKeyword(),
+            if (isPrimary) ktClass?.nameIdentifier else (sourcePsi as? KtSecondaryConstructor)?.getConstructorKeyword(),
             this
         )
     }
@@ -241,7 +248,7 @@ open class KotlinConstructorUMethod(
 
     open protected fun getBodyExpressions(): List<KtExpression> {
         if (isPrimary) return getInitializers()
-        val bodyExpression = (psi.kotlinOrigin as? KtFunction)?.bodyExpression ?: return emptyList()
+        val bodyExpression = (sourcePsi as? KtFunction)?.bodyExpression ?: return emptyList()
         if (bodyExpression is KtBlockExpression) return bodyExpression.statements
         return listOf(bodyExpression)
     }

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

@@ -27,7 +27,9 @@ import com.intellij.psi.impl.compiled.SignatureParsing
 import com.intellij.psi.impl.compiled.StubBuildingVisitor
 import com.intellij.psi.util.PsiTypesUtil
 import org.jetbrains.kotlin.asJava.LightClassUtil
+import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.asJava.elements.FakeFileForLightClass
+import org.jetbrains.kotlin.asJava.findFacadeClass
 import org.jetbrains.kotlin.asJava.toLightClass
 import org.jetbrains.kotlin.asJava.toLightElements
 import org.jetbrains.kotlin.builtins.isBuiltinFunctionalTypeOrSubtype
@@ -45,6 +47,7 @@ import org.jetbrains.kotlin.metadata.jvm.JvmProtoBuf
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMemberSignature
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
@@ -69,6 +72,8 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.expressions.KotlinLocalFunctionUVariable
+import org.jetbrains.uast.kotlin.psi.UastFakeLightMethod
+import org.jetbrains.uast.kotlin.psi.UastFakeLightPrimaryConstructor
 import java.lang.ref.WeakReference
 import java.text.StringCharacterIterator
 
@@ -97,17 +102,29 @@ internal fun resolveSource(context: KtElement, descriptor: DeclarationDescriptor
         && source is KtClassOrObject && source.primaryConstructor == null
         && source.secondaryConstructors.isEmpty()
     ) {
-        return source.toLightClass()?.constructors?.firstOrNull()
+        val lightClass = source.toLightClass() ?: return null
+        lightClass.constructors.firstOrNull()?.let { return it }
+        if (source.isLocal) {
+            return UastFakeLightPrimaryConstructor(source, lightClass)
+        }
+        return null
     }
 
     return when (source) {
-        is KtFunction -> LightClassUtil.getLightClassMethod(source)
+        is KtFunction ->
+            if (source.isLocal)
+                getContainingLightClass(source)?.let { UastFakeLightMethod(source, it) }
+            else
+                LightClassUtil.getLightClassMethod(source)
         is PsiMethod -> source
         null -> resolveDeserialized(context, descriptor) as? PsiMethod
         else -> null
     }
 }
 
+internal fun getContainingLightClass(original: KtDeclaration): KtLightClass? =
+    (original.containingClassOrObject?.toLightClass() ?: original.containingKtFile.findFacadeClass())
+
 internal fun resolveContainingDeserializedClass(context: KtElement, memberDescriptor: DeserializedCallableMemberDescriptor): PsiClass? {
     val containingDeclaration = memberDescriptor.containingDeclaration
     return when (containingDeclaration) {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/psi/UastFakeLightMethod.kt

@@ -12,8 +12,7 @@ import com.intellij.psi.impl.light.LightParameterListBuilder
 import com.intellij.psi.impl.light.LightTypeParameterBuilder
 import org.jetbrains.kotlin.asJava.elements.KotlinLightTypeParameterListBuilder
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
-import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.uast.UastErrorType
@@ -21,18 +20,7 @@ import org.jetbrains.uast.kotlin.analyze
 import org.jetbrains.uast.kotlin.getType
 import org.jetbrains.uast.kotlin.toPsiType
 
-internal class UastFakeLightMethod(internal val original: KtFunction, containingClass: PsiClass) : LightMethodBuilder(
+internal class UastFakeLightMethod(original: KtFunction, containingClass: PsiClass) : UastFakeLightMethodBase<KtFunction>(original, containingClass) {
-    original.manager, original.language, original.name ?: "<no name provided>",
-    LightParameterListBuilder(original.manager, original.language),
-    LightModifierList(original.manager)
-) {
-
-    init {
-        this.containingClass = containingClass
-        if (original.safeAs<KtNamedFunction>()?.isTopLevel == true) {
-            addModifier(PsiModifier.STATIC)
-        }
-    }
 
     private val _buildTypeParameterList by lazy {
         KotlinLightTypeParameterListBuilder(this).also { paramList ->
@@ -86,6 +74,25 @@ internal class UastFakeLightMethod(internal val original: KtFunction, containing
     }
 
     override fun getParameterList(): PsiParameterList = paramsList
+}
+
+internal class UastFakeLightPrimaryConstructor(original: KtClassOrObject, lightClass: PsiClass) : UastFakeLightMethodBase<KtClassOrObject>(original, lightClass){
+    override fun isConstructor(): Boolean = true
+}
+
+internal abstract class UastFakeLightMethodBase<T: KtElement>(internal val original: T, containingClass: PsiClass) : LightMethodBuilder(
+    original.manager, original.language, original.name ?: "<no name provided>",
+    LightParameterListBuilder(original.manager, original.language),
+    LightModifierList(original.manager)
+) {
+
+    init {
+        this.containingClass = containingClass
+        if (original.safeAs<KtNamedFunction>()?.isTopLevel == true) {
+            addModifier(PsiModifier.STATIC)
+        }
+    }
+
 
     override fun getReturnType(): PsiType? {
         val context = original.analyze()
@@ -99,7 +106,7 @@ internal class UastFakeLightMethod(internal val original: KtFunction, containing
         if (this === other) return true
         if (javaClass != other?.javaClass) return false
 
-        other as UastFakeLightMethod
+        other as UastFakeLightMethodBase<*>
 
         if (original != other.original) return false
 

plugins/uast-kotlin/testData/LocalDeclarations.resolved.txt

@@ -1,8 +1,8 @@
 UTypeReferenceExpression (name = boolean) -> USimpleNameReferenceExpression (identifier = Boolean) -> PsiClass:Boolean: Boolean
-    UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))(resolves to null) -> USimpleNameReferenceExpression (identifier = Local) -> KtLightClassForLocalDeclaration:class Local: Local
+    UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))(resolves to LightMethodBuilder:Local) -> USimpleNameReferenceExpression (identifier = Local) -> KtLightClassForLocalDeclaration:class Local: Local
-        UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))(resolves to null) -> USimpleNameReferenceExpression (identifier = Local) -> KtLightClassForLocalDeclaration:class Local: Local
+        UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))(resolves to LightMethodBuilder:Local) -> USimpleNameReferenceExpression (identifier = Local) -> KtLightClassForLocalDeclaration:class Local: Local
     UTypeReferenceExpression (name = int) -> USimpleNameReferenceExpression (identifier = Int) -> PsiClass:Integer: Integer
     UTypeReferenceExpression (name = java.lang.String) -> USimpleNameReferenceExpression (identifier = String) -> PsiClass:String: String
-    UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))(resolves to null) -> USimpleNameReferenceExpression (identifier = bar) -> LightVariableBuilder:bar: bar
+    UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))(resolves to LightMethodBuilder:bar) -> USimpleNameReferenceExpression (identifier = bar) -> LightVariableBuilder:bar: bar
     UBinaryExpression (operator = ==) -> USimpleNameReferenceExpression (identifier = ==) -> null: null
-    UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))(resolves to null) -> USimpleNameReferenceExpression (identifier = Local) -> KtLightClassForLocalDeclaration:class Local: Local
+    UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))(resolves to LightMethodBuilder:Local) -> USimpleNameReferenceExpression (identifier = Local) -> KtLightClassForLocalDeclaration:class Local: Local

plugins/uast-kotlin/tests/AbstractKotlinResolveEverythingTest.kt.191

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2010-2019 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.test.kotlin
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.utils.addToStdlib.cast
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import org.jetbrains.uast.*
+import org.jetbrains.uast.test.common.kotlin.IndentedPrintingVisitor
+import org.jetbrains.uast.test.common.kotlin.visitUFileAndGetResult
+import org.jetbrains.uast.test.env.kotlin.assertEqualsToFile
+import java.io.File
+
+
+abstract class AbstractKotlinResolveEverythingTest : AbstractKotlinUastTest() {
+
+    private fun getTestFile(testName: String, ext: String) =
+        File(File(TEST_KOTLIN_MODEL_DIR, testName).canonicalPath + '.' + ext)
+
+
+    private fun UFile.resolvableWithTargets() = object : IndentedPrintingVisitor(KtBlockExpression::class) {
+        override fun render(element: PsiElement) =
+            sequenceOf(element.toUElementOfType<UReferenceExpression>(), element.toUElementOfType<UCallExpression>()).filterNotNull()
+                .filter {
+                    when (it) {
+                        is UCallExpression -> it.sourcePsi.safeAs<KtCallElement>()?.calleeExpression !is KtSimpleNameExpression
+                        else -> true
+                    }
+                }.takeIf { it.any() }
+                ?.joinTo(StringBuilder(), "\n") { ref ->
+                    StringBuilder().apply {
+                        val parent = ref.uastParent
+                        append(parent?.asLogString())
+                        if (parent is UCallExpression) {
+                            append("(resolves to ${parent.resolve()})")
+                        }
+                        append(" -> ")
+                        append(ref.asLogString())
+                        append(" -> ")
+                        append(ref.cast<UResolvable>().resolve())
+                        append(": ")
+                        append(
+                            when (ref) {
+                                is UReferenceExpression -> ref.resolvedName
+                                is UCallExpression -> ""
+                                else -> "<none>"
+                            }
+                        )
+            }
+        }
+    }.visitUFileAndGetResult(this)
+
+    override fun check(testName: String, file: UFile) {
+        assertEqualsToFile("resolved", getTestFile(testName, "resolved.txt"), file.resolvableWithTargets())
+    }
+}

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -568,11 +568,43 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
                 "bar",
                 localFunction.methodName
             )
-            assertNull(localFunction.resolve())
+            val localFunctionResolved = localFunction.resolve()
+            assertNotNull(localFunctionResolved)
             val receiver = localFunction.receiver ?: kfail("receiver expected")
             assertEquals("UReferenceExpression", receiver.asLogString())
-            val uParameter = (receiver as UReferenceExpression).resolve().toUElement() ?: kfail("uelement expected")
+            val uVariable = (receiver as UReferenceExpression).resolve().toUElement() ?: kfail("uelement expected")
-            assertEquals("ULambdaExpression", uParameter.asLogString())
+            assertEquals("ULambdaExpression", uVariable.asLogString())
+            assertEquals(uVariable, localFunctionResolved.toUElement())
+        }
+    }
+
+    @Test
+    fun testLocalConstructorCall() {
+        doTest("LocalDeclarations") { _, file ->
+            val localFunction = file.findElementByTextFromPsi<UElement>("bar() == Local()").
+                findElementByText<UCallExpression>("Local()")
+            assertEquals(
+                "UIdentifier (Identifier (Local))",
+                localFunction.methodIdentifier?.asLogString()
+            )
+            assertEquals(
+                "Local",
+                localFunction.methodIdentifier?.name
+            )
+            assertEquals(
+                "<init>",
+                localFunction.methodName
+            )
+            val localFunctionResolved = localFunction.resolve()
+            assertNotNull(localFunctionResolved)
+            val classReference = localFunction.classReference ?: kfail("classReference expected")
+            assertEquals("USimpleNameReferenceExpression (identifier = <init>, resolvesTo = Local)", classReference.asLogString())
+            val localClass = classReference.resolve().toUElement() ?: kfail("uelement expected")
+            assertEquals("UClass (name = Local)", localClass.asLogString())
+            assertEquals(localClass, localFunctionResolved.toUElement())
+            val localPrimaryConstructor = localFunctionResolved.toUElementOfType<UMethod>() ?: kfail("constructor expected")
+            assertTrue(localPrimaryConstructor.isConstructor)
+            assertEquals(localClass, localPrimaryConstructor.uastParent)
         }
     }
 

plugins/uast-kotlin/tests/KotlinUastResolveApiTest.kt

@@ -298,7 +298,7 @@ class KotlinUastResolveApiTest : KotlinLightCodeInsightFixtureTestCase() {
         )
 
         val uCallExpression = myFixture.file.findElementAt(myFixture.caretOffset).toUElement().getUCallExpression().orFail("cant convert to UCallExpression")
-        val resolved = uCallExpression.resolve().orFail("cant resolve from $uCallExpression") as PsiVariable
+        val resolved = uCallExpression.resolve().orFail("cant resolve from $uCallExpression")
         TestCase.assertEquals("bar", resolved.name)
     }