Uast: making KotlinClassViaConstructorUSimpleReferenceExpression resolve to class in case of SAM (KT-28272)

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

@@ -16,6 +16,7 @@
 
 package org.jetbrains.uast.kotlin
 
+import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiNamedElement
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.load.java.sam.SamConstructorDescriptor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getAssignmentByLHS
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
@@ -203,12 +205,14 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
         get() = (psi.getResolvedCall(psi.analyze())?.resultingDescriptor as? ConstructorDescriptor)
                 ?.containingDeclaration?.name?.asString()
 
-    override fun resolve(): PsiElement? {
+    override fun resolve(): PsiElement? =
-        val resolvedCall = psi.getResolvedCall(psi.analyze())
+        when (val resultingDescriptor = psi.getResolvedCall(psi.analyze())?.resultingDescriptor) {
-        val resultingDescriptor = resolvedCall?.resultingDescriptor as? ConstructorDescriptor ?: return null
+            is ConstructorDescriptor ->
-        val clazz = resultingDescriptor.containingDeclaration
+                resultingDescriptor.containingDeclaration.toSource()?.getMaybeLightElement(this)
-        return clazz.toSource()?.getMaybeLightElement(this)
+            is SamConstructorDescriptor ->
-    }
+                (resultingDescriptor.returnType?.getFunctionalInterfaceType(this, psi) as? PsiClassType)?.resolve()
+            else -> null
+        }
 }
 
 class KotlinStringUSimpleReferenceExpression(

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

@@ -16,6 +16,7 @@
 
 package org.jetbrains.uast.kotlin
 
+import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiNamedElement
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.load.java.sam.SamConstructorDescriptor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getAssignmentByLHS
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
@@ -206,12 +208,14 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
         get() = (psi.getResolvedCall(psi.analyze())?.resultingDescriptor as? ConstructorDescriptor)
                 ?.containingDeclaration?.name?.asString()
 
-    override fun resolve(): PsiElement? {
+    override fun resolve(): PsiElement? =
-        val resolvedCall = psi.getResolvedCall(psi.analyze())
+        when (val resultingDescriptor = psi.getResolvedCall(psi.analyze())?.resultingDescriptor) {
-        val resultingDescriptor = resolvedCall?.resultingDescriptor as? ConstructorDescriptor ?: return null
+            is ConstructorDescriptor ->
-        val clazz = resultingDescriptor.containingDeclaration
+                resultingDescriptor.containingDeclaration.toSource()?.getMaybeLightElement(this)
-        return clazz.toSource()?.getMaybeLightElement(this)
+            is SamConstructorDescriptor ->
-    }
+                (resultingDescriptor.returnType?.getFunctionalInterfaceType(this, psi) as? PsiClassType)?.resolve()
+            else -> null
+        }
 }
 
 class KotlinStringUSimpleReferenceExpression(

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -1,6 +1,7 @@
 package org.jetbrains.uast.test.kotlin
 
 import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiModifier
 import com.intellij.testFramework.RunAll
 import com.intellij.testFramework.UsefulTestCase
@@ -127,6 +128,12 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
                         "java.lang.Runnable",
                         file.findElementByText<ULambdaExpression>("{ println(\"hello2\") }").functionalInterfaceType?.canonicalText
                     )
+                }, {
+                    val call = file.findElementByText<UCallExpression>("Runnable { println(\"hello2\") }")
+                    assertEquals(
+                        "java.lang.Runnable",
+                        (call.classReference?.resolve() as? PsiClass)?.qualifiedName
+                    )
                 }, {
                     assertEquals(
                         "java.util.function.Supplier<T>",

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

@@ -128,6 +128,12 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
                         "java.lang.Runnable",
                         file.findElementByText<ULambdaExpression>("{ println(\"hello2\") }").functionalInterfaceType?.canonicalText
                     )
+                }, {
+                    val call = file.findElementByText<UCallExpression>("Runnable { println(\"hello2\") }")
+                    assertEquals(
+                        "java.lang.Runnable",
+                        (call.classReference?.resolve() as? PsiClass)?.qualifiedName
+                    )
                 }, {
                     assertEquals(
                         "java.util.function.Supplier<T>",