From 95744062c12b4f74b6c0894dca5a57143f91c043 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Sat, 9 Jul 2022 17:57:12 +0200
Subject: [PATCH] [Symbol LC] add parents for the local classes

---
 .../symbol/classes/FirLightAnonymousClassForSymbol.kt  |  4 ++++
 .../classes/FirLightClassForClassOrObjectSymbol.kt     | 10 +++++++++-
 .../classes/ultraLightClassForLocalDeclaration.kt      |  2 +-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightAnonymousClassForSymbol.kt b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightAnonymousClassForSymbol.kt
index 84d96cc90a6..54a7e5aec06 100644
--- a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightAnonymousClassForSymbol.kt
+++ b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightAnonymousClassForSymbol.kt
@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.asJava.elements.KtLightIdentifier
 import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.analysis.api.symbols.KtAnonymousObjectSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
+import org.jetbrains.kotlin.asJava.classes.getParentForLocalDeclaration
 import org.jetbrains.kotlin.light.classes.symbol.classes.*
 import org.jetbrains.kotlin.load.java.structure.LightClassOriginKind
 import org.jetbrains.kotlin.psi.KtClassOrObject
@@ -89,6 +90,9 @@ internal class FirLightAnonymousClassForSymbol(
 
     override fun getOwnInnerClasses(): List<PsiClass> = _ownInnerClasses
 
+    override fun getParent(): PsiElement? =
+        kotlinOrigin?.let(::getParentForLocalDeclaration)
+
     override fun getScope(): PsiElement? = parent
     override fun getInterfaces(): Array<PsiClass> = PsiClassImplUtil.getInterfaces(this)
     override fun getSuperClass(): PsiClass? = PsiClassImplUtil.getSuperClass(this)
diff --git a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightClassForClassOrObjectSymbol.kt b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightClassForClassOrObjectSymbol.kt
index 5c632173421..dc5abbb5550 100644
--- a/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightClassForClassOrObjectSymbol.kt
+++ b/analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/FirLightClassForClassOrObjectSymbol.kt
@@ -15,7 +15,9 @@ import org.jetbrains.annotations.NonNls
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
 import org.jetbrains.kotlin.asJava.classes.getOutermostClassOrObject
+import org.jetbrains.kotlin.asJava.classes.getParentForLocalDeclaration
 import org.jetbrains.kotlin.asJava.classes.lazyPub
 import org.jetbrains.kotlin.asJava.elements.KtLightField
 import org.jetbrains.kotlin.asJava.elements.KtLightMethod
@@ -31,6 +33,7 @@ import org.jetbrains.kotlin.psi.stubs.KotlinClassOrObjectStub
 import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 import org.jetbrains.kotlin.utils.addToStdlib.ifFalse
 import org.jetbrains.kotlin.utils.addToStdlib.ifTrue
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 context(KtAnalysisSession)
 internal abstract class FirLightClassForClassOrObjectSymbol(
@@ -204,7 +207,12 @@ internal abstract class FirLightClassForClassOrObjectSymbol(
         return null
     }
 
-    override fun getParent(): PsiElement? = containingClass ?: containingFile
+    override fun getParent(): PsiElement? {
+        if (classOrObjectSymbol.safeAs<KtSymbolWithKind>()?.symbolKind == KtSymbolKind.LOCAL) {
+            return kotlinOrigin?.let(::getParentForLocalDeclaration)
+        }
+        return containingClass ?: containingFile
+    }
 
     override fun getScope(): PsiElement? = parent
 
diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightClassForLocalDeclaration.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightClassForLocalDeclaration.kt
index fbb96e1fcb6..d4e26ff336a 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightClassForLocalDeclaration.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightClassForLocalDeclaration.kt
@@ -30,7 +30,7 @@ open class KtUltraLightClassForLocalDeclaration(
     override fun getParent() = _parent
 }
 
-private fun getParentForLocalDeclaration(classOrObject: KtClassOrObject): PsiElement? {
+fun getParentForLocalDeclaration(classOrObject: KtClassOrObject): PsiElement? {
     fun getParentByPsiMethod(method: PsiMethod?, name: String?, forceMethodWrapping: Boolean): PsiElement? {
         if (method == null || name == null) return null