Uast: support for annotated expressions (KT-35801)

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt

@@ -206,11 +206,12 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
         return result.tempVarAssignment
     }
 
-    if (result is KotlinUElvisExpression && parent is KtBinaryExpression) {
+    if (result is KotlinUElvisExpression && parentUnwrapped is KtBinaryExpression) {
-        when (element.psi) {
+        val branch: Sequence<PsiElement?> = element.psi?.parentsWithSelf.orEmpty().takeWhile { it != parentUnwrapped }
-            parent.left -> return result.lhsDeclaration
+        if (branch.contains(parentUnwrapped.left))
-            parent.right -> return result.rhsIfExpression
+            return result.lhsDeclaration
-        }
+        if (branch.contains(parentUnwrapped.right))
+            return result.rhsIfExpression
     }
 
     if ((result is UMethod || result is KotlinLocalFunctionULambdaExpression)

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

@@ -198,6 +198,7 @@ internal object KotlinConverter {
         is KtTypeElement -> unwrapElements(element.parent)
         is KtSuperTypeList -> unwrapElements(element.parent)
         is KtFinallySection -> unwrapElements(element.parent)
+        is KtAnnotatedExpression -> unwrapElements(element.parent)
         else -> element
     }
 
@@ -422,6 +423,11 @@ internal object KotlinConverter {
             else {
                 expr<UDeclarationsExpression>(build(::createLocalFunctionDeclaration))
             }
+            is KtAnnotatedExpression -> {
+                expression.baseExpression
+                    ?.let { convertExpression(it, givenParent, requiredType) }
+                    ?: expr<UExpression>(build(::UnknownKotlinExpression))
+            }
 
             else -> expr<UExpression>(build(::UnknownKotlinExpression))
         }}

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.declarations.KotlinUIdentifier
 import org.jetbrains.uast.kotlin.declarations.KotlinUMethod
@@ -94,6 +95,10 @@ abstract class KotlinUAnnotationBase<T : KtCallElement>(
     }
 
     override fun convertParent(): UElement? {
+        sourcePsi.parent.safeAs<KtAnnotatedExpression>()?.let {
+            return it.baseExpression?.let { KotlinConverter.convertExpression(it, null, DEFAULT_EXPRESSION_TYPES_LIST) }
+        }
+
         val superParent = super.convertParent() ?: return null
         if (annotationUseSiteTarget() == AnnotationUseSiteTarget.RECEIVER) {
             (superParent.uastParent as? KotlinUMethod)?.uastParameters?.firstIsInstance<KotlinReceiverUParameter>()?.let {

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

@@ -337,7 +337,7 @@ open class KotlinUField(
 
 open class KotlinULocalVariable(
         psi: PsiLocalVariable,
-        override val sourcePsi: KtElement,
+        override val sourcePsi: KtElement?,
         givenParent: UElement?
 ) : AbstractKotlinUVariable(givenParent), ULocalVariableExPlaceHolder, PsiLocalVariable by psi {
 

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/ElvisExpression.kt

@@ -57,7 +57,7 @@ private fun createElvisExpressions(
         psiParent: PsiElement): List<UExpression> {
 
     val declaration = KotlinUDeclarationsExpression(containingElement)
-    val tempVariable = KotlinULocalVariable(UastKotlinPsiVariable.create(left, declaration, psiParent), left, declaration)
+    val tempVariable = KotlinULocalVariable(UastKotlinPsiVariable.create(left, declaration, psiParent), null, declaration)
     declaration.declarations = listOf(tempVariable)
 
     val ifExpression = object : UIfExpression, JvmDeclarationUElementPlaceholder {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUFunctionCallExpression.kt

@@ -192,6 +192,7 @@ class KotlinUFunctionCallExpression(
 
     override fun accept(visitor: UastVisitor) {
         if (visitor.visitCallExpression(this)) return
+        annotations.acceptList(visitor)
         methodIdentifier?.accept(visitor)
         classReference.accept(visitor)
         valueArguments.acceptList(visitor)

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUSimpleReferenceExpression.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.serialization.deserialization.descriptors.Deserializ
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
 import org.jetbrains.kotlin.utils.addToStdlib.constant
 import org.jetbrains.uast.*
+import org.jetbrains.uast.internal.acceptList
 import org.jetbrains.uast.internal.log
 import org.jetbrains.uast.kotlin.declarations.KotlinUIdentifier
 import org.jetbrains.uast.kotlin.internal.DelegatedMultiResolve
@@ -81,6 +82,7 @@ open class KotlinUSimpleReferenceExpression(
         if (sourcePsi.parent.destructuringDeclarationInitializer != true) {
             visitAccessorCalls(visitor)
         }
+        annotations.acceptList(visitor)
 
         visitor.afterVisitSimpleNameReferenceExpression(this)
     }
@@ -238,6 +240,10 @@ class KotlinClassViaConstructorUSimpleReferenceExpression(
         }
     }
 
+    override fun accept(visitor: UastVisitor) {
+        super<KotlinAbstractUExpression>.accept(visitor)
+    }
+
     override fun resolve(): PsiElement? = resolved
 
     override fun asLogString(): String = log<USimpleNameReferenceExpression>("identifier = $identifier, resolvesTo = $resolvedName")

plugins/uast-kotlin/testData/AnnotatedExpressions.kt

@@ -0,0 +1,23 @@
+fun foo() {
+
+    @Suppress
+    foo()
+
+    @Suppress
+    val a = 1
+
+    @Suppress
+    var b = 2
+
+    @Suppress
+    b = a
+
+    @Suppress
+    if (a > 2)
+        a
+    else
+        b
+
+    val c = @Suppress a ?: b
+
+}

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

@@ -0,0 +1,40 @@
+UFile (package = )
+    UClass (name = AnnotatedExpressionsKt)
+        UMethod (name = foo)
+            UBlockExpression
+                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                    UAnnotation (fqName = kotlin.Suppress)
+                    UIdentifier (Identifier (foo))
+                    USimpleNameReferenceExpression (identifier = foo, resolvesTo = null)
+                UDeclarationsExpression
+                    ULocalVariable (name = a)
+                        UAnnotation (fqName = kotlin.Suppress)
+                        ULiteralExpression (value = 1)
+                UDeclarationsExpression
+                    ULocalVariable (name = b)
+                        UAnnotation (fqName = kotlin.Suppress)
+                        ULiteralExpression (value = 2)
+                UBinaryExpression (operator = =)
+                    UAnnotation (fqName = kotlin.Suppress)
+                    USimpleNameReferenceExpression (identifier = b)
+                    USimpleNameReferenceExpression (identifier = a)
+                UIfExpression
+                    UAnnotation (fqName = kotlin.Suppress)
+                    UBinaryExpression (operator = >)
+                        USimpleNameReferenceExpression (identifier = a)
+                        ULiteralExpression (value = 2)
+                    USimpleNameReferenceExpression (identifier = a)
+                    USimpleNameReferenceExpression (identifier = b)
+                UDeclarationsExpression
+                    ULocalVariable (name = c)
+                        UExpressionList (elvis)
+                            UDeclarationsExpression
+                                ULocalVariable (name = varae507364)
+                                    USimpleNameReferenceExpression (identifier = a)
+                                        UAnnotation (fqName = kotlin.Suppress)
+                            UIfExpression
+                                UBinaryExpression (operator = !=)
+                                    USimpleNameReferenceExpression (identifier = varae507364)
+                                    ULiteralExpression (value = null)
+                                USimpleNameReferenceExpression (identifier = varae507364)
+                                USimpleNameReferenceExpression (identifier = b)

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

@@ -0,0 +1,13 @@
+public final class AnnotatedExpressionsKt {
+    public static final fun foo() : void {
+        foo()
+        @kotlin.Suppress var a: int = 1
+        @kotlin.Suppress var b: int = 2
+        b = a
+        if (a > 2) a else b
+        var c: int = elvis {
+            var varae507364: int = a
+            if (varae507364 != null) varae507364 else b
+        }
+    }
+}

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

@@ -16,11 +16,9 @@ UFile (package = )
                     UExpressionList (elvis)
                         UDeclarationsExpression
                             ULocalVariable (name = var243c51a0)
-                                UAnnotation (fqName = null)
                                 UExpressionList (elvis)
                                     UDeclarationsExpression
                                         ULocalVariable (name = varc4aef569)
-                                            UAnnotation (fqName = null)
                                             UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
                                                 UIdentifier (Identifier (foo))
                                                 USimpleNameReferenceExpression (identifier = foo, resolvesTo = null)

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

@@ -7,8 +7,8 @@ public final class ElvisKt {
     }
     public static final fun baz() : java.lang.String {
         return elvis {
-            @null var var243c51a0: java.lang.String = elvis {
+            var var243c51a0: java.lang.String = elvis {
-                @null var varc4aef569: java.lang.String = foo("Lorem ipsum")
+                var varc4aef569: java.lang.String = foo("Lorem ipsum")
                 if (varc4aef569 != null) varc4aef569 else foo("dolor sit amet")
             }
             if (var243c51a0 != null) var243c51a0 else bar().toString()

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

@@ -31,7 +31,6 @@ UFile (package = )
                         UExpressionList (elvis)
                             UDeclarationsExpression
                                 ULocalVariable (name = var708e23eb)
-                                    UAnnotation (fqName = org.jetbrains.annotations.Nullable)
                                     USimpleNameReferenceExpression (identifier = arg)
                             UIfExpression
                                 UBinaryExpression (operator = !=)

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

@@ -8,7 +8,7 @@ public final class TypeReferencesKt {
     public static final fun parameterizedFoo(@org.jetbrains.annotations.Nullable arg: T) : void {
         var a: T = arg
         var at: T = elvis {
-            @org.jetbrains.annotations.Nullable var var708e23eb: T = arg
+            var var708e23eb: T = arg
             if (var708e23eb != null) var708e23eb else return
         }
         var tl: java.util.List<? extends T> = listOf(at)

plugins/uast-kotlin/tests/SimpleKotlinRenderLogTest.kt

@@ -144,6 +144,9 @@ class SimpleKotlinRenderLogTest : AbstractKotlinUastTest(), AbstractKotlinRender
 
     @Test
     fun testTryCatch() = doTest("TryCatch")
+
+    @Test
+    fun testAnnotatedExpressions() = doTest("AnnotatedExpressions")
 }
 
 fun withForceUInjectionHostValue(call: () -> Unit) {