Find Usages: Include super declarations when highliting usages in the current file

idea/src/org/jetbrains/kotlin/idea/findUsages/handlers/KotlinFindMemberUsagesHandler.kt

@@ -17,16 +17,21 @@
 package org.jetbrains.kotlin.idea.findUsages.handlers
 
 import com.intellij.codeInsight.highlighting.ReadWriteAccessDetector
+import com.intellij.find.FindManager
 import com.intellij.find.findUsages.AbstractFindUsagesDialog
 import com.intellij.find.findUsages.FindUsagesOptions
+import com.intellij.find.impl.FindManagerImpl
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiReference
+import com.intellij.psi.search.SearchScope
 import com.intellij.psi.search.searches.MethodReferencesSearch
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.usageView.UsageInfo
 import com.intellij.util.*
 import org.jetbrains.kotlin.asJava.toLightMethods
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
 import org.jetbrains.kotlin.idea.findUsages.KotlinCallableFindUsagesOptions
 import org.jetbrains.kotlin.idea.findUsages.KotlinFindUsagesHandlerFactory
 import org.jetbrains.kotlin.idea.findUsages.KotlinFunctionFindUsagesOptions
@@ -40,8 +45,11 @@ import org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinReferencesSearchOpt
 import org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinReferencesSearchParameters
 import org.jetbrains.kotlin.idea.search.usagesSearch.isImportUsage
 import org.jetbrains.kotlin.idea.util.application.runReadAction
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtNamedDeclaration
+import org.jetbrains.kotlin.resolve.getOriginalTopmostOverriddenDescriptors
+import org.jetbrains.kotlin.resolve.source.getPsi
 
 public abstract class KotlinFindMemberUsagesHandler<T : KtNamedDeclaration>
     protected constructor(declaration: T, elementsToSearch: Collection<PsiElement>, factory: KotlinFindUsagesHandlerFactory)
@@ -154,6 +162,22 @@ public abstract class KotlinFindMemberUsagesHandler<T : KtNamedDeclaration>
 
     override fun isSearchForTextOccurencesAvailable(psiElement: PsiElement, isSingleFile: Boolean): Boolean = !isSingleFile
 
