From a162c714c2f5e9553ffb89f7502548c6c187add1 Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <nicolay.mitropolsky@jetbrains.com>
Date: Mon, 30 Jul 2018 18:21:49 +0300
Subject: [PATCH] Uast: a few util methods added

---
 .../src/com/intellij/psi/UastReferenceRegistrar.kt          | 5 +++++
 uast/uast-common/src/org/jetbrains/uast/UastUtils.kt        | 6 ++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt b/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt
index 548e95834395..228713ed310e 100644
--- a/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt
+++ b/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt
@@ -105,6 +105,11 @@ private class UastPatternAdapter(
   override fun getCondition(): ElementPatternCondition<PsiElement> = condition
 }
 
+fun ElementPattern<out UElement>.asPsiPattern(vararg supportedUElementTypes: Class<out UElement>): ElementPattern<PsiElement> = UastPatternAdapter(
+  this::accepts,
+  if (supportedUElementTypes.isNotEmpty()) supportedUElementTypes.toList() else listOf(UElement::class.java)
+)
+
 private class UastReferenceProviderAdapter(val provider: UastReferenceProvider) : PsiReferenceProvider() {
   override fun getReferencesByElement(element: PsiElement, context: ProcessingContext): Array<PsiReference> {
     val uElement = getOrCreateCachedElement(element, context, provider.supportedUElementTypes) ?: return PsiReference.EMPTY_ARRAY
diff --git a/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt b/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt
index 1199a9b082a7..1310432c9dc6 100644
--- a/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt
+++ b/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt
@@ -105,10 +105,12 @@ fun UElement.getContainingVariable(): PsiVariable? = getContainingUVariable()?.p
             replaceWith = ReplaceWith("PsiTreeUtil.getParentOfType(this, PsiClass::class.java)"))
 fun PsiElement?.getContainingClass(): PsiClass? = this?.let { PsiTreeUtil.getParentOfType(it, PsiClass::class.java) }
 
-fun PsiElement?.findContainingUClass(): UClass? {
+fun PsiElement?.findContainingUClass(): UClass? = findContaining(UClass::class.java)
+
+fun <T : UElement> PsiElement?.findContaining(clazz: Class<T>): T? {
   var element = this
   while (element != null && element !is PsiFileSystemItem) {
-    element.toUElementOfType<UClass>()?.let { return it }
+    element.toUElement(clazz)?.let { return it }
     element = element.parent
   }
   return null