FIR IDE: fix not initialized containingClassAttr for copied declaration

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -64,23 +64,32 @@ class RawFirBuilder(
         return reference.accept(Visitor(), Unit) as FirTypeRef
     }
 
-    fun buildFunctionWithBody(function: KtNamedFunction): FirFunction<*> {
+    fun buildFunctionWithBody(function: KtNamedFunction, original: FirFunction<*>?): FirFunction<*> {
-        return buildDeclaration(function) as FirFunction<*>
+        return buildDeclaration(function, original) as FirFunction<*>
     }
 
-    fun buildSecondaryConstructor(secondaryConstructor: KtSecondaryConstructor): FirConstructor {
+    fun buildSecondaryConstructor(secondaryConstructor: KtSecondaryConstructor, original: FirConstructor?): FirConstructor {
-        return buildDeclaration(secondaryConstructor) as FirConstructor
+        return buildDeclaration(secondaryConstructor, original) as FirConstructor
     }
 
-    fun buildPropertyWithBody(property: KtProperty): FirProperty {
+    fun buildPropertyWithBody(property: KtProperty, original: FirProperty?): FirProperty {
         require(!property.isLocal) { "Should not be used to build local properties (variables)" }
-        return buildDeclaration(property) as FirProperty
+        return buildDeclaration(property, original) as FirProperty
     }
 
-    private fun buildDeclaration(declaration: KtDeclaration): FirDeclaration {
+    private fun buildDeclaration(declaration: KtDeclaration, original: FirDeclaration?): FirDeclaration {
         assert(mode ==  RawFirBuilderMode.NORMAL) { "Building FIR declarations isn't supported in stub or lazy mode mode" }
-        setupContextForPosition(declaration)
+        setupContextForPosition(declaration,)
-        return declaration.accept(Visitor(), Unit) as FirDeclaration
+        val firDeclaration = declaration.accept(Visitor(), Unit) as FirDeclaration
+        original?.let { firDeclaration.copyContainingClassAttrFrom(it) }
+        return firDeclaration
+    }
+
+    // TODO this is a (temporary) hack, instead we should properly initialize [context]
+    private fun FirDeclaration.copyContainingClassAttrFrom(from: FirDeclaration) {
+        (this as? FirCallableMemberDeclaration<*>)?.let {
+            it.containingClassAttr = (from as? FirCallableMemberDeclaration<*>)?.containingClassAttr
+        }
     }
 
     private fun setupContextForPosition(position: KtElement) {

compiler/fir/raw-fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/AbstractPartialRawFirBuilderTestCase.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.builder
 
-import junit.framework.TestCase
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.render
@@ -36,17 +35,15 @@ abstract class AbstractPartialRawFirBuilderTestCase : AbstractRawFirBuilderTestC
     private fun testFunctionPartialBuilding(filePath: String, nameToFind: String) {
         testPartialBuilding(
             filePath,
-            { file -> file.findDescendantOfType<KtNamedFunction> { it.name == nameToFind }!! },
+            { file -> file.findDescendantOfType<KtNamedFunction> { it.name == nameToFind }!! }
-            RawFirBuilder::buildFunctionWithBody
+        ) { rawFirBuilder, function -> rawFirBuilder.buildFunctionWithBody(function, original = null) }
-        )
     }
 
     private fun testPropertyPartialBuilding(filePath: String, nameToFind: String) {
         testPartialBuilding(
             filePath,
-            { file -> file.findDescendantOfType<KtProperty> { it.name == nameToFind }!! },
+            { file -> file.findDescendantOfType<KtProperty> { it.name == nameToFind }!! }
-            RawFirBuilder::buildPropertyWithBody
+        ) { rawFirBuilder, property -> rawFirBuilder.buildPropertyWithBody(property, original = null) }
-        )
     }
 
     private fun <T : KtElement> testPartialBuilding(

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForCompletion.kt

@@ -95,7 +95,7 @@ object LowLevelFirApiFacadeForCompletion {
         originalFunction: FirSimpleFunction,
         state: FirModuleResolveState
     ): FirSimpleFunction {
-        val builtFunction = firIdeProvider.buildFunctionWithBody(element)
+        val builtFunction = firIdeProvider.buildFunctionWithBody(element, originalFunction)
 
         // right now we can't resolve builtFunction header properly, as it built right in air,
         // without file, which is now required for running stages other then body resolve, so we
@@ -115,7 +115,7 @@ object LowLevelFirApiFacadeForCompletion {
         originalProperty: FirProperty,
         state: FirModuleResolveState
     ): FirProperty {
-        val builtProperty = firIdeProvider.buildPropertyWithBody(element)
+        val builtProperty = firIdeProvider.buildPropertyWithBody(element, originalProperty)
 
         val originalSetter = originalProperty.setter
         val builtSetter = builtProperty.setter

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt

@@ -81,8 +81,8 @@ internal class ReanalyzableFunctionStructureElement(
         firLazyDeclarationResolver: FirLazyDeclarationResolver,
         firIdeProvider: FirIdeProvider,
     ): ReanalyzableFunctionStructureElement {
-        val newFunction = firIdeProvider.buildFunctionWithBody(newKtDeclaration) as FirSimpleFunction
         val originalFunction = firSymbol.fir as FirSimpleFunction
+        val newFunction = firIdeProvider.buildFunctionWithBody(newKtDeclaration, originalFunction) as FirSimpleFunction
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalFunction, newFunction) {
             firLazyDeclarationResolver.lazyResolveDeclaration(
@@ -119,8 +119,8 @@ internal class ReanalyzablePropertyStructureElement(
         firLazyDeclarationResolver: FirLazyDeclarationResolver,
         firIdeProvider: FirIdeProvider,
     ): ReanalyzablePropertyStructureElement {
-        val newProperty = firIdeProvider.buildPropertyWithBody(newKtDeclaration)
         val originalProperty = firSymbol.fir
+        val newProperty = firIdeProvider.buildPropertyWithBody(newKtDeclaration, originalProperty)
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalProperty, newProperty) {
             firLazyDeclarationResolver.lazyResolveDeclaration(

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -32,7 +32,7 @@ internal object FirLazyBodiesCalculator {
     fun calculateLazyBodiesForFunction(simpleFunction: FirSimpleFunction) {
         if (simpleFunction.body !is FirLazyBlock) return
         val rawFirBuilder = createRawFirBuilder(simpleFunction)
-        val newFunction = rawFirBuilder.buildFunctionWithBody(simpleFunction.psi as KtNamedFunction) as FirSimpleFunction
+        val newFunction = rawFirBuilder.buildFunctionWithBody(simpleFunction.psi as KtNamedFunction, simpleFunction) as FirSimpleFunction
         simpleFunction.apply {
             replaceBody(newFunction.body)
             replaceContractDescription(newFunction.contractDescription)
@@ -43,7 +43,7 @@ internal object FirLazyBodiesCalculator {
         require(!secondaryConstructor.isPrimary)
         if (secondaryConstructor.body !is FirLazyBlock) return
         val rawFirBuilder = createRawFirBuilder(secondaryConstructor)
-        val newFunction = rawFirBuilder.buildSecondaryConstructor(secondaryConstructor.psi as KtSecondaryConstructor)
+        val newFunction = rawFirBuilder.buildSecondaryConstructor(secondaryConstructor.psi as KtSecondaryConstructor, secondaryConstructor)
         secondaryConstructor.apply {
             replaceBody(newFunction.body)
         }
@@ -53,7 +53,7 @@ internal object FirLazyBodiesCalculator {
         if (!needCalculatingLazyBodyForProperty(firProperty)) return
 
         val rawFirBuilder = createRawFirBuilder(firProperty)
-        val newProperty = rawFirBuilder.buildPropertyWithBody(firProperty.psi as KtProperty)
+        val newProperty = rawFirBuilder.buildPropertyWithBody(firProperty.psi as KtProperty, firProperty)
 
         firProperty.getter?.takeIf { it.body is FirLazyBlock }?.let { getter ->
             val newGetter = newProperty.getter!!

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/providers/FirIdeProvider.kt

@@ -95,12 +95,12 @@ internal class FirIdeProvider(
 
     // TODO move out of here
     // used only for completion
-    fun buildFunctionWithBody(ktNamedFunction: KtNamedFunction): FirFunction<*> {
+    fun buildFunctionWithBody(ktNamedFunction: KtNamedFunction, original: FirFunction<*>): FirFunction<*> {
-        return RawFirBuilder(session, kotlinScopeProvider).buildFunctionWithBody(ktNamedFunction)
+        return RawFirBuilder(session, kotlinScopeProvider).buildFunctionWithBody(ktNamedFunction, original)
     }
 
-    fun buildPropertyWithBody(ktNamedFunction: KtProperty): FirProperty {
+    fun buildPropertyWithBody(ktNamedFunction: KtProperty, original: FirProperty): FirProperty {
-        return RawFirBuilder(session, kotlinScopeProvider).buildPropertyWithBody(ktNamedFunction)
+        return RawFirBuilder(session, kotlinScopeProvider).buildPropertyWithBody(ktNamedFunction, original)
     }