diff --git a/plugins/devkit/src/references/TestDataFilesReferencesContributor.java b/plugins/devkit/src/references/TestDataFilesReferencesContributor.java index ed20c2cf8d13..83f32dda09c5 100644 --- a/plugins/devkit/src/references/TestDataFilesReferencesContributor.java +++ b/plugins/devkit/src/references/TestDataFilesReferencesContributor.java @@ -27,11 +27,12 @@ public class TestDataFilesReferencesContributor extends PsiReferenceContributor .registerUastReferenceProvider( registrar, UastPatterns.stringLiteralExpression().inCall(UastPatterns.callExpression()), - new UastReferenceProvider() { + new UastLiteralReferenceProvider() { @NotNull @Override - public PsiReference[] getReferencesByElement(@NotNull UElement element, @NotNull ProcessingContext context) { - ULiteralExpression expression = (ULiteralExpression)element; + public PsiReference[] getReferencesByULiteral(@NotNull ULiteralExpression expression, + @NotNull PsiLanguageInjectionHost host, + @NotNull ProcessingContext context) { UCallExpression call = UastUtils.getUCallExpression(expression); if (call == null) return PsiReference.EMPTY_ARRAY; @@ -45,9 +46,6 @@ public class TestDataFilesReferencesContributor extends PsiReferenceContributor return PsiReference.EMPTY_ARRAY; } - PsiLanguageInjectionHost host = UastLiteralUtils.getPsiLanguageInjectionHost(expression); - if (host == null) return PsiReference.EMPTY_ARRAY; - FileReferenceSet fileReferenceSet = new FileReferenceSet(host); fileReferenceSet.addCustomization( FileReferenceSet.DEFAULT_PATH_EVALUATOR_OPTION, diff --git a/plugins/devkit/src/testAssistant/TestDataReferenceContributor.java b/plugins/devkit/src/testAssistant/TestDataReferenceContributor.java index 927f196b2521..b6cdcda22859 100644 --- a/plugins/devkit/src/testAssistant/TestDataReferenceContributor.java +++ b/plugins/devkit/src/testAssistant/TestDataReferenceContributor.java @@ -33,7 +33,6 @@ import com.intellij.util.ProcessingContext; import com.intellij.util.containers.ContainerUtil; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; -import org.jetbrains.uast.UElement; import org.jetbrains.uast.ULiteralExpression; import org.jetbrains.uast.UastLiteralUtils; @@ -54,16 +53,12 @@ public class TestDataReferenceContributor extends PsiReferenceContributor { new TestDataReferenceProvider(), PsiReferenceRegistrar.DEFAULT_PRIORITY); } - private static class TestDataReferenceProvider extends UastReferenceProvider { + private static class TestDataReferenceProvider extends UastLiteralReferenceProvider { @NotNull @Override - public PsiReference[] getReferencesByElement(@NotNull final UElement element, @NotNull final ProcessingContext context) { - if (!(element instanceof ULiteralExpression)) return PsiReference.EMPTY_ARRAY; - - ULiteralExpression literalExpression = (ULiteralExpression)element; - - PsiLanguageInjectionHost host = UastLiteralUtils.getPsiLanguageInjectionHost(literalExpression); - if (host == null) return PsiReference.EMPTY_ARRAY; + public PsiReference[] getReferencesByULiteral(@NotNull final ULiteralExpression literalExpression, + @NotNull PsiLanguageInjectionHost host, + @NotNull final ProcessingContext context) { TextRange range = ElementManipulators.getValueTextRange(host); diff --git a/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt b/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt index 378053eea949..5ef40b3270c3 100644 --- a/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt +++ b/uast/uast-common/src/com/intellij/psi/UastReferenceRegistrar.kt @@ -30,44 +30,69 @@ import org.jetbrains.uast.toUElement fun PsiReferenceRegistrar.registerUastReferenceProvider(pattern: (UElement, ProcessingContext) -> Boolean, provider: UastReferenceProvider, priority: Double = PsiReferenceRegistrar.DEFAULT_PRIORITY) { - this.registerReferenceProvider(UastPatternAdapter(pattern), UastReferenceProviderAdapter(provider), priority) + this.registerReferenceProvider(UastPatternAdapter(pattern, provider.supportedUElementTypes), + UastReferenceProviderAdapter(provider), + priority) } fun PsiReferenceRegistrar.registerUastReferenceProvider(pattern: ElementPattern, provider: UastReferenceProvider, priority: Double = PsiReferenceRegistrar.DEFAULT_PRIORITY) { - this.registerReferenceProvider(UastPatternAdapter(pattern::accepts), UastReferenceProviderAdapter(provider), priority) + this.registerReferenceProvider(UastPatternAdapter(pattern::accepts, provider.supportedUElementTypes), + UastReferenceProviderAdapter(provider), priority) } abstract class UastReferenceProvider { + open val supportedUElementTypes: List> = listOf(UElement::class.java) + abstract fun getReferencesByElement(element: UElement, context: ProcessingContext): Array } +abstract class UastLiteralReferenceProvider : UastReferenceProvider() { + + override val supportedUElementTypes: List> = listOf(ULiteralExpression::class.java) + + override fun getReferencesByElement(element: UElement, context: ProcessingContext): Array { + val uLiteral = element as? ULiteralExpression ?: return PsiReference.EMPTY_ARRAY + val host = uLiteral.psiLanguageInjectionHost ?: return PsiReference.EMPTY_ARRAY + return getReferencesByULiteral(uLiteral, host, context) + } + + abstract fun getReferencesByULiteral(uLiteral: ULiteralExpression, + host: PsiLanguageInjectionHost, + context: ProcessingContext): Array + +} + fun uastLiteralReferenceProvider(provider: (ULiteralExpression, PsiLanguageInjectionHost) -> Array) = - object : UastReferenceProvider() { - override fun getReferencesByElement(element: UElement, context: ProcessingContext): Array { - val uLiteral = element as? ULiteralExpression ?: return PsiReference.EMPTY_ARRAY - val host = uLiteral.psiLanguageInjectionHost ?: return PsiReference.EMPTY_ARRAY - return provider(uLiteral, host) - } + object : UastLiteralReferenceProvider() { + + override fun getReferencesByULiteral(uLiteral: ULiteralExpression, + host: PsiLanguageInjectionHost, + context: ProcessingContext): Array = provider(uLiteral, host) } 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 fun getOrCreateCachedElement(element: PsiElement, + context: ProcessingContext?, + supportedUElementTypes: List>): UElement? = + element as? UElement ?: context?.get(cachedUElement) ?: supportedUElementTypes.asSequence().mapNotNull { + element.toUElement(it) + }.firstOrNull()?.also { context?.put(cachedUElement, it) } -private class UastPatternAdapter(val predicate: (UElement, ProcessingContext) -> Boolean) : ElementPattern { +private class UastPatternAdapter( + val predicate: (UElement, ProcessingContext) -> Boolean, + val supportedUElementTypes: List> +) : 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) + getOrCreateCachedElement(o, context, supportedUElementTypes) ?.let { predicate(it, context ?: ProcessingContext()) } ?: false else -> false @@ -82,7 +107,7 @@ private class UastPatternAdapter(val predicate: (UElement, ProcessingContext) -> private class UastReferenceProviderAdapter(val provider: UastReferenceProvider) : PsiReferenceProvider() { override fun getReferencesByElement(element: PsiElement, context: ProcessingContext): Array { - val uElement = getOrCreateCachedElement(element, context) ?: return PsiReference.EMPTY_ARRAY + val uElement = getOrCreateCachedElement(element, context, provider.supportedUElementTypes) ?: return PsiReference.EMPTY_ARRAY return provider.getReferencesByElement(uElement, context) } diff --git a/uast/uast-tests/test/com/intellij/psi/UastReferenceRegistrarTest.kt b/uast/uast-tests/test/com/intellij/psi/UastReferenceRegistrarTest.kt new file mode 100644 index 000000000000..bc30a031896c --- /dev/null +++ b/uast/uast-tests/test/com/intellij/psi/UastReferenceRegistrarTest.kt @@ -0,0 +1,71 @@ +// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. + +package com.intellij.psi + +import com.intellij.patterns.ElementPattern +import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase +import com.intellij.util.ProcessingContext +import junit.framework.TestCase +import org.junit.Test + +class UastReferenceRegistrarTest : LightCodeInsightFixtureTestCase() { + + @Test + fun testUastReferenceContributorCalledOnlyOnLiterals() { + + /** + * This implementation is used instead of [ReferenceProvidersRegistry] to emulate behaviour that [PsiReferenceProvider] will be called on + * each [PsiElement] not only [ContributedReferenceHost]-s as it is usually done for Java + */ + val psiReferenceRegistrar = object : PsiReferenceRegistrar() { + + private val providers = mutableListOf, PsiReferenceProvider>>() + + fun getReferencesFor(psiElement: PsiElement): List { + val context = ProcessingContext() + return providers.flatMap { (pattern, provider) -> + if (pattern.accepts(psiElement, context) && provider.acceptsTarget(psiElement)) + provider.getReferencesByElement(psiElement, context).toList() + else emptyList() + } + } + + override fun registerReferenceProvider(pattern: ElementPattern, + provider: PsiReferenceProvider, + priority: Double) { + providers += pattern to provider + } + + } + + val expectedInvocationCount = 10 + var invocationCount = 0 + + psiReferenceRegistrar.registerUastReferenceProvider( + { _, _ -> invocationCount++; true }, + uastLiteralReferenceProvider { _, _ -> PsiReference.EMPTY_ARRAY } + ) + + myFixture.addClass(""" + class MyClass { + String foo(){ + return new java.lang.StringBuilder() + ${(1..10).joinToString("\n") { ".append(\"$it\")" }} + .toString(); + } + } + """.trimIndent()) + + + myFixture.configureByFile("MyClass.java") + myFixture.file.accept(object : JavaRecursiveElementWalkingVisitor() { + override fun visitElement(element: PsiElement) { + psiReferenceRegistrar.getReferencesFor(element) + super.visitElement(element) + } + }) + TestCase.assertEquals(expectedInvocationCount, invocationCount) + + } + +} \ No newline at end of file