UAST: Properly handle annotations on local variables

idea/testData/android/lint/wrongAnnotation.kt

@@ -15,24 +15,27 @@ class WrongAnnotation2 {
         @SuppressLint("NewApi") // Valid: class-file check on method
         fun foobar(view: View, @SuppressLint("NewApi") foo: Int) {
             // Invalid: class-file check
-            @SuppressLint("NewApi") // Invalid
+            <error descr="The `@SuppressLint` annotation cannot be used on a local variable with the lint check 'NewApi': move out to the surrounding method">@SuppressLint("NewApi")</error> // Invalid
             val a: Boolean
-            @SuppressLint("SdCardPath", "NewApi") // Invalid: class-file based check on local variable
+
+            @SuppressLint("SdCardPath", "NewApi") // TODO: Invalid, class-file based check on local variable
             val b: Boolean
-            @android.annotation.SuppressLint("SdCardPath", "NewApi") // Invalid (FQN)
+
+            @android.annotation.SuppressLint("SdCardPath", "NewApi") // TDOD: Invalid (FQN)
             val c: Boolean
+
             @SuppressLint("SdCardPath") // Valid: AST-based check
             val d: Boolean
         }
 
         init {
             // Local variable outside method: invalid
-            @SuppressLint("NewApi")
+            <error descr="The `@SuppressLint` annotation cannot be used on a local variable with the lint check 'NewApi': move out to the surrounding method">@SuppressLint("NewApi")</error>
             val localvar = 5
         }
 
         private fun test() {
-            @SuppressLint("NewApi") // Invalid
+            <error descr="The `@SuppressLint` annotation cannot be used on a local variable with the lint check 'NewApi': move out to the surrounding method">@SuppressLint("NewApi")</error> // Invalid
             val a = View.MEASURED_STATE_MASK
         }
     }

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

@@ -224,7 +224,7 @@ internal object KotlinConverter {
                 val parent = if (parentCallback == null) null else (parentCallback() ?: return null)
                 KotlinUDeclarationsExpression(parent).apply {
                     declarations = element.parameters.mapIndexed { i, p ->
-                        KotlinUVariable.create(UastKotlinPsiParameter.create(p, element, parent!!, i), this)
+                        KotlinUParameter(UastKotlinPsiParameter.create(p, element, parent!!, i), this)
                     }
                 }
             }
@@ -307,13 +307,12 @@ internal object KotlinConverter {
             is KtDestructuringDeclaration -> expr<UDeclarationsExpression> {
                 val parent = if (parentCallback == null) null else (parentCallback() ?: return null)
                 KotlinUDeclarationsExpression(parent).apply {
-                    val tempAssignment = KotlinUVariable.create(UastKotlinPsiVariable.create(expression, parent!!), parent)
+                    val tempAssignment = KotlinULocalVariable(UastKotlinPsiVariable.create(expression, parent!!), parent)
                     val destructuringAssignments = expression.entries.mapIndexed { i, entry ->
                         val psiFactory = KtPsiFactory(expression.project)
                         val initializer = psiFactory.createAnalyzableExpression("${tempAssignment.name}.component${i + 1}()",
                                                                                 expression.containingFile)
-                        KotlinUVariable.create(UastKotlinPsiVariable.create(
+                        KotlinULocalVariable(UastKotlinPsiVariable.create(entry, tempAssignment.psi, parent, initializer), parent)
-                                entry, tempAssignment.psi, parent, initializer), parent)
                     }
                     declarations = listOf(tempAssignment) + destructuringAssignments
                 }
@@ -368,7 +367,6 @@ internal object KotlinConverter {
                 expr<UDeclarationsExpression>(build(::createLocalFunctionDeclaration))
             }
 
-
             else -> expr<UExpression>(build(::UnknownKotlinExpression))
         }}
     }
@@ -403,6 +401,9 @@ private fun convertVariablesDeclaration(
         parent: UElement?
 ): UDeclarationsExpression {
     val parentPsiElement = parent?.psi
-    val variable = KotlinUVariable.create(UastKotlinPsiVariable.create(psi, parentPsiElement, parent!!), parent)
+    val variable = KotlinUAnnotatedLocalVariable(
+            UastKotlinPsiVariable.create(psi, parentPsiElement, parent!!), parent) { annotationParent ->
+        psi.annotationEntries.map { KotlinUAnnotation(it, annotationParent) }
+    }
     return KotlinUDeclarationsExpression(parent).apply { declarations = listOf(variable) }
 }

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

@@ -96,7 +96,7 @@ class KotlinUClass private constructor(
                     }
                 }
             } else {
-                getLanguagePlugin().convert<UMethod>(psiMethod, this)
+                getLanguagePlugin().convert(psiMethod, this)
             }
         }
 

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

@@ -24,16 +24,16 @@ import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtVariableDeclaration
 import org.jetbrains.uast.*
 import org.jetbrains.uast.internal.acceptList
-import org.jetbrains.uast.java.AbstractJavaUVariable
 import org.jetbrains.uast.java.JavaAbstractUExpression
 import org.jetbrains.uast.java.JavaUAnnotation
 import org.jetbrains.uast.java.annotations
 import org.jetbrains.uast.kotlin.declarations.UastLightIdentifier
+import org.jetbrains.uast.kotlin.internal.KotlinUElementWithComments
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 import org.jetbrains.uast.visitor.UastVisitor
 