+    override fun findReferencesToHighlight(target: PsiElement, searchScope: SearchScope): Collection<PsiReference> {
+        val callableDescriptor = (target as? KtCallableDeclaration)?.resolveToDescriptorIfAny() as? CallableDescriptor
+        val baseDescriptors = callableDescriptor?.getOriginalTopmostOverriddenDescriptors() ?: emptyList<CallableDescriptor>()
+        val baseDeclarations = baseDescriptors.map { it.source.getPsi() }.filter { it != null && it != target }
+
+        return if (baseDeclarations.isNotEmpty()) {
+            baseDeclarations.flatMap {
+                val handler = (FindManager.getInstance(project) as FindManagerImpl).findUsagesManager.getFindUsagesHandler(it!!, true)
+                handler?.findReferencesToHighlight(it!!, searchScope) ?: emptyList()
+            }
+        }
+        else {
+            super.findReferencesToHighlight(target, searchScope)
+        }
+    }
+
     companion object {
 
         public fun getInstance(declaration: KtNamedDeclaration,

idea/testData/findUsages/kotlin/findFunctionUsages/highlightingOfSuperUsages.0.kt

@@ -0,0 +1,17 @@
+// PSI_ELEMENT: org.jetbrains.kotlin.psi.KtNamedFunction
+// OPTIONS: usages, skipImports
+// HIGHLIGHTING
+
+open class A : I {
+    open fun foo() {}
+}
+
+class B: A() {
+    override fun <caret>foo() {}
+}
+
+fun test(i: I) {
+    i.foo()
+    A().foo()
+    B().foo()
+}

idea/testData/findUsages/kotlin/findFunctionUsages/highlightingOfSuperUsages.1.java

@@ -0,0 +1,3 @@
+interface I {
+    void foo();
+}

idea/testData/findUsages/kotlin/findFunctionUsages/highlightingOfSuperUsages.results.txt

@@ -0,0 +1,3 @@
+Function call (14: 7) i.foo()
+Function call (15: 9) A().foo()
+Function call (16: 9) B().foo()

idea/tests/org/jetbrains/kotlin/findUsages/AbstractJetFindUsagesTest.java

@@ -376,7 +376,9 @@ public abstract class AbstractJetFindUsagesTest extends JetLightCodeInsightFixtu
             @NotNull T caretElement,
             @Nullable FindUsagesOptions options
     ) throws ClassNotFoundException, InstantiationException, IllegalAccessException {
-        Collection<UsageInfo> usageInfos = findUsages(caretElement, options);
+        boolean highlightingMode = InTextDirectivesUtils.isDirectiveDefined(mainFileText, "// HIGHLIGHTING");
+
+        Collection<UsageInfo> usageInfos = findUsages(caretElement, options, highlightingMode);
 
         Collection<UsageFilteringRule> filteringRules = instantiateClasses(mainFileText, "// FILTERING_RULES: ");
         final Collection<UsageGroupingRule> groupingRules = instantiateClasses(mainFileText, "// GROUPING_RULES: ");
@@ -431,7 +433,11 @@ public abstract class AbstractJetFindUsagesTest extends JetLightCodeInsightFixtu
         JetTestUtils.assertEqualsToFile(new File(rootPath, prefix + "results.txt"), StringUtil.join(finalUsages, "\n"));
     }
 
-    protected Collection<UsageInfo> findUsages(@NotNull PsiElement targetElement, @Nullable FindUsagesOptions options) {
+    protected Collection<UsageInfo> findUsages(
+            @NotNull PsiElement targetElement,
+            @Nullable FindUsagesOptions options,
+            boolean highlightingMode
+    ) {
         Project project = getProject();
 
         FindUsagesHandler handler;
@@ -456,7 +462,14 @@ public abstract class AbstractJetFindUsagesTest extends JetLightCodeInsightFixtu
         PsiElement[] psiElements = ArrayUtil.mergeArrays(handler.getPrimaryElements(), handler.getSecondaryElements());
 
         for (PsiElement psiElement : psiElements) {
-            handler.processElementUsages(psiElement, processor, options);
+            if (highlightingMode) {
+                for (PsiReference reference : handler.findReferencesToHighlight(psiElement, options.searchScope)) {
+                    processor.process(new UsageInfo(reference));
+                }
+            }
+            else {
+                handler.processElementUsages(psiElement, processor, options);
+            }
         }
 
         return processor.getResults();

idea/tests/org/jetbrains/kotlin/findUsages/JetFindUsagesTestGenerated.java

@@ -540,6 +540,12 @@ public class JetFindUsagesTestGenerated extends AbstractJetFindUsagesTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("highlightingOfSuperUsages.0.kt")
+            public void testHighlightingOfSuperUsages() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/testData/findUsages/kotlin/findFunctionUsages/highlightingOfSuperUsages.0.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("javaAndKotlinOverrides.0.kt")
             public void testJavaAndKotlinOverrides() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/testData/findUsages/kotlin/findFunctionUsages/javaAndKotlinOverrides.0.kt");

idea/tests/org/jetbrains/kotlin/findUsages/KotlinFindUsagesWithLibraryCustomTest.kt

@@ -26,7 +26,7 @@ public class KotlinFindUsagesWithLibraryCustomTest : AbstractKotlinFindUsagesWit
         val libraryFile = FilenameIndex.getFilesByName(getProject(), "library.kt", myFixture.getModule().getModuleWithLibrariesScope()).first()
         val indexOf = libraryFile.getText().indexOf("localClassProperty")
         val jetParameter = libraryFile.findElementAt(indexOf)!!.getStrictParentOfType<KtParameter>()!!
-        val usages = findUsages(jetParameter.getOriginalElement(), null)
+        val usages = findUsages(jetParameter.getOriginalElement(), null, false)
         assertEquals(2, usages.size())
     }
 }