From 29c16f956163d4f978258847895d068c0f69a396 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Thu, 19 Apr 2018 13:27:36 +0300
Subject: [PATCH] Improve KotlinUElementWithComments & use it in most
 KotlinUElements

Related to KT-22135
---
 .../uast/kotlin/KotlinAbstractUElement.kt     |  9 ++++---
 .../uast/kotlin/declarations/KotlinUMethod.kt |  3 +++
 .../kotlin/declarations/KotlinUVariable.kt    |  2 +-
 .../internal/KotlinUElementWithComments.kt    | 25 ++++++++++++++++++-
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt
index 3a3df78d801..1f6878c5ae3 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt
@@ -29,9 +29,10 @@ import org.jetbrains.kotlin.psi.psiUtil.isPropertyParameter
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
+import org.jetbrains.uast.kotlin.internal.KotlinUElementWithComments
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
-abstract class KotlinAbstractUElement(private val givenParent: UElement?) : UElement, JvmDeclarationUElement {
+abstract class KotlinAbstractUElement(private val givenParent: UElement?) : KotlinUElementWithComments, UElement, JvmDeclarationUElement {
 
     final override val uastParent: UElement? by lz {
         givenParent ?: convertParent()
@@ -194,8 +195,10 @@ private fun findAnnotationClassFromConstructorParameter(parameter: KtParameter):
     return null
 }
 
-abstract class KotlinAbstractUExpression(givenParent: UElement?)
-    : KotlinAbstractUElement(givenParent), UExpression, JvmDeclarationUElement {
+abstract class KotlinAbstractUExpression(givenParent: UElement?) :
+    KotlinAbstractUElement(givenParent),
+    UExpression,
+    JvmDeclarationUElement {
 
     override val javaPsi: PsiElement? = null
 
diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt
index 5a2d576619e..7d68de53c03 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt
@@ -36,6 +36,9 @@ open class KotlinUMethod(
         psi: KtLightMethod,
         givenParent: UElement?
 ) : KotlinAbstractUElement(givenParent), UAnnotationMethod, JavaUElementWithComments, PsiMethod by psi {
+    override val comments: List<UComment>
+        get() = super<KotlinAbstractUElement>.comments
+
     override val psi: KtLightMethod = unwrap<UMethod, KtLightMethod>(psi)
 
     override val javaPsi = psi
diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt
index 78a1665e80c..6926a5e9353 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt
@@ -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() {
diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinUElementWithComments.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinUElementWithComments.kt
index ae834dd348f..a7812a27c86 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinUElementWithComments.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinUElementWithComments.kt
@@ -17,14 +17,37 @@
 package org.jetbrains.uast.kotlin.internal
 
 import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.uast.JvmDeclarationUElement
 import org.jetbrains.uast.UComment
 import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UExpression
 
 interface KotlinUElementWithComments : UElement, JvmDeclarationUElement {
+
     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
+    }
+
 }
\ No newline at end of file