191: Uast: making KotlinUObjectLiteralExpression provide a proper referenceNameElement (KT-30534)

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

@@ -1,13 +1,12 @@
 package org.jetbrains.uast.test.kotlin
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.uast.UFile
-import org.jetbrains.uast.UReferenceExpression
+import org.jetbrains.uast.*
 import org.jetbrains.uast.test.common.UElementToParentMap
 import org.jetbrains.uast.test.common.kotlin.IdentifiersTestBase
 import org.jetbrains.uast.test.env.assertEqualsToFile
-import org.jetbrains.uast.toUElementOfType
 import java.io.File
+import kotlin.test.assertNotNull
 
 
 abstract class AbstractKotlinIdentifiersTest : AbstractKotlinUastTest(), IdentifiersTestBase {
@@ -23,6 +22,15 @@ abstract class AbstractKotlinIdentifiersTest : AbstractKotlinUastTest(), Identif
     }
 }
 
-fun UFile.asRefNames() = object : UElementToParentMap({ it.toUElementOfType<UReferenceExpression>()?.referenceNameElement }) {
+private fun refNameRetriever(psiElement: PsiElement): UElement? =
+    when (val uElement = psiElement.toUElementOfExpectedTypes(UCallExpression::class.java, UReferenceExpression::class.java)) {
+        is UReferenceExpression -> uElement.referenceNameElement
+        is UCallExpression -> uElement.classReference?.referenceNameElement
+        else -> null
+    }?.also {
+        assertNotNull(it.sourcePsi, "referenceNameElement should have physical source, origin = $psiElement")
+    }
+
+fun UFile.asRefNames() = object : UElementToParentMap(::refNameRetriever) {
     override fun renderSource(element: PsiElement): String = element.javaClass.simpleName
 }.visitUFileAndGetResult(this)