191: Uast: support convertToAlternatives for KtClassOrObject

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

@@ -181,11 +181,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
                 is KtEnumEntry -> el<UEnumConstant> {
                     convertEnumEntry(original, givenParent)
                 }
-                is KtClassOrObject -> el<UClass> {
+                is KtClassOrObject -> convertClassOrObject(original, givenParent, this).firstOrNull()
-                    original.toLightClass()?.let { lightClass ->
-                        KotlinUClass.create(lightClass, givenParent)
-                    }
-                }
                 is KtFunction ->
                     if (original.isLocal) {
                         el<ULambdaExpression> {
@@ -275,6 +271,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
     override fun <T : UElement> convertToAlternatives(element: PsiElement, requiredTypes: Array<out Class<out T>>): Sequence<T> = when {
         (element is KtProperty && !element.isLocal) -> convertNonLocalProperty(element, null, requiredTypes) as Sequence<T>
         element is KtParameter -> convertParameter(element, null, requiredTypes) as Sequence<T>
+        element is KtClassOrObject -> convertClassOrObject(element, null, requiredTypes) as Sequence<T>
         else -> sequenceOf(convertElementWithParent(element, requiredTypes.nonEmptyOr(DEFAULT_TYPES_LIST)) as? T).filterNotNull()
     }
 }
@@ -332,6 +329,25 @@ private fun convertParameter(
     *convertToPropertyAlternatives(LightClassUtil.getLightClassPropertyMethods(element), givenParent)
 )
 
+private fun convertClassOrObject(
+    element: KtClassOrObject,
+    givenParent: UElement?,
+    requiredTypes: Array<out Class<out UElement>>
+): Sequence<UElement> {
+    val ktLightClass = element.toLightClass() ?: return emptySequence()
+    val uClass = KotlinUClass.create(ktLightClass, givenParent)
+    return requiredTypes.accommodate(
+        alternative { uClass },
+        alternative primaryConstructor@{
+            val primaryConstructor = element.primaryConstructor ?: return@primaryConstructor null
+            uClass.methods.asSequence()
+                .filter { it.sourcePsi == primaryConstructor }
+                .firstOrNull()
+        }
+    )
+}
+
+
 internal object KotlinConverter {
     internal tailrec fun unwrapElements(element: PsiElement?): PsiElement? = when (element) {
         is KtValueArgumentList -> unwrapElements(element.parent)

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

@@ -6,6 +6,7 @@
 package org.jetbrains.uast.test.kotlin
 
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.test.testFramework.KtUsefulTestCase.assertInstanceOf
@@ -125,4 +126,34 @@ class KotlinUastAlternativesTest : AbstractKotlinUastTest() {
         }
     }
 
+    @Test
+    fun testPrimaryConstructorAlternatives() {
+        doTest("ManyAlternatives") { name, file ->
+            val index = file.psi.text.indexOf("ClassA")
+            val ktProperty = PsiTreeUtil.getParentOfType(file.psi.findElementAt(index), KtClass::class.java)!!
+            val plugin = UastLanguagePlugin.byLanguage(ktProperty.language)!!
+
+            plugin.convertToAlternatives<UElement>(ktProperty, arrayOf(UClass::class.java)).let {
+                assertEquals("public final class ClassA {", it.joinToString { it.asRenderString().lineSequence().first() })
+            }
+
+            plugin.convertToAlternatives<UElement>(ktProperty, arrayOf(UClass::class.java, UMethod::class.java)).let {
+                assertEquals(
+                    "public final class ClassA {, " +
+                            "public fun ClassA(@org.jetbrains.annotations.NotNull justParam: int, @org.jetbrains.annotations.NotNull paramAndProp: java.lang.String) = UastEmptyExpression",
+                    it.joinToString { it.asRenderString().lineSequence().first() }
+                )
+            }
+
+            plugin.convertToAlternatives<UElement>(ktProperty, arrayOf(UElement::class.java)).let {
+                assertEquals(
+                    "public final class ClassA {, " +
+                            "public fun ClassA(@org.jetbrains.annotations.NotNull justParam: int, @org.jetbrains.annotations.NotNull paramAndProp: java.lang.String) = UastEmptyExpression",
+                    it.joinToString { it.asRenderString().lineSequence().first() }
+                )
+            }
+
+        }
+    }
+
 }