Implement getTypeElement for UastKotlinPsiVariable

#KT-16849 Fixed

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/psi/UastKotlinPsiVariable.kt

@@ -1,17 +1,20 @@
 package org.jetbrains.uast.kotlin.psi
 
-import com.intellij.lang.Language
 import com.intellij.psi.*
+import com.intellij.psi.impl.light.LightTypeElement
 import org.jetbrains.kotlin.asJava.elements.LightVariableBuilder
+import org.jetbrains.kotlin.builtins.createFunctionType
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.uast.UDeclaration
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UastErrorType
 import org.jetbrains.uast.getParentOfType
 import org.jetbrains.uast.kotlin.analyze
+import org.jetbrains.uast.kotlin.lz
 import org.jetbrains.uast.kotlin.orAnonymous
 import org.jetbrains.uast.kotlin.toPsiType
 
@@ -19,16 +22,31 @@ class UastKotlinPsiVariable(
         manager: PsiManager,
         name: String,
         type: PsiType,
-        language: Language,
         val ktInitializer: KtExpression?,
         val psiParent: PsiElement?,
         val containingElement: UElement,
         val ktElement: KtElement
-) : LightVariableBuilder(manager, name, type, language), PsiLocalVariable {
+) : LightVariableBuilder(manager, name, type, KotlinLanguage.INSTANCE), PsiLocalVariable {
+
+    private val psiTypeElement: PsiTypeElement by lz {
+        LightTypeElement(manager, type)
+    }
+
+    private val psiInitializer: PsiExpression? by lz {
+        ktInitializer?.let { KotlinUastPsiExpression(it, containingElement) }
+    }
+
     override fun getParent() = psiParent
 
     override fun hasInitializer() = ktInitializer != null
-    override fun getInitializer(): PsiExpression? = ktInitializer?.let { KotlinUastPsiExpression(it, containingElement) }
+
+    override fun getInitializer(): PsiExpression? = psiInitializer
+
+    override fun getTypeElement() = psiTypeElement
+
+    override fun setInitializer(initializer: PsiExpression?) = throw NotImplementedError()
+
+    override fun getContainingFile(): PsiFile? = ktElement.containingFile
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
@@ -36,11 +54,6 @@ class UastKotlinPsiVariable(
         return ktElement == (other as? UastKotlinPsiVariable)?.ktElement
     }
 
-    override fun getTypeElement() = throw NotImplementedError()
-    override fun setInitializer(initializer: PsiExpression?) = throw NotImplementedError()
-
-    override fun getContainingFile(): PsiFile? = ktElement.containingFile
-
     override fun hashCode() = ktElement.hashCode()
 
     companion object {
@@ -55,7 +68,6 @@ class UastKotlinPsiVariable(
                     declaration.manager,
                     declaration.name.orAnonymous("unnamed"),
                     declaration.typeReference.toPsiType(containingElement),
-                    KotlinLanguage.INSTANCE,
                     initializer ?: declaration.initializer,
                     psiParent,
                     containingElement,
@@ -67,8 +79,7 @@ class UastKotlinPsiVariable(
             return UastKotlinPsiVariable(
                     declaration.manager,
                     "var" + Integer.toHexString(declaration.getHashCode()),
-                    UastErrorType, //TODO,
+                    declaration.initializer?.getType(containingElement) ?: UastErrorType,
-                    KotlinLanguage.INSTANCE,
                     declaration.initializer,
                     psiParent,
                     containingElement,
@@ -80,8 +91,7 @@ class UastKotlinPsiVariable(
             return UastKotlinPsiVariable(
                     initializer.manager,
                     "var" + Integer.toHexString(initializer.getHashCode()),
-                    UastErrorType, //TODO,
+                    initializer.getType(containingElement) ?: UastErrorType,
-                    KotlinLanguage.INSTANCE,
                     initializer,
                     psiParent,
                     containingElement,
@@ -93,8 +103,7 @@ class UastKotlinPsiVariable(
             return UastKotlinPsiVariable(
                     localFunction.manager,
                     name,
-                    UastErrorType, // TODO,
+                    localFunction.getFunctionType(containingElement) ?: UastErrorType,
-                    KotlinLanguage.INSTANCE,
                     localFunction,
                     psiParent,
                     containingElement,
@@ -104,12 +113,26 @@ class UastKotlinPsiVariable(
 }
 
 private class KotlinUastPsiExpression(val ktExpression: KtExpression, val parent: UElement) : PsiElement by ktExpression, PsiExpression {
-    override fun getType(): PsiType? {
+    override fun getType(): PsiType? = ktExpression.getType(parent)
-        val ktType = ktExpression.analyze()[BindingContext.EXPRESSION_TYPE_INFO, ktExpression]?.type ?: return null
-        return ktType.toPsiType(parent, ktExpression, boxed = false)
-    }
 }
 
+private fun KtFunction.getFunctionType(parent: UElement): PsiType? {
+    val descriptor = analyze()[BindingContext.FUNCTION, this] ?: return null
+    val returnType = descriptor.returnType ?: return null
+
+    return createFunctionType(
+            descriptor.builtIns,
+            descriptor.annotations,
+            descriptor.extensionReceiverParameter?.type,
+            descriptor.valueParameters.map { it.type },
+            descriptor.valueParameters.map { it.name },
+            returnType
+    ).toPsiType(parent, this, boxed = false)
+}
+
+private fun KtExpression.getType(parent: UElement): PsiType? =
+        analyze()[BindingContext.EXPRESSION_TYPE_INFO, this]?.type?.toPsiType(parent, this, boxed = false)
+
 private fun PsiElement.getHashCode(): Int {
     var result = 42
     result = 41 * result + containingFile.name.hashCode()

plugins/uast-kotlin/testData/Elvis.kt

@@ -1,6 +1,8 @@
 
-fun foo(bar: String): Any? = null
+fun foo(bar: String): String? = null
 
-fun baz(): Any? {
+fun bar() = 42
-    return foo("Lorem ipsum") ?: foo("dolor sit amet") ?: foo("consectetuer adipiscing elit")
+
+fun baz(): String? {
+    return foo("Lorem ipsum") ?: foo("dolor sit amet") ?: bar().toString()
 }

plugins/uast-kotlin/testData/Elvis.log.txt

@@ -4,34 +4,39 @@ UFile (package = )
             UParameter (name = bar)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             ULiteralExpression (value = null)
+        UAnnotationMethod (name = bar)
+            ULiteralExpression (value = 42)
         UAnnotationMethod (name = baz)
             UBlockExpression
                 UReturnExpression
                     UExpressionList (elvis)
                         UDeclarationsExpression
-                            ULocalVariable (name = var243c4e1a)
+                            ULocalVariable (name = var243c51a0)
                                 UExpressionList (elvis)
                                     UDeclarationsExpression
-                                        ULocalVariable (name = varc4aef1e3)
+                                        ULocalVariable (name = varc4aef569)
                                             UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
                                                 UIdentifier (Identifier (foo))
                                                 USimpleNameReferenceExpression (identifier = foo)
                                                 ULiteralExpression (value = "Lorem ipsum")
                                     UIfExpression
                                         UBinaryExpression (operator = !=)
-                                            USimpleNameReferenceExpression (identifier = varc4aef1e3)
+                                            USimpleNameReferenceExpression (identifier = varc4aef569)
                                             ULiteralExpression (value = null)
-                                        USimpleNameReferenceExpression (identifier = varc4aef1e3)
+                                        USimpleNameReferenceExpression (identifier = varc4aef569)
                                         UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
                                             UIdentifier (Identifier (foo))
                                             USimpleNameReferenceExpression (identifier = foo)
                                             ULiteralExpression (value = "dolor sit amet")
                         UIfExpression
                             UBinaryExpression (operator = !=)
-                                USimpleNameReferenceExpression (identifier = var243c4e1a)
+                                USimpleNameReferenceExpression (identifier = var243c51a0)
                                 ULiteralExpression (value = null)
-                            USimpleNameReferenceExpression (identifier = var243c4e1a)
+                            USimpleNameReferenceExpression (identifier = var243c51a0)
-                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
+                            UQualifiedReferenceExpression
-                                UIdentifier (Identifier (foo))
+                                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
-                                USimpleNameReferenceExpression (identifier = foo)
+                                    UIdentifier (Identifier (bar))
-                                ULiteralExpression (value = "consectetuer adipiscing elit")
+                                    USimpleNameReferenceExpression (identifier = bar)
+                                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                                    UIdentifier (Identifier (toString))
+                                    USimpleNameReferenceExpression (identifier = toString)

plugins/uast-kotlin/testData/Elvis.render.txt

@@ -1,12 +1,13 @@
 public final class ElvisKt {
-    public static final fun foo(bar: java.lang.String) : java.lang.Object = null
+    public static final fun foo(bar: java.lang.String) : java.lang.String = null
-    public static final fun baz() : java.lang.Object {
+    public static final fun bar() : int = 42
+    public static final fun baz() : java.lang.String {
         return elvis {
-            var var243c4e1a: <ErrorType> = elvis {
+            var var243c51a0: java.lang.String = elvis {
-                var varc4aef1e3: <ErrorType> = foo("Lorem ipsum")
+                var varc4aef569: java.lang.String = foo("Lorem ipsum")
-                if (varc4aef1e3 != null) varc4aef1e3 else foo("dolor sit amet")
+                if (varc4aef569 != null) varc4aef569 else foo("dolor sit amet")
             }
-            if (var243c4e1a != null) var243c4e1a else foo("consectetuer adipiscing elit")
+            if (var243c51a0 != null) var243c51a0 else bar().toString()
         }
     }
 }

plugins/uast-kotlin/testData/LocalDeclarations.kt

@@ -6,5 +6,7 @@ fun foo(): Boolean {
         Local()
     }
 
+    fun Int.someLocalFun(text: String) = 42
+
     return bar() == Local()
 }

plugins/uast-kotlin/testData/LocalDeclarations.log.txt

@@ -18,6 +18,11 @@ UFile (package = )
                                 UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))
                                     UIdentifier (Identifier (Local))
                                     USimpleNameReferenceExpression (identifier = <init>)
+                UDeclarationsExpression
+                    UVariable (name = someLocalFun)
+                        ULambdaExpression
+                            UParameter (name = text)
+                            ULiteralExpression (value = 42)
                 UReturnExpression
                     UBinaryExpression (operator = ==)
                         UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))

plugins/uast-kotlin/testData/LocalDeclarations.render.txt

@@ -9,6 +9,9 @@ public final class LocalDeclarationsKt {
         var baz: <ErrorType> = fun () {
                 <init>()
         }
+        var someLocalFun: kotlin.jvm.functions.Function2<? super java.lang.Integer,? super java.lang.String,? extends java.lang.Integer> = fun (var text: java.lang.String) {
+            42
+        }
         return bar() == <init>()
     }
 }

plugins/uast-kotlin/testData/LocalDeclarations.types.txt

@@ -18,6 +18,11 @@ UFile (package = ) [public final class LocalDeclarationsKt {...]
                                 UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0)) [<init>()] : PsiType:<ErrorType>
                                     UIdentifier (Identifier (Local)) [UIdentifier (Identifier (Local))]
                                     USimpleNameReferenceExpression (identifier = <init>) [<init>] : PsiType:<ErrorType>
+                UDeclarationsExpression [var someLocalFun: kotlin.jvm.functions.Function2<? super java.lang.Integer,? super java.lang.String,? extends java.lang.Integer> = fun (var text: java.lang.String) {...}]
+                    UVariable (name = someLocalFun) [var someLocalFun: kotlin.jvm.functions.Function2<? super java.lang.Integer,? super java.lang.String,? extends java.lang.Integer> = fun (var text: java.lang.String) {...}]
+                        ULambdaExpression [fun (var text: java.lang.String) {...}]
+                            UParameter (name = text) [var text: java.lang.String]
+                            ULiteralExpression (value = 42) [42] : PsiType:int
                 UReturnExpression [return bar() == <init>()] : PsiType:Void
                     UBinaryExpression (operator = ==) [bar() == <init>()] : PsiType:boolean
                         UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0)) [bar()] : PsiType:<ErrorType>

plugins/uast-kotlin/testData/LocalDeclarations.values.txt

@@ -18,6 +18,11 @@ UFile (package = ) [public final class LocalDeclarationsKt {...]
                                 UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0)) [<init>()] = external <init>()()
                                     UIdentifier (Identifier (Local)) [UIdentifier (Identifier (Local))]
                                     USimpleNameReferenceExpression (identifier = <init>) [<init>] = external <init>()
+                UDeclarationsExpression [var someLocalFun: kotlin.jvm.functions.Function2<? super java.lang.Integer,? super java.lang.String,? extends java.lang.Integer> = fun (var text: java.lang.String) {...}] = Undetermined
+                    UVariable (name = someLocalFun) [var someLocalFun: kotlin.jvm.functions.Function2<? super java.lang.Integer,? super java.lang.String,? extends java.lang.Integer> = fun (var text: java.lang.String) {...}]
+                        ULambdaExpression [fun (var text: java.lang.String) {...}] = Undetermined
+                            UParameter (name = text) [var text: java.lang.String]
+                            ULiteralExpression (value = 42) [42] = 42
                 UReturnExpression [return bar() == <init>()] = Nothing
                     UBinaryExpression (operator = ==) [bar() == <init>()] = Undetermined
                         UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0)) [bar()] = external bar()()