From 75a511fb48d7884b79303e8420255e1ba62dd8bc Mon Sep 17 00:00:00 2001
From: Vyacheslav Gerasimov <vyacheslav.gerasimov@jetbrains.com>
Date: Tue, 1 Aug 2017 22:45:57 +0300
Subject: [PATCH] UAST: Fix getUastParentOfType, skip psi parents without
 UElement

---
 .../src/org/jetbrains/uast/UastContext.kt            | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/uast/uast-common/src/org/jetbrains/uast/UastContext.kt b/uast/uast-common/src/org/jetbrains/uast/UastContext.kt
index 5d6a54eda809..6b304ecf7221 100644
--- a/uast/uast-common/src/org/jetbrains/uast/UastContext.kt
+++ b/uast/uast-common/src/org/jetbrains/uast/UastContext.kt
@@ -130,15 +130,13 @@ fun <T : UElement> PsiFile.findUElementAt(offset: Int, cls: Class<out T>): T? {
  * Finds an UAST element of the given type among the parents of the given PSI element.
  */
 @JvmOverloads
-fun <T : UElement> PsiElement?.getUastParentOfType(cls: Class<out T>, strict: Boolean = false): T? {
-    val uElement = this.toUElement() ?: return null
-    val sequence = if (strict)
-        (uElement.uastParent?.withContainingElements ?: emptySequence())
-    else
-        uElement.withContainingElements
+fun <T : UElement> PsiElement?.getUastParentOfType(cls: Class<out T>, strict: Boolean = false): T? = this?.run {
+    val startingElement = if (strict) this.parent else this
+    val parentSequence = generateSequence(startingElement, PsiElement::getParent)
+    val firstUElement = parentSequence.mapNotNull { it.toUElement() }.firstOrNull() ?: return null
 
     @Suppress("UNCHECKED_CAST")
-    return sequence.firstOrNull { cls.isInstance(it) } as T?
+    return firstUElement.withContainingElements.firstOrNull { cls.isInstance(it) } as T?
 }
 
 inline fun <reified T : UElement> PsiElement?.getUastParentOfType(strict: Boolean = false): T? = getUastParentOfType(T::class.java, strict)