Uast: support for Kotlin array literals

2018-01-31 18:48:15 +03:00
parent 358bc62fc1
commit 8ac95b54a2
7 changed files with 127 additions and 3 deletions
@@ -323,8 +323,20 @@ internal object KotlinConverter {
                }
            is KtImportDirective -> el<UImportStatement>(build(::KotlinUImportStatement))
            else -> {
-                if (element is LeafPsiElement && element.elementType == KtTokens.IDENTIFIER) {
+                if (element is LeafPsiElement) {
+                    if (element.elementType == KtTokens.IDENTIFIER)
                    el<UIdentifier>(build(::UIdentifier))
+                    else if (element.elementType == KtTokens.LBRACKET && element.parent is KtCollectionLiteralExpression)
+                        el<UIdentifier> {
+                            UIdentifier(
+                                element,
+                                KotlinUCollectionLiteralExpression(
+                                    element.parent as KtCollectionLiteralExpression,
+                                    null
+                                )
+                            )
+                        }
+                    else null
                } else {
                    null
                }
@@ -395,6 +407,7 @@ internal object KotlinConverter {
            is KtSafeQualifiedExpression -> expr<UQualifiedReferenceExpression>(build(::KotlinUSafeQualifiedExpression))
            is KtSimpleNameExpression -> expr<USimpleNameReferenceExpression>(build(::KotlinUSimpleReferenceExpression))
            is KtCallExpression -> expr<UCallExpression>(build(::KotlinUFunctionCallExpression))
+            is KtCollectionLiteralExpression -> expr<UCallExpression>(build(::KotlinUCollectionLiteralExpression))
            is KtBinaryExpression -> {
                if (expression.operationToken == KtTokens.ELVIS) {
                    expr<UExpressionList>(build(::createElvisExpression))
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. 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.expressions
+
+import com.intellij.psi.PsiArrayType
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiType
+import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
+import org.jetbrains.uast.*
+import org.jetbrains.uast.kotlin.KotlinAbstractUExpression
+import org.jetbrains.uast.kotlin.KotlinConverter
+import org.jetbrains.uast.kotlin.KotlinUElementWithType
+
+class KotlinUCollectionLiteralExpression(
+    override val sourcePsi: KtCollectionLiteralExpression,
+    givenParent: UElement?
+) : KotlinAbstractUExpression(givenParent), UCallExpression, KotlinUElementWithType {
+
+    override val classReference: UReferenceExpression? get() = null
+
+    override val kind: UastCallKind = UastCallKind.NESTED_ARRAY_INITIALIZER
+
+    override val methodIdentifier: UIdentifier? by lazy { UIdentifier(sourcePsi.leftBracket, this) }
+
+    override val methodName: String? get() = null
+
+    override val receiver: UExpression? get() = null
+
+    override val receiverType: PsiType? get() = null
+
+    override val returnType: PsiType? get() = getExpressionType()
+
+    override val typeArgumentCount: Int get() = typeArguments.size
+
+    override val typeArguments: List<PsiType> get() = listOfNotNull((returnType as? PsiArrayType)?.componentType)
+
+    override val valueArgumentCount: Int
+        get() = sourcePsi.getInnerExpressions().size
+
+    override val valueArguments by lazy {
+        sourcePsi.getInnerExpressions().map { KotlinConverter.convertOrEmpty(it, this) }
+    }
+
+    override fun asRenderString(): String = "collectionLiteral[" + valueArguments.joinToString { it.asRenderString() } + "]"
+
+    override fun resolve(): PsiMethod? = null
+
+    override val psi: PsiElement get() = sourcePsi
+
+}
@@ -3,6 +3,8 @@ annotation class IntRange(val from: Long, val to: Long)

 annotation class RequiresPermission(val anyOf: IntArray)

+annotation class RequiresStrPermission(val strs: Array<String>)
+
 annotation class WithDefaultValue(val value: Int = 42)

 annotation class SuppressLint(vararg val value: String)
@@ -15,4 +17,10 @@ fun foo(): Int = 5

@IntRange(0, 100)
@SuppressLint("Lorem", "Ipsum", "Dolor")
-fun bar() = Unit
+fun bar() = Unit
+
+@RequiresPermission(anyOf = [1, 2, 3])
+fun fooWithArrLiteral(): Int = 5
+
+@RequiresStrPermission(strs = ["a", "b", "c"])
+fun fooWithStrArrLiteral(): Int = 3
@@ -32,11 +32,31 @@ UFile (package = )
                        ULiteralExpression (value = "Ipsum")
                        ULiteralExpression (value = "Dolor")
            USimpleNameReferenceExpression (identifier = Unit)
+        UAnnotationMethod (name = fooWithArrLiteral)
+            UAnnotation (fqName = RequiresPermission)
+                UNamedExpression (name = anyOf)
+                    UCallExpression (kind = UastCallKind(name='array_initializer'), argCount = 3))
+                        UIdentifier (Identifier ([))
+                        ULiteralExpression (value = 1)
+                        ULiteralExpression (value = 2)
+                        ULiteralExpression (value = 3)
+            ULiteralExpression (value = 5)
+        UAnnotationMethod (name = fooWithStrArrLiteral)
+            UAnnotation (fqName = RequiresStrPermission)
+                UNamedExpression (name = strs)
+                    UCallExpression (kind = UastCallKind(name='array_initializer'), argCount = 3))
+                        UIdentifier (Identifier ([))
+                        ULiteralExpression (value = "a")
+                        ULiteralExpression (value = "b")
+                        ULiteralExpression (value = "c")
+            ULiteralExpression (value = 3)
    UClass (name = IntRange)
        UAnnotationMethod (name = from)
        UAnnotationMethod (name = to)
    UClass (name = RequiresPermission)
        UAnnotationMethod (name = anyOf)
+    UClass (name = RequiresStrPermission)
+        UAnnotationMethod (name = strs)
    UClass (name = WithDefaultValue)
        UAnnotationMethod (name = value)
            ULiteralExpression (value = 42)
@@ -7,6 +7,10 @@ public final class AnnotationParametersKt {
    @IntRange(from = 0, to = 100)
    @SuppressLint(value = <noref>("Lorem", "Ipsum", "Dolor"))
    public static final fun bar() : void = Unit
+    @RequiresPermission(anyOf = collectionLiteral[1, 2, 3])
+    public static final fun fooWithArrLiteral() : int = 5
+    @RequiresStrPermission(strs = collectionLiteral["a", "b", "c"])
+    public static final fun fooWithStrArrLiteral() : int = 3
 }

 public abstract annotation IntRange {
@@ -18,6 +22,10 @@ public abstract annotation RequiresPermission {
    public abstract fun anyOf() : int[] = UastEmptyExpression
 }

+public abstract annotation RequiresStrPermission {
+    public abstract fun strs() : java.lang.String[] = UastEmptyExpression
+}
+
 public abstract annotation WithDefaultValue {
    public abstract fun value() : int = UastEmptyExpression
 }
@@ -3,6 +3,7 @@ package org.jetbrains.uast.test.kotlin
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiRecursiveElementVisitor
+import com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.addToStdlib.assertedCast
@@ -108,7 +109,9 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog

                jvmDeclaration.sourcePsi?.let {
                    assertTrue("sourcePsi should be physical but ${it.javaClass} found for [${it.text}] " +
-                               "for ${jvmDeclaration.javaClass}->${jvmDeclaration.uastParent?.javaClass}", it is KtElement)
+                                       "for ${jvmDeclaration.javaClass}->${jvmDeclaration.uastParent?.javaClass}",
+                               it is KtElement || it is LeafPsiElement
+                    )
                }
                jvmDeclaration.javaPsi?.let {
                    assertTrue("javaPsi should be light but ${it.javaClass} found for [${it.text}] " +
@@ -257,6 +257,24 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
            )
        }
    }
+
+    @Test
+    fun testLiteralArraysTypes() {
+        doTest("AnnotationParameters") { _, file ->
+            file.findElementByTextFromPsi<UCallExpression>("intArrayOf(1, 2, 3)").let { field ->
+                Assert.assertEquals("PsiType:int[]", field.returnType.toString())
+            }
+            file.findElementByTextFromPsi<UCallExpression>("[1, 2, 3]").let { field ->
+                Assert.assertEquals("PsiType:int[]", field.returnType.toString())
+                Assert.assertEquals("PsiType:int", field.typeArguments.single().toString())
+            }
+            file.findElementByTextFromPsi<UCallExpression>("[\"a\", \"b\", \"c\"]").let { field ->
+                Assert.assertEquals("PsiType:String[]", field.returnType.toString())
+                Assert.assertEquals("PsiType:String", field.typeArguments.single().toString())
+            }
+
+        }
+    }
 }

 fun <T, R> Iterable<T>.assertedFind(value: R, transform: (T) -> R): T = find { transform(it) == value } ?: throw AssertionError("'$value' not found, only ${this.joinToString { transform(it).toString() }}")