Uast: expression type for assigning arrays (KT-34187)

JetBrains · Apr 14, 2020 · 738e2f8 · 738e2f8
1 parent 9674643
commit 738e2f8
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/...ins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUArrayAccessExpression.kt b/...ins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUArrayAccessExpression.kt
@@ -16,7 +16,9 @@
 
 package org.jetbrains.uast.kotlin
 
+import com.intellij.psi.PsiType
 import org.jetbrains.kotlin.psi.KtArrayAccessExpression
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.uast.UArrayAccessExpression
 import org.jetbrains.uast.UElement
 
@@ -26,4 +28,14 @@ class KotlinUArrayAccessExpression(
 ) : KotlinAbstractUExpression(givenParent), UArrayAccessExpression, KotlinUElementWithType, KotlinEvaluatableUElement {
     override val receiver by lz { KotlinConverter.convertOrEmpty(sourcePsi.arrayExpression, this) }
     override val indices by lz { sourcePsi.indexExpressions.map { KotlinConverter.convertOrEmpty(it, this) } }
+
+    override fun getExpressionType(): PsiType? {
+        super<KotlinUElementWithType>.getExpressionType()?.let { return it }
+
+        // for unknown reason in assigment position there is no `EXPRESSION_TYPE_INFO` so we getting it from the array type
+        val arrayExpression = sourcePsi.arrayExpression ?: return null
+        val arrayType = arrayExpression.analyze()[BindingContext.EXPRESSION_TYPE_INFO, arrayExpression]?.type ?: return null
+        return arrayType.arguments.firstOrNull()?.type?.toPsiType(this, arrayExpression, false )
+    }
+
 }
diff --git a/plugins/uast-kotlin/tests/KotlinUastResolveApiTest.kt b/plugins/uast-kotlin/tests/KotlinUastResolveApiTest.kt
@@ -315,5 +315,30 @@ class KotlinUastResolveApiTest : KotlinLightCodeInsightFixtureTestCase() {
         val resolved = compiledAnnotationParameter.resolve() as PsiMethod
         TestCase.assertEquals("message", resolved.name)
     }
+
+    fun testAssigningArrayElementType() {
+        myFixture.configureByText(
+            "MyClass.kt", """ 
+            fun foo() {
+                val arr = arrayOfNulls<List<*>>(10)
+                arr[0] = emptyList<Any>()
+                
+                val lst = mutableListOf<List<*>>()
+                lst[0] = emptyList<Any>()
+            }
+        """
+        )
+
+        val uFile = myFixture.file.toUElement()!!
+
+        TestCase.assertEquals(
+            "PsiType:List<?>",
+            uFile.findElementByTextFromPsi<UExpression>("arr[0]").getExpressionType().toString()
+        )
+        TestCase.assertEquals(
+            "PsiType:List<?>",
+            uFile.findElementByTextFromPsi<UExpression>("lst[0]").getExpressionType().toString()
+        )
+    }
 }