From 767af0dae0b4d97cfbd5dbd21b2e0cd63fb46e9f Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Thu, 29 Jul 2021 13:55:32 +0200
Subject: [PATCH] HL API: fix super declarations search for local classes

---
 .../fir/symbols/KtFirOverrideInfoProvider.kt   | 18 +++++++++---------
 .../overridenDeclarations/inLocalClass.kt      |  9 +++++++++
 .../overridenDeclarations/inLocalClass.txt     |  4 ++++
 ...riddenDeclarationProviderTestGenerated.java |  6 ++++++
 4 files changed, 28 insertions(+), 9 deletions(-)
 create mode 100644 idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.kt
 create mode 100644 idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.txt

diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt
index da3f8e40ac2..6ef1807818a 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirOverrideInfoProvider.kt
@@ -15,22 +15,22 @@ import org.jetbrains.kotlin.fir.originalForIntersectionOverrideAttr
 import org.jetbrains.kotlin.fir.originalForSubstitutionOverride
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
+import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.components.KtOverrideInfoProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.components.KtFirAnalysisSessionComponent
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.util.ImplementationStatus
 
-class KtFirOverrideInfoProvider(
-    override val analysisSession: KtAnalysisSession,
+internal class KtFirOverrideInfoProvider(
+    override val analysisSession: KtFirAnalysisSession,
     override val token: ValidityToken,
-) : KtOverrideInfoProvider() {
-
-    private val firAnalysisSession = analysisSession as KtFirAnalysisSession
+) : KtOverrideInfoProvider(), KtFirAnalysisSessionComponent {
 
     override fun isVisible(memberSymbol: KtCallableSymbol, classSymbol: KtClassOrObjectSymbol): Boolean {
         require(memberSymbol is KtFirSymbol<*>)
@@ -60,7 +60,7 @@ class KtFirOverrideInfoProvider(
                 if (parentClassFir !is FirClass) return@inner null
 
                 memberFir.symbol.getImplementationStatus(
-                    SessionHolderImpl(firAnalysisSession.rootModuleSession, ScopeSession()),
+                    SessionHolderImpl(rootModuleSession, ScopeSession()),
                     parentClassFir.symbol
                 )
             }
@@ -71,9 +71,9 @@ class KtFirOverrideInfoProvider(
         require(symbol is KtFirSymbol<*>)
         return symbol.firRef.withFir(FirResolvePhase.STATUS) { firDeclaration ->
             if (firDeclaration !is FirCallableDeclaration) return@withFir null
-            with(analysisSession) {
-                getOriginalOverriddenSymbol(firDeclaration)?.containingClass()?.classId?.getCorrespondingToplevelClassOrObjectSymbol()
-            }
+            val containingClass =
+                getOriginalOverriddenSymbol(firDeclaration)?.containingClass()?.toSymbol(rootModuleSession) ?: return@withFir null
+            analysisSession.firSymbolBuilder.classifierBuilder.buildClassLikeSymbol(containingClass.fir) as? KtClassOrObjectSymbol
         }
     }
 
diff --git a/idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.kt b/idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.kt
new file mode 100644
index 00000000000..e8dde0fdd15
--- /dev/null
+++ b/idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.kt
@@ -0,0 +1,9 @@
+fun foo() {
+    open class A {
+        open fun foo() {}
+    }
+
+    class B : A() {
+        override fun fo<caret>o() {}
+    }
+}
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.txt b/idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.txt
new file mode 100644
index 00000000000..064868264e5
--- /dev/null
+++ b/idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.txt
@@ -0,0 +1,4 @@
+ALL:
+  foo.A.foo(): Unit
+DIRECT:
+  foo.A.foo(): Unit
diff --git a/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/OverriddenDeclarationProviderTestGenerated.java b/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/OverriddenDeclarationProviderTestGenerated.java
index 863ddd71b75..3c0e1080299 100644
--- a/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/OverriddenDeclarationProviderTestGenerated.java
+++ b/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/OverriddenDeclarationProviderTestGenerated.java
@@ -24,6 +24,12 @@ public class OverriddenDeclarationProviderTestGenerated extends AbstractOverridd
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/overridenDeclarations"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("inLocalClass.kt")
+    public void testInLocalClass() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/overridenDeclarations/inLocalClass.kt");
+    }
+
     @Test
     @TestMetadata("inOtherFile.kt")
     public void testInOtherFile() throws Exception {