From abc93aa1d4d2d9b5f0bf9872c58f8a27c1853609 Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <nicolay.mitropolsky@jetbrains.com>
Date: Wed, 15 Nov 2017 15:11:21 +0300
Subject: [PATCH] Java-Uast: properly unwrapping call expressions for
 lazy-parent (IDEA-181557)

---
 .../JavaUQualifiedReferenceExpression.kt         | 16 ++++++++++++++++
 .../JavaUSimpleNameReferenceExpression.kt        | 11 +++++++++++
 2 files changed, 27 insertions(+)

diff --git a/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUQualifiedReferenceExpression.kt b/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUQualifiedReferenceExpression.kt
index 83c0e3fbb01e..2b82e3f3f08b 100644
--- a/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUQualifiedReferenceExpression.kt
+++ b/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUQualifiedReferenceExpression.kt
@@ -31,6 +31,8 @@ class JavaUQualifiedReferenceExpression(
     JavaUSimpleNameReferenceExpression(psi.referenceNameElement, psi.referenceName ?: "<error>", this, psi)
   }
 
+  override fun convertParent(): UElement? = super.convertParent().let(this::unwrapCompositeQualifiedReference)
+
   override val accessType: UastQualifiedExpressionAccessType
     get() = UastQualifiedExpressionAccessType.SIMPLE
 
@@ -38,4 +40,18 @@ class JavaUQualifiedReferenceExpression(
     get() = (psi.resolve() as? PsiNamedElement)?.name
 
   override fun resolve() = psi.resolve()
+}
+
+internal fun UElement.unwrapCompositeQualifiedReference(uParent: UElement?): UElement? = when (uParent) {
+  is UQualifiedReferenceExpression -> {
+    if (uParent.receiver == this || uParent.selector == this)
+      uParent
+    else {
+      val selector = uParent.selector
+      if (this is JavaUSimpleNameReferenceExpression || (selector is JavaUCallExpression && selector.valueArgumentCount > 0))
+        selector
+      else uParent.uastParent
+    }
+  }
+  else -> uParent
 }
\ No newline at end of file
diff --git a/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUSimpleNameReferenceExpression.kt b/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUSimpleNameReferenceExpression.kt
index 91380fd1ed59..fc3cd25365cc 100644
--- a/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUSimpleNameReferenceExpression.kt
+++ b/uast/uast-java/src/org/jetbrains/uast/java/expressions/JavaUSimpleNameReferenceExpression.kt
@@ -29,6 +29,17 @@ class JavaUSimpleNameReferenceExpression(
   override fun resolve() = (reference ?: psi as? PsiReference)?.resolve()
   override val resolvedName: String?
     get() = ((reference ?: psi as? PsiReference)?.resolve() as? PsiNamedElement)?.name
+
+  override fun getPsiParentForLazyConversion(): PsiElement? {
+    val parent = super.getPsiParentForLazyConversion()
+    if (parent is PsiReferenceExpression && parent.parent is PsiMethodCallExpression) {
+      return parent.parent
+    }
+    return parent
+  }
+
+  override fun convertParent(): UElement? = super.convertParent().let(this::unwrapCompositeQualifiedReference)
+
 }
 
 class JavaUTypeReferenceExpression(