Up KotlinAbstractUElement / UElementWithComments / UVariable in branches

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

@@ -28,7 +28,7 @@ import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
-abstract class KotlinAbstractUElement(private val givenParent: UElement?) : UElement {
+abstract class KotlinAbstractUElement(private val givenParent: UElement?) : KotlinUElementWithComments, UElement {
 
     override val uastParent: UElement? by lz {
         givenParent ?: convertParent()

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

@@ -35,7 +35,7 @@ import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 import org.jetbrains.uast.visitor.UastVisitor
 
 abstract class AbstractKotlinUVariable(givenParent: UElement?)
-    : KotlinAbstractUElement(givenParent), PsiVariable, UVariable, KotlinUElementWithComments {
+    : KotlinAbstractUElement(givenParent), PsiVariable, UVariable {
 
     override val uastInitializer: UExpression?
         get() {

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

@@ -39,8 +39,7 @@ import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 import org.jetbrains.uast.visitor.UastVisitor
 
-abstract class AbstractKotlinUVariable(givenParent: UElement?) : KotlinAbstractUElement(givenParent), PsiVariable, UVariable, UAnchorOwner,
+abstract class AbstractKotlinUVariable(givenParent: UElement?) : KotlinAbstractUElement(givenParent), PsiVariable, UVariable, UAnchorOwner {
-    KotlinUElementWithComments {
 
     override val uastInitializer: UExpression?
         get() {

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

@@ -39,7 +39,7 @@ import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 import org.jetbrains.uast.visitor.UastVisitor
 
 abstract class AbstractKotlinUVariable(givenParent: UElement?)
-    : KotlinAbstractUElement(givenParent), PsiVariable, UVariable, KotlinUElementWithComments {
+    : KotlinAbstractUElement(givenParent), PsiVariable, UVariable {
 
     override val uastInitializer: UExpression?
         get() {

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

@@ -21,9 +21,28 @@ 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) }
+            val childrenComments = psi.children.filterIsInstance<PsiComment>().map { UComment(it, this) }
+            if (this !is UExpression) return childrenComments
+            val childrenAndSiblingComments = childrenComments +
+                    psi.nearestCommentSibling(forward = true)?.let { listOf(UComment(it, this)) }.orEmpty() +
+                    psi.nearestCommentSibling(forward = false)?.let { listOf(UComment(it, this)) }.orEmpty()
+            val parent = psi.parent as? KtValueArgument ?: return childrenAndSiblingComments
+
+            return childrenAndSiblingComments +
+                    parent.nearestCommentSibling(forward = true)?.let { listOf(UComment(it, this)) }.orEmpty() +
+                    parent.nearestCommentSibling(forward = false)?.let { listOf(UComment(it, this)) }.orEmpty()
         }
+
+    private fun PsiElement.nearestCommentSibling(forward: Boolean): PsiComment? {
+        var sibling = if (forward) nextSibling else prevSibling
+        while (sibling is PsiWhiteSpace && !sibling.text.contains('\n')) {
+            sibling = if (forward) sibling.nextSibling else sibling.prevSibling
+        }
+        return sibling as? PsiComment
+    }
+
 }