diff --git a/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt b/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt new file mode 100644 index 000000000000..854d5562fa9b --- /dev/null +++ b/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt @@ -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 + +} + +private val cachedUElement = Key.create("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 { + + 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::class.java) { + override fun accepts(o: Any?, context: ProcessingContext?): Boolean = this@UastPatternAdapter.accepts(o, context) + }) + + override fun getCondition(): ElementPatternCondition = condition +} + +private class UastReferenceProviderAdapter(val provider: UastReferenceProvider) : PsiReferenceProvider() { + override fun getReferencesByElement(element: PsiElement, context: ProcessingContext): Array { + val uElement = getOrCreateCachedElement(element, context) ?: return PsiReference.EMPTY_ARRAY + return provider.getReferencesByElement(uElement, context) + } + + override fun acceptsTarget(target: PsiElement): Boolean = true +} \ No newline at end of file