Uast: fix for UVariables missing typeReferences

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt

@@ -92,9 +92,11 @@ abstract class AbstractKotlinUVariable(givenParent: UElement?) : KotlinAbstractU
     }
 
 
-    abstract protected fun acceptsAnnotationTarget(target: AnnotationUseSiteTarget?): Boolean
+    protected abstract fun acceptsAnnotationTarget(target: AnnotationUseSiteTarget?): Boolean
 
-    override val typeReference by lz { getLanguagePlugin().convertOpt<UTypeReferenceExpression>(psi.typeElement, this) }
+    override val typeReference: UTypeReferenceExpression? by lz {
+        KotlinUTypeReferenceExpression(type, (sourcePsi as? KtCallableDeclaration)?.typeReference, this)
+    }
 
     override val uastAnchor: UIdentifier?
         get() {
@@ -158,8 +160,6 @@ class KotlinUVariable(
 
     override val psi = javaPsi
 
-    override val typeReference by lz { getLanguagePlugin().convertOpt<UTypeReferenceExpression>(psi.typeElement, this) }
-
     override fun acceptsAnnotationTarget(target: AnnotationUseSiteTarget?): Boolean = true
 
     override fun getInitializer(): PsiExpression? {

plugins/uast-kotlin/testData/TypeReferences.kt

@@ -0,0 +1,4 @@
+fun foo(parameter: Int) {
+    val varWithType: String? = "Not Null"
+    val varWithoutType = "lorem ipsum"
+}

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -495,6 +495,36 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
         }
     }
 
+    @Test
+    fun testVariablesTypeReferences() {
+        doTest("TypeReferences") { _, file ->
+            run {
+                val localVariable = file.findElementByTextFromPsi<UVariable>("val varWithType: String? = \"Not Null\"")
+                val typeReference = localVariable.typeReference
+                assertEquals("java.lang.String", typeReference?.getQualifiedName())
+                val sourcePsi = typeReference?.sourcePsi ?: kotlin.test.fail("no sourcePsi")
+                assertTrue("sourcePsi = $sourcePsi should be physical", sourcePsi.isPhysical)
+                assertEquals("String?", sourcePsi.text)
+            }
+
+            run {
+                val localVariable = file.findElementByTextFromPsi<UVariable>("val varWithoutType = \"lorem ipsum\"")
+                val typeReference = localVariable.typeReference
+                assertEquals("java.lang.String", typeReference?.getQualifiedName())
+                assertNull(typeReference?.sourcePsi)
+            }
+
+            run {
+                val localVariable = file.findElementByTextFromPsi<UVariable>("parameter: Int")
+                val typeReference = localVariable.typeReference
+                assertEquals("int", typeReference?.type?.presentableText)
+                val sourcePsi = typeReference?.sourcePsi ?: kotlin.test.fail("no sourcePsi")
+                assertTrue("sourcePsi = $sourcePsi should be physical", sourcePsi.isPhysical)
+                assertEquals("Int", sourcePsi.text)
+            }
+        }
+    }
+
 }
 
 fun <T, R> Iterable<T>.assertedFind(value: R, transform: (T) -> R): T =

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

@@ -18,6 +18,7 @@ import org.jetbrains.uast.test.env.kotlin.findElementByTextFromPsi
 import org.jetbrains.uast.visitor.AbstractUastVisitor
 import org.junit.Assert
 import org.junit.Test
+import kotlin.test.fail as kfail
 
 
 class KotlinUastApiTest : AbstractKotlinUastTest() {
@@ -510,6 +511,35 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
         }
     }
 
+    @Test
+    fun testVariablesTypeReferences() {
+        doTest("TypeReferences") { _, file ->
+            run {
+                val localVariable = file.findElementByTextFromPsi<UVariable>("val varWithType: String? = \"Not Null\"")
+                val typeReference = localVariable.typeReference
+                assertEquals("java.lang.String", typeReference?.getQualifiedName())
+                val sourcePsi = typeReference?.sourcePsi ?: kotlin.test.fail("no sourcePsi")
+                assertTrue("sourcePsi = $sourcePsi should be physical", sourcePsi.isPhysical)
+                assertEquals("String?", sourcePsi.text)
+            }
+
+            run {
+                val localVariable = file.findElementByTextFromPsi<UVariable>("val varWithoutType = \"lorem ipsum\"")
+                val typeReference = localVariable.typeReference
+                assertEquals("java.lang.String", typeReference?.getQualifiedName())
+                assertNull(typeReference?.sourcePsi)
+            }
+
+            run {
+                val localVariable = file.findElementByTextFromPsi<UVariable>("parameter: Int")
+                val typeReference = localVariable.typeReference
+                assertEquals("int", typeReference?.type?.presentableText)
+                val sourcePsi = typeReference?.sourcePsi ?: kotlin.test.fail("no sourcePsi")
+                assertTrue("sourcePsi = $sourcePsi should be physical", sourcePsi.isPhysical)
+                assertEquals("Int", sourcePsi.text)
+            }
+        }
+    }
 
 }