From 4e2ac7f23852555f3e9b63040d88e7b0013a4432 Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <Nicolay.Mitropolsky@jetbrains.com>
Date: Fri, 26 Apr 2019 13:52:33 +0300
Subject: [PATCH] Uast-pattens: registering Uast-patterns only on
 `PsiLanguageInjectionHosts`

as long as now we can rely on `UInjectionHost` existence for both Java and Kotlin

GitOrigin-RevId: 121e36a242e4155cf8ee46f970ed67e0f3dd6d89
---
 .../intellij/psi/UastReferenceRegistrar.kt    | 21 +++++++++++++++++--
 1 file changed, 19 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 7d54fb7c0a93..38e8f40e7a02 100644
--- a/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt
+++ b/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt
@@ -23,15 +23,17 @@ import com.intellij.openapi.util.RecursionManager
 import com.intellij.patterns.ElementPattern
 import com.intellij.patterns.ElementPatternCondition
 import com.intellij.patterns.InitialPatternCondition
+import com.intellij.patterns.StandardPatterns
 import com.intellij.util.ProcessingContext
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UExpression
+import org.jetbrains.uast.expressions.UInjectionHost
 import org.jetbrains.uast.toUElement
 
 fun PsiReferenceRegistrar.registerUastReferenceProvider(pattern: (UElement, ProcessingContext) -> Boolean,
                                                         provider: UastReferenceProvider,
                                                         priority: Double = PsiReferenceRegistrar.DEFAULT_PRIORITY) {
-  this.registerReferenceProvider(UastPatternAdapter(pattern, provider.supportedUElementTypes),
+  this.registerReferenceProvider(adaptPattern(pattern, provider.supportedUElementTypes),
                                  UastReferenceProviderAdapter(provider),
                                  priority)
 }
@@ -39,7 +41,7 @@ fun PsiReferenceRegistrar.registerUastReferenceProvider(pattern: (UElement, Proc
 fun PsiReferenceRegistrar.registerUastReferenceProvider(pattern: ElementPattern<out UElement>,
                                                         provider: UastReferenceProvider,
                                                         priority: Double = PsiReferenceRegistrar.DEFAULT_PRIORITY) {
-  this.registerReferenceProvider(UastPatternAdapter(pattern::accepts, provider.supportedUElementTypes),
+  this.registerReferenceProvider(adaptPattern(pattern::accepts, provider.supportedUElementTypes),
                                  UastReferenceProviderAdapter(provider), priority)
 }
 
@@ -80,6 +82,21 @@ private fun getOrCreateCachedElement(element: PsiElement,
     element.toUElement(it)
   }.firstOrNull()?.also { context?.put(cachedUElement, it) }
 
+
+private fun adaptPattern(
+  predicate: (UElement, ProcessingContext) -> Boolean,
+  supportedUElementTypes: List<Class<out UElement>>
+): ElementPattern<out PsiElement> {
+  val uastPatternAdapter = UastPatternAdapter(predicate, supportedUElementTypes)
+
+  // optimisation until IDEA-211738 is implemented
+  if (supportedUElementTypes == listOf(UInjectionHost::class.java)) {
+    return StandardPatterns.instanceOf(PsiLanguageInjectionHost::class.java).and(uastPatternAdapter)
+  }
+
+  return uastPatternAdapter
+}
+
 private class UastPatternAdapter(
   val predicate: (UElement, ProcessingContext) -> Boolean,
   val supportedUElementTypes: List<Class<out UElement>>