UastReferenceRegistrar to register UastReferenceProviders

2026-09-27 10:03:11 +07:00 · 2017-09-18 14:30:41 +03:00
parent 6f2c9b9623
commit 0b357185f0
1 changed files with 72 additions and 0 deletions
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+@file:JvmName("UastReferenceRegistrar")
+
+package com.intellij.psi
+
+import com.intellij.openapi.util.Key
+import com.intellij.patterns.ElementPattern
+import com.intellij.patterns.ElementPatternCondition
+import com.intellij.patterns.InitialPatternCondition
+import com.intellij.util.ProcessingContext
+import org.jetbrains.uast.UElement
+import org.jetbrains.uast.toUElement
+
+fun PsiReferenceRegistrar.registerUastReferenceProvider(pattern: (UElement, ProcessingContext) -> Boolean,
+                                                        provider: UastReferenceProvider) {
+  this.registerReferenceProvider(UastPatternAdapter(pattern), UastReferenceProviderAdapter(provider))
+}
+
+abstract class UastReferenceProvider {
+
+  abstract fun getReferencesByElement(element: UElement, context: ProcessingContext): Array<PsiReference>
+
+}
+
+private val cachedUElement = Key.create<UElement>("UastReferenceRegistrar.cachedUElement")
+
+private fun getOrCreateCachedElement(element: PsiElement, context: ProcessingContext?): UElement? =
+  element as? UElement ?:
+  context?.get(cachedUElement) ?:
+  element.toUElement()?.also { context?.put(cachedUElement, it) }
+
+private class UastPatternAdapter(val predicate: (UElement, ProcessingContext) -> Boolean) : ElementPattern<PsiElement> {
+
+  override fun accepts(o: Any?): Boolean = accepts(o, null)
+
+  override fun accepts(o: Any?, context: ProcessingContext?): Boolean = when (o) {
+    is PsiElement ->
+      getOrCreateCachedElement(o, context)
+        ?.let { predicate(it, context ?: ProcessingContext()) }
+      ?: false
+    else -> false
+  }
+
+  private val condition = ElementPatternCondition(object : InitialPatternCondition<PsiElement>(PsiElement::class.java) {
+    override fun accepts(o: Any?, context: ProcessingContext?): Boolean = this@UastPatternAdapter.accepts(o, context)
+  })
+
+  override fun getCondition(): ElementPatternCondition<PsiElement> = condition
+}
+
+private class UastReferenceProviderAdapter(val provider: UastReferenceProvider) : PsiReferenceProvider() {
+  override fun getReferencesByElement(element: PsiElement, context: ProcessingContext): Array<PsiReference> {
+    val uElement = getOrCreateCachedElement(element, context) ?: return PsiReference.EMPTY_ARRAY
+    return provider.getReferencesByElement(uElement, context)
+  }
+
+  override fun acceptsTarget(target: PsiElement): Boolean = true
+}