[FIR2IR] Get rid of all usages of IrSymbol.owner from Fir2IrVisitor

^KT-60924

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/ConversionUtils.kt

@@ -716,7 +716,7 @@ fun IrActualizedResult?.extractFirDeclarations(): Set<FirDeclaration>? {
 // This method is intended to be used for default values of annotation parameters (compile-time strings, numbers, enum values, KClasses)
 // where they are needed and may produce incorrect results for values that may be encountered outside annotations.
 fun FirExpression.asCompileTimeIrInitializer(components: Fir2IrComponents): IrExpressionBody? {
-    return when (val elem = this.accept(Fir2IrVisitor(components, Fir2IrConversionScope()), null)) {
+    return when (val elem = this.accept(Fir2IrVisitor(components, Fir2IrConversionScope(components.configuration)), null)) {
         is IrExpressionBody -> elem
         is IrExpression -> components.irFactory.createExpressionBody(elem)
         else -> null

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConversionScope.kt

@@ -11,15 +11,18 @@ import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
+import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbolInternals
 import org.jetbrains.kotlin.ir.util.isSetter
 import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.util.PrivateForInline
 
 @OptIn(PrivateForInline::class)
-class Fir2IrConversionScope {
+class Fir2IrConversionScope(val configuration: Fir2IrConfiguration) {
     @PublishedApi
     @PrivateForInline
     internal val parentStack = mutableListOf<IrDeclarationParent>()
@@ -30,7 +33,7 @@ class Fir2IrConversionScope {
 
     @PublishedApi
     @PrivateForInline
-    internal val currentlyGeneratedDelegatedConstructors = mutableMapOf<IrClass, IrConstructor>()
+    internal val currentlyGeneratedDelegatedConstructors = mutableMapOf<IrClassSymbol, IrConstructor>()
 
     inline fun <T : IrDeclarationParent, R> withParent(parent: T, f: T.() -> R): R {
         parentStack += parent
@@ -42,16 +45,16 @@ class Fir2IrConversionScope {
     }
 
     internal fun <T> forDelegatingConstructorCall(constructor: IrConstructor, irClass: IrClass, f: () -> T): T {
-        currentlyGeneratedDelegatedConstructors[irClass] = constructor
+        currentlyGeneratedDelegatedConstructors[irClass.symbol] = constructor
         try {
             return f()
         } finally {
-            currentlyGeneratedDelegatedConstructors.remove(irClass)
+            currentlyGeneratedDelegatedConstructors.remove(irClass.symbol)
         }
     }
 
-    fun getConstructorForCurrentlyGeneratedDelegatedConstructor(itClass: IrClass): IrConstructor? =
-        currentlyGeneratedDelegatedConstructors[itClass]
+    fun getConstructorForCurrentlyGeneratedDelegatedConstructor(itClassSymbol: IrClassSymbol): IrConstructor? =
+        currentlyGeneratedDelegatedConstructors[itClassSymbol]
 
     fun containingFileIfAny(): IrFile? = parentStack.getOrNull(0) as? IrFile
 
@@ -81,6 +84,27 @@ class Fir2IrConversionScope {
         error("Accessor of property ${property.render()} not found on parent stack")
     }
 
+    @OptIn(IrSymbolInternals::class)
+    @Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
+    inline fun <reified D : IrDeclaration> findDeclarationInParentsStack(symbol: IrSymbol): @kotlin.internal.NoInfer D {
+        if (!AbstractTypeChecker.RUN_SLOW_ASSERTIONS) {
+            return symbol.owner as D
+        }
+        for (parent in parentStack.asReversed()) {
+            if ((parent as? IrDeclaration)?.symbol == symbol) {
+                return parent as D
+            }
+        }
+        /*
+         * In case of IDE (when allowNonCachedDeclarations is set to true) we may be in scope of some already compiled class,
+         *   for which we have Fir2IrLazyClass in symbol
+         */
+        if (configuration.allowNonCachedDeclarations) {
+            return symbol.owner as D
+        }
+        error("Declaration with symbol $symbol is not found in parents stack")
+    }
+
     fun <T : IrDeclaration> applyParentFromStackTo(declaration: T): T {
         declaration.parent = parentStack.last()
         return declaration

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt

@@ -568,7 +568,7 @@ class Fir2IrConverter(
                 components, fir2IrConfiguration.languageVersionSettings, moduleDescriptor, irMangler
             )
             components.irBuiltIns = irBuiltIns
-            val conversionScope = Fir2IrConversionScope()
+            val conversionScope = Fir2IrConversionScope(components.configuration)
             val fir2irVisitor = Fir2IrVisitor(components, conversionScope)
             components.builtIns = Fir2IrBuiltIns(components, specialSymbolProvider)
             components.annotationGenerator = AnnotationGenerator(components)

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.contracts.description.LogicOperationKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.descriptors.isObject
 import org.jetbrains.kotlin.diagnostics.findChildByType
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.backend.generators.ClassMemberGenerator
@@ -28,6 +29,7 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.references.*
+import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedConeType
 import org.jetbrains.kotlin.fir.resolve.isIteratorNext
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -44,15 +46,12 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
-import org.jetbrains.kotlin.ir.symbols.IrSymbolInternals
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrErrorClassImpl
 import org.jetbrains.kotlin.ir.types.impl.IrErrorTypeImpl
 import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.defaultConstructor
-import org.jetbrains.kotlin.ir.util.defaultType
-import org.jetbrains.kotlin.ir.util.isObject
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
@@ -652,14 +651,20 @@ class Fir2IrVisitor(
     }
 
     // Note that this mimics psi2ir [StatementGenerator#shouldGenerateReceiverAsSingletonReference].
-    private fun shouldGenerateReceiverAsSingletonReference(irClass: IrClass): Boolean {
+    private fun shouldGenerateReceiverAsSingletonReference(irClassSymbol: IrClassSymbol): Boolean {
         val scopeOwner = conversionScope.parent()
-        return irClass.isObject &&
-                scopeOwner != irClass && // For anonymous initializers
-                !((scopeOwner is IrFunction || scopeOwner is IrProperty || scopeOwner is IrField) && (scopeOwner as IrDeclaration).parent == irClass) // Members of object
+        // For anonymous initializers
+        if ((scopeOwner as? IrDeclaration)?.symbol == irClassSymbol) return false
+        // Members of object
+        return when (scopeOwner) {
+            is IrFunction, is IrProperty, is IrField -> {
+                val parent = (scopeOwner as IrDeclaration).parent as? IrDeclaration
+                parent?.symbol != irClassSymbol
+            }
+            else -> true
+        }
     }
 
-    @OptIn(IrSymbolInternals::class)
     override fun visitThisReceiverExpression(
         thisReceiverExpression: FirThisReceiverExpression,
         data: Any?
@@ -676,21 +681,40 @@ class Fir2IrVisitor(
             is FirClassSymbol -> {
                 // Object case
                 val firClass = boundSymbol.fir as FirClass
-                val irClass = if (firClass.origin == FirDeclarationOrigin.Source) {
+                val irClassSymbol = if (firClass.origin.fromSource || firClass.origin.generated) {
                     // We anyway can use 'else' branch as fallback, but
                     // this is an additional check of FIR2IR invariants
                     // (source classes should be already built when we analyze bodies)
-                    classifierStorage.getCachedIrClass(firClass)!!
+                    classifierStorage.getCachedIrClass(firClass)!!.symbol
                 } else {
-                    classifierStorage.getIrClassSymbol(boundSymbol).owner
+                    /*
+                     * The only case when we can refer to non-source this is resolution to companion object of parent
+                     *   class in some constructor scope:
+                     *
+                     * // MODULE: lib
+                     * abstract class Base {
+                     *     companion object {
+                     *         fun foo(): Int = 1
+                     *     }
+                     * }
+                     *
+                     * // MODULE: app(lib)
+                     * class Derived(
+                     *     val x: Int = foo() // this: Base.Companion
+                     * ) : Base()
+                     */
+                    classifierStorage.getIrClassSymbol(boundSymbol)
                 }
 
-                if (shouldGenerateReceiverAsSingletonReference(irClass)) {
+                if (firClass.classKind.isObject && shouldGenerateReceiverAsSingletonReference(irClassSymbol)) {
                     return thisReceiverExpression.convertWithOffsets { startOffset, endOffset ->
-                        IrGetObjectValueImpl(startOffset, endOffset, irClass.defaultType, irClass.symbol)
+                        val irType = boundSymbol.defaultType().toIrType()
+                        IrGetObjectValueImpl(startOffset, endOffset, irType, irClassSymbol)
                     }
                 }
 
+                val irClass = conversionScope.findDeclarationInParentsStack<IrClass>(irClassSymbol)
+
                 val dispatchReceiver = conversionScope.dispatchReceiverParameter(irClass)
                 if (dispatchReceiver != null) {
                     return thisReceiverExpression.convertWithOffsets { startOffset, endOffset ->
@@ -699,7 +723,7 @@ class Fir2IrVisitor(
                         } ?: IrGetValueImpl(startOffset, endOffset, dispatchReceiver.type, dispatchReceiver.symbol)
                         if (calleeReference.contextReceiverNumber != -1) {
                             val constructorForCurrentlyGeneratedDelegatedConstructor =
-                                conversionScope.getConstructorForCurrentlyGeneratedDelegatedConstructor(irClass)
+                                conversionScope.getConstructorForCurrentlyGeneratedDelegatedConstructor(irClass.symbol)
 
                             if (constructorForCurrentlyGeneratedDelegatedConstructor != null) {
                                 val constructorParameter =
@@ -740,7 +764,10 @@ class Fir2IrVisitor(
             }
             is FirCallableSymbol -> {
                 val irFunction = when (boundSymbol) {
-                    is FirFunctionSymbol -> declarationStorage.getIrFunctionSymbol(boundSymbol).owner
+                    is FirFunctionSymbol -> {
+                        val functionSymbol = declarationStorage.getIrFunctionSymbol(boundSymbol)
+                        conversionScope.findDeclarationInParentsStack<IrSimpleFunction>(functionSymbol)
+                    }
                     is FirPropertySymbol -> {
                         val property = declarationStorage.getIrPropertySymbol(boundSymbol) as? IrPropertySymbol
                         property?.let { conversionScope.parentAccessorOfPropertyFromStack(it) }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/FirIrProvider.kt

@@ -26,7 +26,7 @@ class FirIrProvider(val fir2IrComponents: Fir2IrComponents) : IrProvider {
     private val symbolProvider = fir2IrComponents.session.symbolProvider
     private val declarationStorage = fir2IrComponents.declarationStorage
     private val classifierStorage = fir2IrComponents.classifierStorage
-    private val fakeOverrideGenerator = FakeOverrideGenerator(fir2IrComponents, Fir2IrConversionScope())
+    private val fakeOverrideGenerator = FakeOverrideGenerator(fir2IrComponents, Fir2IrConversionScope(components.configuration))
 
     override fun getDeclaration(symbol: IrSymbol): IrDeclaration? {
         val signature = symbol.signature ?: return null