FIR IDE: fix containing KtFirSymbolContainingDeclarationProvider for implicit constructors

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt

@@ -7,16 +7,20 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+import org.jetbrains.kotlin.fir.FirRealSourceElementKind
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.parentOfType
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.components.KtSymbolContainingDeclarationProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithKind
 import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtObjectLiteralExpression
 import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 
@@ -54,19 +58,9 @@ internal class KtFirSymbolContainingDeclarationProvider(
 
     private fun getContainingDeclarationForKotlinInSourceSymbol(symbol: KtSymbolWithKind): KtSymbolWithKind = with(analysisSession) {
         require(symbol.origin == KtSymbolOrigin.SOURCE || symbol.origin == KtSymbolOrigin.SOURCE_MEMBER_GENERATED)
+        require(symbol is KtFirSymbol<*>)
 
-        val psi = when (val psi = symbol.getPsi()) {
+        val containingDeclaration = getContainingPsi(symbol)
-            is KtDeclaration -> psi
-            is KtObjectLiteralExpression -> psi.objectDeclaration
-            else -> error { "PSI of kotlin declaration should be KtDeclaration but was ${psi::class.simpleName}" }
-        }
-
-        val containingDeclaration = when (symbol.origin) {
-            KtSymbolOrigin.SOURCE -> psi.parentOfType()
-                ?: error("Containing declaration should present for non-toplevel declaration")
-            KtSymbolOrigin.SOURCE_MEMBER_GENERATED -> psi
-            else -> error("Unsupported declaration origin ${symbol.origin}")
-        }
 
         return with(analysisSession) {
             val containingSymbol = containingDeclaration.getSymbol()
@@ -75,16 +69,23 @@ internal class KtFirSymbolContainingDeclarationProvider(
         }
     }
 
-    private fun KtSymbolWithKind.getPsi(): PsiElement {
+    private fun getContainingPsi(symbol: KtFirSymbol<*>): KtDeclaration {
-        require(this is KtFirSymbol<*>)
+        val source = symbol.firRef.withFir(action = FirDeclaration::source)
-        return getPropertyByParameterPsi()
+        val thisSource = when (source?.kind) {
-            ?: psi
+            null -> error("PSI should present for declaration built by Kotlin code")
-            ?: error("PSI should present for declaration built by Kotlin code")
+            FirFakeSourceElementKind.ImplicitConstructor ->
-    }
+                return source.psi as KtDeclaration
+            FirFakeSourceElementKind.PropertyFromParameter -> return source.psi?.parentOfType<KtPrimaryConstructor>()!!
+            FirRealSourceElementKind -> source.psi!!
+            else -> error("Unexpected FirSourceElement: kind=${source.kind} element=${source.psi!!::class.simpleName}")
+        }
 
-    private fun KtFirSymbol<*>.getPropertyByParameterPsi() = firRef.withFir { fir ->
+        return when (symbol.origin) {
-        if (fir.source?.kind == FirFakeSourceElementKind.PropertyFromParameter) fir.psi?.parentOfType<KtPrimaryConstructor>()
+            KtSymbolOrigin.SOURCE -> thisSource.parentOfType()
-        else null
+                ?: error("Containing declaration should present for non-toplevel declaration")
+            KtSymbolOrigin.SOURCE_MEMBER_GENERATED -> thisSource as KtDeclaration
+            else -> error("Unsupported declaration origin ${symbol.origin}")
+        }
     }
 
     private fun getContainingDeclarationForLibrarySymbol(symbol: KtSymbolWithKind): KtSymbolWithKind = with(analysisSession) {
@@ -108,6 +109,10 @@ internal class KtFirSymbolContainingDeclarationProvider(
                 val fqName = symbol.callableIdIfNonLocal ?: error("fqName should not be null for non-local declaration")
                 fqName.classId
             }
+            is KtConstructorSymbol -> {
+                symbol.containingClassIdIfNonLocal
+                    ?: error("fqName should not be null for non-local declaration")
+            }
             else -> error("We should not have a ${symbol::class} from a library")
         } ?: error("outerClassId should not be null for member declaration")
         val containingClass = containingClassId.getCorrespondingToplevelClassOrObjectSymbol()