FIR: Make resolution logic for local classes the same as for top-level

Run different phases through common transformers sequentially
It should make code more reusable and may help with avoid common bugs

compiler/fir/analysis-tests/testData/resolve/localClasses/implicitInLocalClasses.txt

@@ -3,24 +3,24 @@ FILE: implicitInLocalClasses.kt
     }
     public final fun main(): R|kotlin/Unit| {
         local final class A : R|kotlin/Any| {
-            public constructor(): R|A| {
+            public[local] constructor(): R|A| {
                 super<R|kotlin/Any|>()
             }
 
-            public final fun foo(x: R|kotlin/Int|): R|kotlin/Boolean| {
+            public[local] final fun foo(x: R|kotlin/Int|): R|kotlin/Boolean| {
                 ^foo this@R|/A|.R|/A.bar|(R|<local>/x|)
             }
 
-            public final fun bar(y: R|kotlin/Int|): R|kotlin/Boolean| {
+            public[local] final fun bar(y: R|kotlin/Int|): R|kotlin/Boolean| {
                 ^bar CMP(>, this@R|/A|.R|kotlin/Any.hashCode|().R|kotlin/Int.plus|(R|<local>/y|).R|kotlin/Int.compareTo|(Int(0)))
             }
 
-            public final val w: R|kotlin/Boolean|
+            public[local] final val w: R|kotlin/Boolean|
                 public get(): R|kotlin/Boolean| {
                     ^ this@R|/A|.R|/A.z|
                 }
 
-            public final val z: R|kotlin/Boolean|
+            public[local] final val z: R|kotlin/Boolean|
                 public get(): R|kotlin/Boolean| {
                     ^ ==(this@R|/A|.R|kotlin/Any.hashCode|(), Int(0))
                 }
@@ -33,32 +33,32 @@ FILE: implicitInLocalClasses.kt
         R|/useBoolean|(R|<local>/a|.R|/A.w|)
         R|/useBoolean|(R|<local>/a|.R|/A.z|)
         local final class B : R|kotlin/Any| {
-            public constructor(): R|B| {
+            public[local] constructor(): R|B| {
                 super<R|kotlin/Any|>()
             }
 
-            public final fun foo(x: R|kotlin/Int|): R|kotlin/Boolean| {
+            public[local] final fun foo(x: R|kotlin/Int|): R|kotlin/Boolean| {
                 ^foo this@R|/B|.R|/B.inner|.R|/B.Inner.w|
             }
 
-            public final fun bar(y: R|kotlin/Int|): R|kotlin/Boolean| {
+            public[local] final fun bar(y: R|kotlin/Int|): R|kotlin/Boolean| {
                 ^bar CMP(>, this@R|/B|.R|kotlin/Any.hashCode|().R|kotlin/Int.plus|(R|<local>/y|).R|kotlin/Int.compareTo|(Int(0)))
             }
 
-            public final val inner: R|B.Inner| = this@R|/B|.R|/B.Inner.Inner|()
+            public[local] final val inner: R|B.Inner| = this@R|/B|.R|/B.Inner.Inner|()
                 public get(): R|B.Inner|
 
             local final inner class Inner : R|kotlin/Any| {
-                public constructor(): R|B.Inner| {
+                public[local] constructor(): R|B.Inner| {
                     super<R|kotlin/Any|>()
                 }
 
-                public final val w: R|kotlin/Boolean|
+                public[local] final val w: R|kotlin/Boolean|
                     public get(): R|kotlin/Boolean| {
                         ^ this@R|/B.Inner|.R|/B.Inner.z|
                     }
 
-                public final val z: R|kotlin/Boolean|
+                public[local] final val z: R|kotlin/Boolean|
                     public get(): R|kotlin/Boolean| {
                         ^ this@R|/B|.R|/B.bar|(Int(1))
                     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt

@@ -33,7 +33,13 @@ class FirStatusResolveTransformerAdapter : FirTransformer<Nothing?>() {
     }
 }
 
-class FirStatusResolveTransformer(override val session: FirSession) :
+fun <F : FirClass<F>> F.runStatusResolveForLocalClass(session: FirSession): F {
+    val transformer = FirStatusResolveTransformer(session)
+
+    return this.transform<F, Nothing?>(transformer, null).single
+}
+
+private class FirStatusResolveTransformer(override val session: FirSession) :
     FirAbstractTreeTransformer<FirDeclarationStatus?>(phase = FirResolvePhase.STATUS) {
     private val classes = mutableListOf<FirClass<*>>()
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.LocalClassesNavigationInfo
 import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.createImportingScopes
@@ -42,6 +43,26 @@ class FirSupertypeResolverTransformer(
     }
 }
 
+fun <F : FirClass<F>> F.runSupertypeResolvePhaseForLocalClass(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    currentScopeList: List<FirScope>,
+    localClassesNavigationInfo: LocalClassesNavigationInfo,
+): F {
+    val supertypeComputationSession = SupertypeComputationSession()
+    val applySupertypesTransformer = FirApplySupertypesTransformer(supertypeComputationSession)
+    val supertypeResolverVisitor = FirSupertypeResolverVisitor(
+        session, supertypeComputationSession, scopeSession,
+        FirImmutableCompositeScope(ImmutableList.ofAll(currentScopeList)),
+        localClassesNavigationInfo
+    )
+
+    this.accept(supertypeResolverVisitor)
+    supertypeComputationSession.breakLoops(session)
+
+    return this.transform<F, Nothing?>(applySupertypesTransformer, null).single
+}
+
 private class FirApplySupertypesTransformer(
     private val supertypeComputationSession: SupertypeComputationSession
 ) : FirDefaultTransformer<Nothing?>() {
@@ -125,7 +146,9 @@ fun FirRegularClass.resolveSupertypesInTheAir(session: FirSession): List<FirType
 private class FirSupertypeResolverVisitor(
     private val session: FirSession,
     private val supertypeComputationSession: SupertypeComputationSession,
-    private val scopeSession: ScopeSession
+    private val scopeSession: ScopeSession,
+    private val scopeForLocalClass: FirImmutableCompositeScope? = null,
+    private val localClassesNavigationInfo: LocalClassesNavigationInfo? = null
 ) : FirDefaultVisitorVoid() {
     override fun visitElement(element: FirElement) {}
 
@@ -170,12 +193,28 @@ private class FirSupertypeResolverVisitor(
 
     private fun prepareScope(classLikeDeclaration: FirClassLikeDeclaration<*>): FirImmutableCompositeScope {
         val classId = classLikeDeclaration.symbol.classId
-        val outerClassFir = classId.outerClassId?.let(session.firProvider::getFirClassifierByFqName) as? FirRegularClass
 
-        val result = if (outerClassFir == null) {
+        val result = when {
-            prepareFileScope(session.firProvider.getFirClassifierContainerFile(classId))
+            classId.isLocal -> {
-        } else {
+                // Local type aliases are not supported
-            prepareScopeForNestedClasses(outerClassFir)
+                if (classLikeDeclaration !is FirClass<*>) return FirImmutableCompositeScope.EMPTY
+
+                // Local classes should be treated specially and supplied with localClassesNavigationInfo, normally
+                // But it seems to be too strict to add an assertion here
+                val navigationInfo = localClassesNavigationInfo ?: return FirImmutableCompositeScope.EMPTY
+
+                val parent = localClassesNavigationInfo.parentForClass[classLikeDeclaration]
+
+                when {
+                    parent != null -> prepareScopeForNestedClasses(parent)
+                    else -> scopeForLocalClass ?: return FirImmutableCompositeScope.EMPTY
+                }
+            }
+            classId.isNestedClass -> {
+                val outerClassFir = classId.outerClassId?.let(session.firProvider::getFirClassifierByFqName) as? FirRegularClass
+                prepareScopeForNestedClasses(outerClassFir ?: return FirImmutableCompositeScope.EMPTY)
+            }
+            else -> prepareFileScope(session.firProvider.getFirClassifierContainerFile(classId))
         }
 
         return result.childScope(classLikeDeclaration.typeParametersScope())
@@ -369,7 +408,7 @@ private class SupertypeComputationSession {
 
 fun FirTypeRef.firClassLike(session: FirSession): FirClassLikeDeclaration<*>? {
     val type = coneTypeSafe<ConeClassLikeType>() ?: return null
-    return session.firProvider.getFirClassifierByFqName(type.lookupTag.classId)
+    return type.lookupTag.toSymbol(session)?.fir
 }
 
 sealed class SupertypeComputationStatus {
@@ -384,6 +423,11 @@ private typealias ImmutableList<E> = javaslang.collection.List<E>
 private class FirImmutableCompositeScope(
     override val scopes: ImmutableList<FirScope>
 ) : FirIterableScope() {
+    
+    companion object {
+        val EMPTY = FirImmutableCompositeScope(ImmutableList.empty())
+    }
+    
     fun childScope(newScope: FirScope?) = newScope?.let { FirImmutableCompositeScope(scopes.push(newScope)) } ?: this
     fun childScope(newScopes: Collection<FirScope>) = FirImmutableCompositeScope(scopes.pushAll(newScopes))
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirTypeResolveTransformer.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.createImportingScopes
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -33,14 +34,29 @@ class FirTypeResolveTransformerAdapter(
     }
 }
 
-class FirTypeResolveTransformer(
+fun <F : FirClass<F>> F.runTypeResolvePhaseForLocalClass(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    currentScopeList: List<FirScope>,
+): F {
+    val transformer = FirTypeResolveTransformer(scopeSession, session, currentScopeList)
+
+    return this.transform<F, Nothing?>(transformer, null).single
+}
+
+private class FirTypeResolveTransformer(
     private val scopeSession: ScopeSession,
     override val session: FirSession,
+    initialScopes: List<FirScope> = emptyList()
 ) : FirAbstractTreeTransformerWithSuperTypes(
     phase = FirResolvePhase.TYPES,
     reversedScopePriority = true
 ) {
 
+    init {
+        towerScope.addScopes(initialScopes.asReversed())
+    }
+
     private val typeResolverTransformer: FirSpecificTypeResolverTransformer = FirSpecificTypeResolverTransformer(towerScope, session)
 
     override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirFile> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -270,43 +270,6 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
         )
     }
 
-    private fun prepareLocalClassForBodyResolve(klass: FirClass<*>) {
-        if (klass.superTypeRefs.any { it !is FirResolvedTypeRef }) {
-            klass.replaceSuperTypeRefs(
-                klass.superTypeRefs.map { superTypeRef ->
-                    this.transformer.transformTypeRef(superTypeRef, ResolutionMode.ContextIndependent).single
-                }
-            )
-        }
-        if (klass is FirRegularClass) {
-            klass.transformStatus(transformer, klass.resolveStatus(klass.status).mode())
-        }
-        // This is necessary because of possible jumps from implicit bodies inside
-        for (declaration in klass.declarations) {
-            when (declaration) {
-                is FirRegularClass -> {
-                    prepareLocalClassForBodyResolve(declaration)
-                }
-                is FirCallableMemberDeclaration<*> -> {
-                    withTypeParametersOf(declaration) {
-                        prepareSignatureForBodyResolve(declaration)
-                        declaration.transformStatus(transformer, declaration.resolveStatus(declaration.status, klass).mode())
-                        if (declaration is FirProperty) {
-                            declaration.getter?.let { it.transformStatus(this, it.resolveStatus(it.status, klass).mode()) }
-                            declaration.setter?.let { it.transformStatus(this, it.resolveStatus(it.status, klass).mode()) }
-                        }
-                    }
-                }
-            }
-            declaration.replaceResolvePhase(FirResolvePhase.STATUS)
-        }
-        if (klass is FirRegularClass) {
-            for (typeParameter in klass.typeParameters) {
-                typeParameter.replaceResolvePhase(FirResolvePhase.STATUS)
-            }
-        }
-    }
-
     override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): CompositeTransformResult<FirStatement> {
         if (regularClass.symbol.classId.isLocal && regularClass !in context.targetedLocalClasses) {
             return regularClass.runAllPhasesForLocalClass(components, data).also {
@@ -316,10 +279,6 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
 
         context.storeClass(regularClass)
         return withTypeParametersOf(regularClass) {
-            if (regularClass.symbol.classId.isLocal) {
-                prepareLocalClassForBodyResolve(regularClass)
-            }
-
             val oldConstructorScope = primaryConstructorParametersScope
             val oldContainingClass = containingClass
             primaryConstructorParametersScope = null
@@ -351,7 +310,6 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             return anonymousObject.runAllPhasesForLocalClass(components, data).compose()
         }
 
-        prepareLocalClassForBodyResolve(anonymousObject)
         val type = anonymousObject.defaultType()
         anonymousObject.resultType = buildResolvedTypeRef {
             source = anonymousObject.source

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/LocalClassesResolution.kt

@@ -7,12 +7,28 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.transformers.runStatusResolveForLocalClass
+import org.jetbrains.kotlin.fir.resolve.transformers.runSupertypeResolvePhaseForLocalClass
+import org.jetbrains.kotlin.fir.resolve.transformers.runTypeResolvePhaseForLocalClass
+import org.jetbrains.kotlin.fir.scopes.impl.createCurrentScopeList
 
 fun <F : FirClass<F>> F.runAllPhasesForLocalClass(
     components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
     resolutionMode: ResolutionMode
 ): F {
     val localClassesNavigationInfo = collectLocalClassesNavigationInfo()
+    runSupertypeResolvePhaseForLocalClass(
+        components.session,
+        components.scopeSession,
+        components.createCurrentScopeList(),
+        localClassesNavigationInfo
+    )
+    runTypeResolvePhaseForLocalClass(
+        components.session,
+        components.scopeSession,
+        components.createCurrentScopeList()
+    )
+    runStatusResolveForLocalClass(components.session)
     runBodiesResolutionForLocalClass(components, resolutionMode, localClassesNavigationInfo)
     return this
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirTypeResolveScopeForBodyResolve.kt

@@ -14,13 +14,16 @@ class FirTypeResolveScopeForBodyResolve(
     private val components: BodyResolveComponents
 ) : FirIterableScope() {
     override val scopes: Iterable<FirScope>
-        get() = mutableListOf<FirScope>().apply {
+        get() = components.createCurrentScopeList()
-            addAll(components.localScopes.asReversed())
-            components.implicitReceiverStack.receiversAsReversed().mapNotNullTo(this) {
-                (it as? ImplicitDispatchReceiverValue)?.implicitScope
-            }
-
-            addAll(components.typeParametersScopes.asReversed())
-            addAll(components.fileImportsScope.asReversed())
-        }
 }
+
+fun BodyResolveComponents.createCurrentScopeList(): List<FirScope> =
+    mutableListOf<FirScope>().apply {
+        addAll(localScopes.asReversed())
+        implicitReceiverStack.receiversAsReversed().mapNotNullTo(this) {
+            (it as? ImplicitDispatchReceiverValue)?.implicitScope
+        }
+
+        addAll(typeParametersScopes.asReversed())
+        addAll(fileImportsScope.asReversed())
+    }

compiler/testData/diagnostics/tests/dataFlow/local/LocalClassInMemberOfLocalClass.fir.kt

@@ -6,8 +6,8 @@ fun test(x: Any) {
       if (y !is String) return
       class Local {
         init {
-          x.<!UNRESOLVED_REFERENCE!>length<!>
+          x.length
-          y.<!UNRESOLVED_REFERENCE!>length<!>
+          y.length
         }
       }
     }

compiler/testData/diagnostics/tests/dataFlow/local/LocalClassInitializer.fir.kt

@@ -4,7 +4,7 @@ fun f(a: Any?) {
   if (a is B) {
     class C : X(a) {
       init {
-        a.<!UNRESOLVED_REFERENCE!>foo<!>()
+        a.foo()
       }
     }
   }

compiler/testData/diagnostics/tests/dataFlow/local/NestedLocalClass.fir.kt

@@ -4,7 +4,7 @@ fun test(x: Any) {
   class LocalOuter {
     inner class Local {
       init {
-        x.<!UNRESOLVED_REFERENCE!>length<!>
+        x.length
       }
     }
   }

compiler/testData/diagnostics/tests/regressions/itselfAsUpperBoundLocal.fir.kt

@@ -1,5 +1,5 @@
 // !WITH_NEW_INFERENCE
 fun bar() {
     fun <T: T?> foo() {}
-    foo()
+    <!INAPPLICABLE_CANDIDATE!>foo<!>()
 }