-abstract class AbstractKotlinUVariable : AbstractJavaUVariable() {
+abstract class AbstractKotlinUVariable : PsiVariable, UVariable, KotlinUElementWithComments {
     override val uastInitializer: UExpression?
         get() {
             val psi = psi
@@ -68,6 +68,17 @@ abstract class AbstractKotlinUVariable : AbstractJavaUVariable() {
         val kotlinOrigin = (psi as? KtLightElement<*, *>)?.kotlinOrigin
         return UastLightIdentifier(psi, kotlinOrigin as KtNamedDeclaration?)
     }
+
+    override val annotations: List<UAnnotation> by lz { psi.annotations.map { JavaUAnnotation(it, this) } }
+
+    override val typeReference by lz { getLanguagePlugin().convertOpt<UTypeReferenceExpression>(psi.typeElement, this) }
+
+    override val uastAnchor: UElement?
+        get() = UIdentifier(psi.nameIdentifier, this)
+
+    override fun equals(other: Any?) = other is AbstractKotlinUVariable && psi == other.psi
+
+    override fun hashCode() = psi.hashCode()
 }
 
 class KotlinUVariable(
@@ -77,6 +88,7 @@ class KotlinUVariable(
     override val psi = unwrap<UVariable, PsiVariable>(psi)
 
     override val annotations by lz { psi.annotations.map { JavaUAnnotation(it, this) } }
+
     override val typeReference by lz { getLanguagePlugin().convertOpt<UTypeReferenceExpression>(psi.typeElement, this) }
 
     override fun getContainingFile(): PsiFile? = (psi as? KtLightElement<*, *>)?.kotlinOrigin?.containingFile ?: psi.containingFile
@@ -92,18 +104,6 @@ class KotlinUVariable(
     override fun getNameIdentifier(): PsiIdentifier {
         return super.getNameIdentifier()
     }
-
-    companion object {
-        fun create(psi: PsiVariable, containingElement: UElement?): UVariable {
-            return when (psi) {
-                is PsiEnumConstant -> KotlinUEnumConstant(psi, containingElement)
-                is PsiLocalVariable -> KotlinULocalVariable(psi, containingElement)
-                is PsiParameter -> KotlinUParameter(psi, containingElement)
-                is PsiField -> KotlinUField(psi, containingElement)
-                else -> KotlinUVariable(psi, containingElement)
-            }
-        }
-    }
 }
 
 open class KotlinUParameter(
@@ -188,6 +188,15 @@ open class KotlinULocalVariable(
     }
 }
 
+open class KotlinUAnnotatedLocalVariable(
+        psi: PsiLocalVariable,
+        uastParent: UElement?,
+        computeAnnotations: (parent: UElement) -> List<UAnnotation>
+) : KotlinULocalVariable(psi, uastParent) {
+
+    override val annotations: List<UAnnotation> by lz { computeAnnotations(this) }
+}
+
 open class KotlinUEnumConstant(
         psi: PsiEnumConstant,
         override val uastParent: UElement?

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

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

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinUElementWithComments.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.uast.kotlin.internal
+
+import com.intellij.psi.PsiComment
+import org.jetbrains.uast.UComment
+import org.jetbrains.uast.UElement
+
+interface KotlinUElementWithComments : UElement {
+    override val comments: List<UComment>
+        get() {
+            val psi = psi ?: return emptyList()
+            return psi.children.filter { it is PsiComment }.map { UComment(it, this) }
+        }
+}

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

@@ -9,6 +9,7 @@ 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.kotlin.utils.ifEmpty
 import org.jetbrains.uast.UDeclaration
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UastErrorType
@@ -62,7 +63,7 @@ class UastKotlinPsiVariable(
                 parent: PsiElement?, 
                 containingElement: UElement, 
                 initializer: KtExpression? = null
-        ): PsiVariable {
+        ): PsiLocalVariable {
             val psiParent = containingElement.getParentOfType<UDeclaration>()?.psi ?: parent
             return UastKotlinPsiVariable(
                     declaration.manager,
@@ -74,7 +75,7 @@ class UastKotlinPsiVariable(
                     declaration)
         }
         
-        fun create(declaration: KtDestructuringDeclaration, containingElement: UElement): PsiVariable {
+        fun create(declaration: KtDestructuringDeclaration, containingElement: UElement): PsiLocalVariable {
             val psiParent = containingElement.getParentOfType<UDeclaration>()?.psi ?: declaration.parent
             return UastKotlinPsiVariable(
                     declaration.manager,
@@ -86,7 +87,7 @@ class UastKotlinPsiVariable(
                     declaration)
         }
 
-        fun create(initializer: KtExpression, containingElement: UElement, parent: PsiElement): PsiVariable {
+        fun create(initializer: KtExpression, containingElement: UElement, parent: PsiElement): PsiLocalVariable {
             val psiParent = containingElement.getParentOfType<UDeclaration>()?.psi ?: parent
             return UastKotlinPsiVariable(
                     initializer.manager,
@@ -98,7 +99,7 @@ class UastKotlinPsiVariable(
                     initializer)
         }
 
-        fun create(name: String, localFunction: KtFunction, containingElement: UElement): PsiVariable {
+        fun create(name: String, localFunction: KtFunction, containingElement: UElement): PsiLocalVariable {
             val psiParent = containingElement.getParentOfType<UDeclaration>()?.psi ?: localFunction.parent
             return UastKotlinPsiVariable(
                     localFunction.manager,