[FIR] Change API for generating nested classes from plugins

`FirDeclarationGenerationExtension.generateClassLikeDeclaration` was split
  into two functions: one for generating top level classes, and one for
  nested classes. Such change reduces verbosity and error-proness of
  this extension and also allows to smoothly run plugins on local classes

^KT-55248 Fixed

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.FirLazyValue
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
-import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassDeclaredMemberScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
@@ -37,7 +36,9 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
      *
      * If classId has `outerClassId.Companion` format then generated class should be a companion object
      */
-    open fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
+    open fun generateTopLevelClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
+
+    open fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? = null
 
     // Can be called on STATUS stage
     open fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> = emptyList()

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirExtensionDeclarationsSymbolProvider.kt

@@ -5,16 +5,15 @@
 
 package org.jetbrains.kotlin.fir.extensions
 
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.caches.*
 import org.jetbrains.kotlin.fir.declarations.validate
 import org.jetbrains.kotlin.fir.ownerGenerator
-import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
-import org.jetbrains.kotlin.fir.scopes.impl.groupExtensionsByName
+import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
+import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
@@ -59,17 +58,6 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
             extensions.flatGroupBy { it.topLevelClassIdsCache.getValue() }
         }
 
-    private val extensionsByNestedClassifierClassId: FirCache<ClassId, Map<ClassId, List<FirDeclarationGenerationExtension>>, Nothing?> =
-        session.firCachesFactory.createCache cache@{ outerClassId, _ ->
-            val outerClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(outerClassId) as? FirClassSymbol<*>
-                ?: return@cache emptyMap()
-            session.groupExtensionsByName(
-                outerClassSymbol.fir,
-                nameExtractor = { nestedClassifierNamesCache.getValue(outerClassSymbol) },
-                nameTransformer = { outerClassId.createNestedClassId(it) }
-            )
-        }
-
     private val extensionsByTopLevelCallableId: FirLazyValue<Map<CallableId, List<FirDeclarationGenerationExtension>>, Nothing?> =
         session.firCachesFactory.createLazyValue {
             extensions.flatGroupBy { it.topLevelCallableIdsCache.getValue() }
@@ -78,21 +66,34 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
     // ------------------------------------------ generators ------------------------------------------
 
     private fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        val matchedExtensions = when {
-            classId.isNestedClass -> extensionsByNestedClassifierClassId.getValue(classId.outerClassId!!)[classId]
-            else -> extensionsByTopLevelClassId.getValue()[classId]
-        } ?: return null
-        val generatedClasses = matchedExtensions
-            .mapNotNull { generatorExtension ->
-                generatorExtension.generateClassLikeDeclaration(classId)?.also { symbol ->
-                    symbol.fir.ownerGenerator = generatorExtension
+        return when {
+            classId.isLocal -> null
+            classId.isNestedClass -> {
+                val owner = session.symbolProvider.getClassLikeSymbolByClassId(classId.outerClassId!!) as? FirClassSymbol<*> ?: return null
+                val nestedClassifierScope = session.nestedClassifierScope(owner.fir) ?: return null
+                var result: FirClassLikeSymbol<*>? = null
+                nestedClassifierScope.processClassifiersByName(classId.shortClassName) {
+                    if (it is FirClassLikeSymbol<*>) {
+                        result = it
+                    }
+                }
+                result
+            }
+            else -> {
+                val matchedExtensions = extensionsByTopLevelClassId.getValue()[classId] ?: return null
+                val generatedClasses = matchedExtensions
+                    .mapNotNull { generatorExtension ->
+                        generatorExtension.generateTopLevelClassLikeDeclaration(classId)?.also { symbol ->
+                            symbol.fir.ownerGenerator = generatorExtension
+                        }
+                    }
+                    .onEach { it.fir.validate() }
+                when (generatedClasses.size) {
+                    0 -> null
+                    1 -> generatedClasses.first()
+                    else -> error("Multiple plugins generated classes with same classId $classId\n${generatedClasses.joinToString("\n") { it.fir.render() }}")
                 }
             }
-            .onEach { it.fir.validate() }
-        return when (generatedClasses.size) {
-            0 -> null
-            1 -> generatedClasses.first()
-            else -> error("Multiple plugins generated classes with same classId $classId\n${generatedClasses.joinToString("\n") { it.fir.render() }}")
         }
     }
 

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirSwitchableExtensionDeclarationsSymbolProvider.kt

@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.fir.extensions
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -24,7 +21,7 @@ import org.jetbrains.kotlin.name.Name
  * This is also legal, because plugins can not generate annotation classes which can influence other plugins or this plugin itself
  */
 class FirSwitchableExtensionDeclarationsSymbolProvider private constructor(
-    private val delegate: FirSymbolProvider
+    private val delegate: FirExtensionDeclarationsSymbolProvider
 ) : FirSymbolProvider(delegate.session) {
     companion object {
         fun create(session: FirSession): FirSwitchableExtensionDeclarationsSymbolProvider? {

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirGeneratedClassDeclaredMemberScope.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.declarations.validate
 import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.ownerGenerator
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
@@ -160,16 +161,18 @@ internal fun FirSession.getExtensionsForClass(klass: FirClass): List<FirDeclarat
 class FirGeneratedClassNestedClassifierScope private constructor(
     useSiteSession: FirSession,
     klass: FirClass,
-    private val nestedClassifierNames: Set<Name>
+    private val extensionsByName: Map<Name, List<FirDeclarationGenerationExtension>>,
 ) : FirNestedClassifierScope(klass, useSiteSession) {
     companion object {
         @OptIn(FirExtensionApiInternals::class)
         fun create(useSiteSession: FirSession, klass: FirClass): FirGeneratedClassNestedClassifierScope? {
-            val extensions = useSiteSession.getExtensionsForClass(klass)
             val symbol = klass.symbol
-            val classifierNames = extensions.flatMapTo(mutableSetOf()) { it.nestedClassifierNamesCache.getValue(symbol) }
-            if (classifierNames.isEmpty()) return null
-            return FirGeneratedClassNestedClassifierScope(useSiteSession, klass, classifierNames)
+
+            val extensionsByName = useSiteSession.getExtensionsForClass(klass).flatGroupBy {
+                it.nestedClassifierNamesCache.getValue(symbol)
+            }
+            if (extensionsByName.isEmpty()) return null
+            return FirGeneratedClassNestedClassifierScope(useSiteSession, klass, extensionsByName)
         }
     }
 
@@ -186,10 +189,25 @@ class FirGeneratedClassNestedClassifierScope private constructor(
             }
         }
 
-        if (name !in nestedClassifierNames) return null
-        val generatedClass = useSiteSession.generatedDeclarationsSymbolProvider
-            ?.getClassLikeSymbolByClassId(klass.classId.createNestedClassId(name))
-        require(generatedClass is FirRegularClassSymbol?) { "Only regular class are allowed as nested classes" }
+        val extensions = extensionsByName[name] ?: return null
+
+        val generatedClasses = extensions.mapNotNull { extension ->
+            extension.generateNestedClassLikeDeclaration(klass.symbol, name)?.also { symbol ->
+                symbol.fir.ownerGenerator = extension
+            }
+        }
+
+        val generatedClass = when (generatedClasses.size) {
+            0 -> return null
+            1 -> generatedClasses.first()
+            else -> error(
+                """
+                     Multiple plugins generated nested class with same name $name for class ${klass.classId}:
+                    ${generatedClasses.joinToString("\n") { it.fir.render() }}
+                """.trimIndent()
+            )
+        }
+        require(generatedClass is FirRegularClassSymbol) { "Only regular class are allowed as nested classes" }
         return generatedClass
     }
 
@@ -198,10 +216,10 @@ class FirGeneratedClassNestedClassifierScope private constructor(
     }
 
     override fun isEmpty(): Boolean {
-        return nestedClassifierNames.isEmpty()
+        return extensionsByName.isEmpty()
     }
 
     override fun getClassifierNames(): Set<Name> {
-        return nestedClassifierNames
+        return extensionsByName.keys
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProviderImpl.kt

@@ -10,9 +10,11 @@ import com.google.common.collect.Multimap
 import kotlinx.collections.immutable.PersistentList
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.NoMutableState
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.extensions.predicate.AbstractPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
@@ -145,6 +147,9 @@ class FirPredicateBasedProviderImpl(private val session: FirSession) : FirPredic
         // ------------------------------------ Utilities ------------------------------------
 
         private fun matchWith(declaration: FirDeclaration, annotations: Set<AnnotationFqn>): Boolean {
+            if (declaration is FirClass && declaration.isLocal) {
+                return declaration.annotations.any { it.fqName(session) in annotations }
+            }
             return cache.annotationsOfDeclaration[declaration].any { it in annotations }
         }
 

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

@@ -9,6 +9,8 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.createCurrentScopeList
 import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.runCompilerRequiredAnnotationsResolvePhaseForLocalClass
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.runCompanionGenerationPhaseForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runStatusResolveForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runSupertypeResolvePhaseForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runTypeResolvePhaseForLocalClass
@@ -32,6 +34,14 @@ fun <F : FirClassLikeDeclaration> F.runAllPhasesForLocalClass(
         components.context.outerLocalClassForNested[nested.symbol] = outer.symbol
     }
 
+    runCompilerRequiredAnnotationsResolvePhaseForLocalClass(
+        components.session,
+        components.scopeSession,
+        localClassesNavigationInfo,
+        components.file,
+        components.containingDeclarations,
+    )
+    runCompanionGenerationPhaseForLocalClass(components.session)
     runSupertypeResolvePhaseForLocalClass(
         components.session,
         components.scopeSession,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/AbstractFirSpecificAnnotationResolveTransformer.kt

@@ -45,7 +45,8 @@ import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled
 internal abstract class AbstractFirSpecificAnnotationResolveTransformer(
     protected val session: FirSession,
     protected val scopeSession: ScopeSession,
-    protected val computationSession: CompilerRequiredAnnotationsComputationSession
+    protected val computationSession: CompilerRequiredAnnotationsComputationSession,
+    containingDeclarations: List<FirDeclaration> = emptyList()
 ) : FirDefaultTransformer<Nothing?>() {
     companion object {
         private val REQUIRED_ANNOTATIONS: Set<ClassId> = setOf(Deprecated, DeprecatedSinceKotlin, WasExperimental, JvmRecord)
@@ -58,6 +59,9 @@ internal abstract class AbstractFirSpecificAnnotationResolveTransformer(
     private val annotationsFromPlugins: Set<AnnotationFqn> = session.registeredPluginAnnotations.annotations
     private val metaAnnotationsFromPlugins: Set<AnnotationFqn> = session.registeredPluginAnnotations.metaAnnotations
 
+    protected open val shouldRecordIntoPredicateBasedProvider: Boolean
+        get() = true
+
     @PrivateForInline
     val typeResolverTransformer: FirSpecificTypeResolverTransformer = FirSpecificTypeResolverTransformer(
         session,
@@ -72,7 +76,13 @@ internal abstract class AbstractFirSpecificAnnotationResolveTransformer(
     )
 
     private var owners: PersistentList<FirDeclaration> = persistentListOf()
-    private val classDeclarationsStack = ArrayDeque<FirClass>()
+    private val classDeclarationsStack = ArrayDeque<FirClass>().apply {
+        for (declaration in containingDeclarations) {
+            if (declaration is FirClass) {
+                add(declaration)
+            }
+        }
+    }
 
     @OptIn(PrivateForInline::class)
     override fun transformAnnotationCall(annotationCall: FirAnnotationCall, data: Nothing?): FirStatement {
@@ -163,7 +173,9 @@ internal abstract class AbstractFirSpecificAnnotationResolveTransformer(
     @OptIn(FirExtensionApiInternals::class)
     override fun transformDeclaration(declaration: FirDeclaration, data: Nothing?): FirDeclaration {
         return (transformAnnotationContainer(declaration, data) as FirDeclaration).also {
-            predicateBasedProvider.registerAnnotatedDeclaration(declaration, owners)
+            if (shouldRecordIntoPredicateBasedProvider) {
+                predicateBasedProvider.registerAnnotatedDeclaration(declaration, owners)
+            }
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompanionGenerationProcessor.kt

@@ -9,14 +9,20 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.extensions.FirSwitchableExtensionDeclarationsSymbolProvider
 import org.jetbrains.kotlin.fir.extensions.generatedDeclarationsSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
+import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
+import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.name.SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
 
 class FirCompanionGenerationProcessor(
     session: FirSession,
@@ -26,8 +32,8 @@ class FirCompanionGenerationProcessor(
 }
 
 class FirCompanionGenerationTransformer(val session: FirSession) : FirTransformer<Nothing?>() {
-
-    lateinit var generatedDeclarationProvider: FirSwitchableExtensionDeclarationsSymbolProvider
+    private val generatedDeclarationProvider: FirSwitchableExtensionDeclarationsSymbolProvider? =
+        session.generatedDeclarationsSymbolProvider
 
     override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
         return element
@@ -35,26 +41,45 @@ class FirCompanionGenerationTransformer(val session: FirSession) : FirTransforme
 
     override fun transformFile(file: FirFile, data: Nothing?): FirFile {
         // I don't want to use laziness here to prevent possible multi-threading problems
-        generatedDeclarationProvider = session.generatedDeclarationsSymbolProvider ?: return file
+        if (generatedDeclarationProvider == null) return file
         return file.transformDeclarations(this, data)
     }
 
     override fun transformRegularClass(regularClass: FirRegularClass, data: Nothing?): FirStatement {
-        generateCompanion(regularClass)
+        val companionSymbol = generateCompanion(regularClass)
+        if (companionSymbol != null) {
+            regularClass.replaceCompanionObjectSymbol(companionSymbol)
+        }
         return regularClass.transformDeclarations(this, data)
     }
 
-    private fun generateCompanion(regularClass: FirRegularClass) {
-        // TODO: add proper error reporting
-        generatedDeclarationProvider = session.generatedDeclarationsSymbolProvider ?: return
-        val companionClassId = regularClass.classId.createNestedClassId(SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT)
-        when (val generatedCompanion = generatedDeclarationProvider.getClassLikeSymbolByClassId(companionClassId)) {
-            null -> {}
+    private fun generateCompanion(regularClass: FirRegularClass): FirRegularClassSymbol? {
+        if (generatedDeclarationProvider == null) return null
+        val generatedCompanion = if (regularClass.isLocal) {
+            var result: FirClassLikeSymbol<*>? = null
+            session.nestedClassifierScope(regularClass)?.processClassifiersByName(DEFAULT_NAME_FOR_COMPANION_OBJECT) {
+                if (it is FirClassLikeSymbol<*> && it.origin.generated) {
+                    result = it
+                }
+            }
+            result
+        } else {
+            val companionClassId = regularClass.classId.createNestedClassId(DEFAULT_NAME_FOR_COMPANION_OBJECT)
+            generatedDeclarationProvider.getClassLikeSymbolByClassId(companionClassId)?.takeIf { it.origin.generated }
+        }
+
+        return when (generatedCompanion) {
+            null -> null
             is FirRegularClassSymbol -> when {
                 regularClass.companionObjectSymbol != null -> error("Plugin generated companion object for class $regularClass, but it is already present in class")
-                else -> regularClass.replaceCompanionObjectSymbol(generatedCompanion)
+                else -> generatedCompanion
             }
             else -> error("Plugin generated non regular class as companion object")
         }
     }
 }
+
+fun <F : FirClassLikeDeclaration> F.runCompanionGenerationPhaseForLocalClass(session: FirSession): F {
+    val transformer = FirCompanionGenerationTransformer(session)
+    return this.transformSingle(transformer, null)
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompilerRequiredAnnotationsResolveTransformer.kt

@@ -14,15 +14,11 @@ import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
 import org.jetbrains.kotlin.fir.resolve.transformers.*
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.LocalClassesNavigationInfo
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.fir.withFileAnalysisExceptionWrapping
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.name.StandardClassIds.Annotations.Deprecated
-import org.jetbrains.kotlin.name.StandardClassIds.Annotations.DeprecatedSinceKotlin
-import org.jetbrains.kotlin.name.StandardClassIds.Annotations.JvmRecord
-import org.jetbrains.kotlin.name.StandardClassIds.Annotations.WasExperimental
 
 class FirCompilerRequiredAnnotationsResolveProcessor(
     session: FirSession,
@@ -174,3 +170,42 @@ private class FirDesignatedSpecificAnnotationResolveTransformer(
         return !designation.shouldSkipClass(declaration)
     }
 }
+
+fun <F : FirClassLikeDeclaration> F.runCompilerRequiredAnnotationsResolvePhaseForLocalClass(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    localClassesNavigationInfo: LocalClassesNavigationInfo,
+    useSiteFile: FirFile,
+    containingDeclarations: List<FirDeclaration>,
+): F {
+    val computationSession = CompilerRequiredAnnotationsComputationSession()
+    val annotationsResolveTransformer = FirSpecificAnnotationForLocalClassesResolveTransformer(
+        session,
+        scopeSession,
+        computationSession,
+        containingDeclarations,
+        localClassesNavigationInfo
+    )
+    return annotationsResolveTransformer.withFileScopes(useSiteFile) {
+        this.transformSingle(annotationsResolveTransformer, null)
+    }
+}
+
+private class FirSpecificAnnotationForLocalClassesResolveTransformer(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    computationSession: CompilerRequiredAnnotationsComputationSession,
+    containingDeclarations: List<FirDeclaration>,
+    private val localClassesNavigationInfo: LocalClassesNavigationInfo
+) : AbstractFirSpecificAnnotationResolveTransformer(session, scopeSession, computationSession, containingDeclarations) {
+    override fun shouldTransformDeclaration(declaration: FirDeclaration): Boolean {
+        localClassesNavigationInfo.allMembers
+        return when (declaration) {
+            is FirClassLikeDeclaration -> declaration in localClassesNavigationInfo.parentForClass
+            else -> true
+        }
+    }
+
+    override val shouldRecordIntoPredicateBasedProvider: Boolean
+        get() = false
+}

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AdditionalMembersGenerator.kt

@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.fir.plugin.fqn
 import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
 
@@ -58,10 +57,8 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         return listOf(buildMaterializeFunction(matchedClassSymbol, callableId, Key).symbol)
     }
 
-    override fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        if (classId.shortClassName != NESTED_NAME) return null
-        val parentClassId = classId.parentClassId ?: return null
-        if (matchedClasses.none { it.classId == parentClassId }) return null
+    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+        if (matchedClasses.none { it == owner }) return null
         return buildRegularClass {
             resolvePhase = FirResolvePhase.BODY_RESOLVE
             moduleData = session.moduleData
@@ -69,16 +66,14 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
             classKind = ClassKind.CLASS
             scopeProvider = session.kotlinScopeProvider
             status = FirResolvedDeclarationStatusImpl(Visibilities.Public, Modality.FINAL, EffectiveVisibility.Public)
-            name = classId.shortClassName
-            symbol = FirRegularClassSymbol(classId)
+            this.name = name
+            symbol = FirRegularClassSymbol(owner.classId.createNestedClassId(name))
             superTypeRefs += session.builtinTypes.anyType
         }.symbol
     }
 
     override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
-        val ownerClassId = context.owner.classId
-        assert(ownerClassId in nestedClassIds)
-        return listOf(buildConstructor(ownerClassId, isInner = false, Key).symbol)
+        return listOf(buildConstructor(context.owner, isInner = false, Key).symbol)
     }
 
     override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/CompanionGenerator.kt

@@ -17,11 +17,8 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.origin
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
-import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
+import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.plugin.fqn
 import org.jetbrains.kotlin.fir.resolve.defaultType
@@ -29,7 +26,6 @@ import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
 
@@ -42,15 +38,8 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
         private val FOO_NAME = Name.identifier("foo")
     }
 
-    private val matchedClasses by lazy {
-        session.predicateBasedProvider.getSymbolsByPredicate(PREDICATE)
-            .filterIsInstance<FirRegularClassSymbol>()
-    }
-
-    override fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        if (classId.shortClassName != SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT) return null
-        val owner = matchedClasses.firstOrNull { it.classId == classId.outerClassId } ?: return null
-        if (owner.companionObjectSymbol != null) return null
+    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+        if (name != SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT) return null
         val regularClass = buildRegularClass {
             resolvePhase = FirResolvePhase.BODY_RESOLVE
             moduleData = session.moduleData
@@ -64,8 +53,8 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
             ).apply {
                 isCompanion = true
             }
-            name = SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
-            symbol = FirRegularClassSymbol(classId)
+            this.name = name
+            symbol = FirRegularClassSymbol(owner.classId.createNestedClassId(name))
             superTypeRefs += session.builtinTypes.anyType
         }
         return regularClass.symbol
@@ -94,7 +83,7 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
     }
 
     override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
-        val constructor = buildConstructor(context.owner.classId, isInner = false, Key)
+        val constructor = buildConstructor(context.owner, isInner = false, Key)
         return listOf(constructor.symbol)
     }
 

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/ExternalClassGenerator.kt

@@ -61,30 +61,26 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
         }
     }
 
-    override fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        val owner = classId.outerClassId?.let { session.symbolProvider.getClassLikeSymbolByClassId(it) } as? FirClassSymbol<*>
-        return when {
-            owner != null -> when (val origin = owner.origin) {
-                is FirDeclarationOrigin.Plugin -> when (origin.key) {
-                    Key -> generateNestedClass(classId, owner)
-                    else -> null
-                }
-                else -> null
-            }
-            else -> generateAllOpenGeneratedClass(classId)
-        }
-    }
-
-    private fun generateAllOpenGeneratedClass(classId: ClassId): FirClassLikeSymbol<*>? {
+    override fun generateTopLevelClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
         if (classId != GENERATED_CLASS_ID) return null
         if (matchedClasses.isEmpty()) return null
         return buildClass(classId).symbol
     }
 
+    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+        return when (val origin = owner.origin) {
+            is FirDeclarationOrigin.Plugin -> when (origin.key) {
+                Key -> generateNestedClass(owner.classId.createNestedClassId(name), owner)
+                else -> null
+            }
+            else -> null
+        }
+    }
+
     override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
         val classId = context.owner.classId
         if (classId != GENERATED_CLASS_ID && classId !in classIdsForMatchedClasses) return emptyList()
-        return listOf(buildConstructor(classId, isInner = false, Key).symbol)
+        return listOf(buildConstructor(context.owner, isInner = false, Key).symbol)
     }
 
     private fun generateNestedClass(classId: ClassId, owner: FirClassSymbol<*>): FirClassLikeSymbol<*>? {

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/generationUtils.kt

@@ -19,10 +19,7 @@ import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
-import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeKotlinTypeProjection
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
@@ -63,8 +60,9 @@ fun FirDeclarationGenerationExtension.buildMaterializeFunction(
 
 // FIXME: this has to be shared
 @OptIn(SymbolInternals::class)
-fun FirDeclarationGenerationExtension.buildConstructor(classId: ClassId, isInner: Boolean, key: GeneratedDeclarationKey): FirConstructor {
-    val lookupTag = ConeClassLikeLookupTagImpl(classId)
+fun FirDeclarationGenerationExtension.buildConstructor(owner: FirClassSymbol<*>, isInner: Boolean, key: GeneratedDeclarationKey): FirConstructor {
+    val classId = owner.classId
+    val lookupTag = owner.toLookupTag()
     return buildPrimaryConstructor {
         resolvePhase = FirResolvePhase.BODY_RESOLVE
         moduleData = session.moduleData

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/AbstractTransformerForGenerator.kt

@@ -20,7 +20,10 @@ import org.jetbrains.kotlin.ir.util.primaryConstructor
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 
-abstract class AbstractTransformerForGenerator(protected val context: IrPluginContext) : IrElementVisitorVoid {
+abstract class AbstractTransformerForGenerator(
+    protected val context: IrPluginContext,
+    private val visitBodies: Boolean
+) : IrElementVisitorVoid {
     protected val irFactory = context.irFactory
     protected val irBuiltIns = context.irBuiltIns
 
@@ -29,24 +32,38 @@ abstract class AbstractTransformerForGenerator(protected val context: IrPluginCo
     abstract fun generateBodyForConstructor(constructor: IrConstructor, key: GeneratedDeclarationKey): IrBody?
 
     final override fun visitElement(element: IrElement) {
-        when (element) {
-            is IrDeclaration,
-            is IrFile,
-            is IrModuleFragment -> element.acceptChildrenVoid(this)
-            else -> {}
+        if (visitBodies) {
+            element.acceptChildrenVoid(this)
+        } else {
+            when (element) {
+                is IrDeclaration,
+                is IrFile,
+                is IrModuleFragment -> element.acceptChildrenVoid(this)
+                else -> {}
+            }
         }
     }
 
     final override fun visitSimpleFunction(declaration: IrSimpleFunction) {
         val origin = declaration.origin
-        if (origin !is GeneratedByPlugin || !interestedIn(origin.pluginKey)) return
+        if (origin !is GeneratedByPlugin || !interestedIn(origin.pluginKey)) {
+            if (visitBodies) {
+                visitElement(declaration)
+            }
+            return
+        }
         require(declaration.body == null)
         declaration.body = generateBodyForFunction(declaration, origin.pluginKey)
     }
 
     final override fun visitConstructor(declaration: IrConstructor) {
         val origin = declaration.origin
-        if (origin !is GeneratedByPlugin || !interestedIn(origin.pluginKey)) return
+        if (origin !is GeneratedByPlugin || !interestedIn(origin.pluginKey)) {
+            if (visitBodies) {
+                visitElement(declaration)
+            }
+            return
+        }
         require(declaration.body == null)
         declaration.body = generateBodyForConstructor(declaration, origin.pluginKey)
     }

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/TransformerForAdditionalMembersGenerator.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.IrBody
 
-class TransformerForAdditionalMembersGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context) {
+class TransformerForAdditionalMembersGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context, visitBodies = false) {
     override fun interestedIn(key: GeneratedDeclarationKey): Boolean {
         return key == AdditionalMembersGenerator.Key
     }

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/TransformerForCompanionGenerator.kt

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
 
-class TransformerForCompanionGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context) {
+class TransformerForCompanionGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context, visitBodies = true) {
     override fun interestedIn(key: GeneratedDeclarationKey): Boolean {
         return key == CompanionGenerator.Key
     }

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/TransformerForExternalClassGenerator.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.IrBody
 
-class TransformerForExternalClassGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context) {
+class TransformerForExternalClassGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context, visitBodies = false) {
     override fun interestedIn(key: GeneratedDeclarationKey): Boolean {
         return key == ExternalClassGenerator.Key
     }

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/TransformerForTopLevelDeclarationsGenerator.kt

@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
 import org.jetbrains.kotlin.ir.types.classFqName
 
-class TransformerForTopLevelDeclarationsGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context) {
+class TransformerForTopLevelDeclarationsGenerator(context: IrPluginContext) : AbstractTransformerForGenerator(context, visitBodies = false) {
     override fun interestedIn(key: GeneratedDeclarationKey): Boolean {
         return key == TopLevelDeclarationsGenerator.Key
     }

plugins/fir-plugin-prototype/testData/box/localClassWithCompanionObject.fir.txt

@@ -0,0 +1,33 @@
+FILE: localClassWithCompanionObject.kt
+    public final fun box(): R|kotlin/String| {
+        @R|org/jetbrains/kotlin/fir/plugin/CompanionWithFoo|() local final class SomeClass : R|kotlin/Any| {
+            public constructor(): R|SomeClass| {
+                super<R|kotlin/Any|>()
+            }
+
+            public final companion object Companion : R|kotlin/Any| {
+                public final fun foo(): R|kotlin/Int|
+
+                public constructor(): R|SomeClass.Companion|
+
+            }
+
+        }
+
+        local final fun takeInt(x: R|kotlin/Int|): R|kotlin/Unit| {
+            when () {
+                !=(R|<local>/x|, Int(10)) ->  {
+                    throw R|java/lang/IllegalArgumentException.IllegalArgumentException|()
+                }
+            }
+
+        }
+
+        local final fun test(): R|kotlin/Unit| {
+            R|<local>/takeInt|(Q|SomeClass|.R|/SomeClass.Companion.foo|())
+            R|<local>/takeInt|(Q|SomeClass.Companion|.R|/SomeClass.Companion.foo|())
+        }
+
+        R|<local>/test|()
+        ^box String(OK)
+    }

plugins/fir-plugin-prototype/testData/box/localClassWithCompanionObject.kt

@@ -0,0 +1,18 @@
+import org.jetbrains.kotlin.fir.plugin.CompanionWithFoo
+
+fun box(): String {
+    @CompanionWithFoo
+    class SomeClass
+
+    fun takeInt(x: Int) {
+        if (x != 10) throw IllegalArgumentException()
+    }
+
+    fun test() {
+        takeInt(SomeClass.foo())
+        takeInt(SomeClass.Companion.foo())
+    }
+
+    test()
+    return "OK"
+}

plugins/fir-plugin-prototype/testData/diagnostics/memberGen/localClassWithCompanionObject.fir.txt

@@ -0,0 +1,41 @@
+FILE: localClassWithCompanionObject.kt
+    public final fun context(): R|kotlin/Unit| {
+        @R|org/jetbrains/kotlin/fir/plugin/CompanionWithFoo|() local final class SomeClass : R|kotlin/Any| {
+            public constructor(): R|SomeClass| {
+                super<R|kotlin/Any|>()
+            }
+
+            @R|org/jetbrains/kotlin/fir/plugin/CompanionWithFoo|() local final class Nested : R|kotlin/Any| {
+                public constructor(): R|SomeClass.Nested| {
+                    super<R|kotlin/Any|>()
+                }
+
+                public final companion object Companion : R|kotlin/Any| {
+                    public final fun foo(): R|kotlin/Int|
+
+                    public constructor(): R|SomeClass.Nested.Companion|
+
+                }
+
+            }
+
+            public final companion object Companion : R|kotlin/Any| {
+                public final fun foo(): R|kotlin/Int|
+
+                public constructor(): R|SomeClass.Companion|
+
+            }
+
+        }
+
+        local final fun takeInt(x: R|kotlin/Int|): R|kotlin/Unit| {
+        }
+
+        local final fun test(): R|kotlin/Unit| {
+            R|<local>/takeInt|(Q|SomeClass|.R|/SomeClass.Companion.foo|())
+            R|<local>/takeInt|(Q|SomeClass.Companion|.R|/SomeClass.Companion.foo|())
+            R|<local>/takeInt|(Q|SomeClass.Nested|.R|/SomeClass.Nested.Companion.foo|())
+            R|<local>/takeInt|(Q|SomeClass.Nested.Companion|.R|/SomeClass.Nested.Companion.foo|())
+        }
+
+    }

plugins/fir-plugin-prototype/testData/diagnostics/memberGen/localClassWithCompanionObject.kt

@@ -0,0 +1,20 @@
+import org.jetbrains.kotlin.fir.plugin.CompanionWithFoo
+
+fun context() {
+    @CompanionWithFoo
+    class SomeClass {
+
+        @CompanionWithFoo
+        <!NESTED_CLASS_NOT_ALLOWED!>class Nested<!>
+    }
+
+    fun takeInt(x: Int) {}
+
+    fun test() {
+        takeInt(SomeClass.foo())
+        takeInt(SomeClass.Companion.foo())
+
+        takeInt(SomeClass.Nested.foo())
+        takeInt(SomeClass.Nested.Companion.foo())
+    }
+}

plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPluginBlackBoxCodegenTestGenerated.java

@@ -43,6 +43,12 @@ public class FirPluginBlackBoxCodegenTestGenerated extends AbstractFirPluginBlac
         runTest("plugins/fir-plugin-prototype/testData/box/generatedClassWithMembersAndNestedClasses.kt");
     }
 
+    @Test
+    @TestMetadata("localClassWithCompanionObject.kt")
+    public void testLocalClassWithCompanionObject() throws Exception {
+        runTest("plugins/fir-plugin-prototype/testData/box/localClassWithCompanionObject.kt");
+    }
+
     @Test
     @TestMetadata("metaAnnotationFromLibrary.kt")
     public void testMetaAnnotationFromLibrary() throws Exception {

plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPluginDiagnosticTestGenerated.java

@@ -73,6 +73,12 @@ public class FirPluginDiagnosticTestGenerated extends AbstractFirPluginDiagnosti
             runTest("plugins/fir-plugin-prototype/testData/diagnostics/memberGen/generatedClassWithMembersAndNestedClasses.kt");
         }
 
+        @Test
+        @TestMetadata("localClassWithCompanionObject.kt")
+        public void testLocalClassWithCompanionObject() throws Exception {
+            runTest("plugins/fir-plugin-prototype/testData/diagnostics/memberGen/localClassWithCompanionObject.kt");
+        }
+
         @Test
         @TestMetadata("topLevelCallables.kt")
         public void testTopLevelCallables() throws Exception {

plugins/kotlinx-serialization/kotlinx-serialization.k2/src/org/jetbrains/kotlinx/serialization/compiler/fir/FirBuildersUtils.kt

@@ -30,7 +30,7 @@ import org.jetbrains.kotlin.types.Variance
 @OptIn(SymbolInternals::class)
 fun FirDeclarationGenerationExtension.buildPrimaryConstructor(owner: FirClassSymbol<*>, isInner: Boolean, key: GeneratedDeclarationKey, status: FirDeclarationStatus): FirConstructor {
     val classId = owner.classId
-    val lookupTag = ConeClassLikeLookupTagImpl(classId)
+    val lookupTag = owner.toLookupTag()
     return buildPrimaryConstructor {
         moduleData = session.moduleData
         origin = key.origin

plugins/kotlinx-serialization/kotlinx-serialization.k2/src/org/jetbrains/kotlinx/serialization/compiler/fir/SerializationFirResolveExtension.kt

@@ -11,16 +11,16 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.getContainingDeclarationSymbol
+import org.jetbrains.kotlin.fir.containingClassForLocalAttr
 import org.jetbrains.kotlin.fir.copy
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.builder.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.origin
 import org.jetbrains.kotlin.fir.declarations.utils.isCompanion
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
-import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.declarations.utils.isLocal
+import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
@@ -68,19 +68,16 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
         return result
     }
 
-    override fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        return when (classId.shortClassName) {
-            SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT -> {
-                generateCompanionDeclaration(classId)
-            }
-            SerialEntityNames.SERIALIZER_CLASS_NAME -> {
-                addSerializerImplClass(classId)
-            }
-            else -> error("Can't generate class ${classId.asSingleFqName()}")
+    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+        if (owner !is FirRegularClassSymbol) return null
+        if (!session.predicateBasedProvider.matches(FirSerializationPredicates.annotatedWithSerializableOrMeta, owner)) return null
+        return when (name) {
+            SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT -> generateCompanionDeclaration(owner)
+            SerialEntityNames.SERIALIZER_CLASS_NAME -> generateSerializerImplClass(owner)
+            else -> error("Can't generate class ${owner.classId.createNestedClassId(name).asSingleFqName()}")
         }
     }
 
-
     override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
         val classId = classSymbol.classId
         val result = mutableSetOf<Name>()
@@ -285,15 +282,7 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
         return listOf(defaultObjectConstructor.symbol)
     }
 
-    private fun getClassWithAnnotatedWithSerializableOrMeta(classId: ClassId?): FirRegularClassSymbol? {
-        if (classId == null) return null
-        return (session.symbolProvider.getClassLikeSymbolByClassId(classId) as? FirRegularClassSymbol)?.takeIf {
-            session.predicateBasedProvider.matches(FirSerializationPredicates.annotatedWithSerializableOrMeta, it)
-        }
-    }
-
-    fun addSerializerImplClass(classId: ClassId): FirClassLikeSymbol<*>? {
-        val owner = getClassWithAnnotatedWithSerializableOrMeta(classId.outerClassId) ?: return null
+    private fun generateSerializerImplClass(owner: FirRegularClassSymbol): FirClassLikeSymbol<*> {
         val hasTypeParams = owner.typeParameterSymbols.isNotEmpty()
         val serializerKind = if (hasTypeParams) ClassKind.CLASS else ClassKind.OBJECT
         val serializerFirClass = buildRegularClass {
@@ -302,7 +291,7 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
             classKind = serializerKind
             scopeProvider = session.kotlinScopeProvider
             name = SerialEntityNames.SERIALIZER_CLASS_NAME
-            symbol = FirRegularClassSymbol(classId)
+            symbol = FirRegularClassSymbol(owner.classId.createNestedClassId(name))
             status = FirResolvedDeclarationStatusImpl(
                 Visibilities.Public,
                 Modality.FINAL,
@@ -325,13 +314,11 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
                     )
                 ), isNullable = false
             ).toFirResolvedTypeRef()
-        }
+        }.also { it.initLocalAttributeIfNeeded(owner) }
         return serializerFirClass.symbol
     }
 
-    fun generateCompanionDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        if (classId.shortClassName != SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT) return null
-        val owner = getClassWithAnnotatedWithSerializableOrMeta(classId.outerClassId) ?: return null
+    private fun generateCompanionDeclaration(owner: FirRegularClassSymbol): FirRegularClassSymbol? {
         if (owner.companionObjectSymbol != null) return null
         val regularClass = buildRegularClass {
             moduleData = session.moduleData
@@ -346,16 +333,22 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
                 isCompanion = true
             }
             name = SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
-            symbol = FirRegularClassSymbol(classId)
+            symbol = FirRegularClassSymbol(owner.classId.createNestedClassId(name))
             superTypeRefs += session.builtinTypes.anyType
             if (with(session) { owner.companionNeedsSerializerFactory }) {
                 val serializerFactoryClassId = ClassId(SerializationPackages.internalPackageFqName, SERIALIZER_FACTORY_INTERFACE_NAME)
                 superTypeRefs += serializerFactoryClassId.constructClassLikeType(emptyArray(), false).toFirResolvedTypeRef()
             }
-        }
+        }.also { it.initLocalAttributeIfNeeded(owner) }
         return regularClass.symbol
     }
 
+    private fun FirRegularClass.initLocalAttributeIfNeeded(owner: FirRegularClassSymbol) {
+        if (owner.isLocal) {
+            this.containingClassForLocalAttr = owner.toLookupTag()
+        }
+    }
+
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {
         register(FirSerializationPredicates.annotatedWithSerializableOrMeta, FirSerializationPredicates.hasMetaAnnotation)
     }

plugins/kotlinx-serialization/testData/firMembers/serializerInLocalClass.fir.txt

@@ -0,0 +1,50 @@
+FILE: serializerInLocalClass.kt
+    public final fun local(): R|kotlin/String| {
+        @R|kotlinx/serialization/Serializable|() local final data class Carrier : R|kotlin/Any| {
+            public constructor(i: R|kotlin/Int|): R|Carrier| {
+                super<R|kotlin/Any|>()
+            }
+
+            public final val i: R|kotlin/Int| = R|<local>/i|
+                public get(): R|kotlin/Int|
+
+            public final operator fun component1(): R|kotlin/Int|
+
+            public final fun copy(i: R|kotlin/Int| = this@R|/Carrier|.R|<local>/i|): R|Carrier|
+
+            public final companion object Companion : R|kotlin/Any| {
+                public final fun serializer(): R|kotlinx/serialization/KSerializer<Carrier>|
+
+                private constructor(): R|Carrier.Companion|
+
+            }
+
+            public final object $serializer : R|kotlinx/serialization/internal/GeneratedSerializer<Carrier>| {
+                public final fun serialize(encoder: R|kotlinx/serialization/encoding/Encoder|, value: R|Carrier|): R|kotlin/Unit|
+
+                public final fun deserialize(decoder: R|kotlinx/serialization/encoding/Decoder|): R|Carrier|
+
+                public final val descriptor: R|kotlinx/serialization/descriptors/SerialDescriptor|
+                    public get(): R|kotlinx/serialization/descriptors/SerialDescriptor|
+
+                public final fun childSerializers(): R|kotlin/Array<kotlinx/serialization/KSerializer<*>>|
+
+                private constructor(): R|Carrier.$serializer|
+
+            }
+
+        }
+
+        ^local Q|Carrier|.R|/Carrier.Companion.serializer|().R|SubstitutionOverride<kotlinx/serialization/KSerializer.descriptor: R|kotlinx/serialization/descriptors/SerialDescriptor|>|.R|kotlin/Any.toString|()
+    }
+    public final fun box(): R|kotlin/String| {
+        lval expected: R|kotlin/String| = String(local.Carrier(i: kotlin.Int))
+        lval actual: R|kotlin/String| = R|/local|()
+        when () {
+            !=(R|<local>/expected|, R|<local>/actual|) ->  {
+                ^box <strcat>(String(Fail: ), R|<local>/actual|)
+            }
+        }
+
+        ^box String(OK)
+    }

plugins/kotlinx-serialization/testData/firMembers/serializerInLocalClass.kt

@@ -0,0 +1,20 @@
+// TARGET_BACKEND: JVM_IR
+// WITH_STDLIB
+
+import kotlinx.serialization.*
+
+fun local(): String {
+    @Serializable
+    data class Carrier(val i: Int)
+
+    return Carrier.serializer().descriptor.toString()
+}
+
+fun box(): String {
+    val expected = "local.Carrier(i: kotlin.Int)"
+    val actual = local()
+    if (expected != actual) {
+        return "Fail: $actual"
+    }
+    return "OK"
+}

plugins/kotlinx-serialization/tests-gen/org/jetbrains/kotlinx/serialization/runners/SerializationFirBlackBoxTestGenerated.java

@@ -265,6 +265,12 @@ public class SerializationFirBlackBoxTestGenerated extends AbstractSerialization
             runTest("plugins/kotlinx-serialization/testData/firMembers/serializableWithCompanion.kt");
         }
 
+        @Test
+        @TestMetadata("serializerInLocalClass.kt")
+        public void testSerializerInLocalClass() throws Exception {
+            runTest("plugins/kotlinx-serialization/testData/firMembers/serializerInLocalClass.kt");
+        }
+
         @Test
         @TestMetadata("serializerViaCompanion.kt")
         public void testSerializerViaCompanion() throws Exception {

plugins/kotlinx-serialization/tests-gen/org/jetbrains/kotlinx/serialization/runners/SerializationFirDiagnosticTestGenerated.java

@@ -240,6 +240,12 @@ public class SerializationFirDiagnosticTestGenerated extends AbstractSerializati
             runTest("plugins/kotlinx-serialization/testData/firMembers/serializableWithCompanion.kt");
         }
 
+        @Test
+        @TestMetadata("serializerInLocalClass.kt")
+        public void testSerializerInLocalClass() throws Exception {
+            runTest("plugins/kotlinx-serialization/testData/firMembers/serializerInLocalClass.kt");
+        }
+
         @Test
         @TestMetadata("serializerViaCompanion.kt")
         public void testSerializerViaCompanion() throws Exception {

plugins/lombok/lombok.k2/src/org/jetbrains/kotlin/lombok/k2/generators/BuilderGenerator.kt

@@ -78,11 +78,9 @@ class BuilderGenerator(session: FirSession) : FirDeclarationGenerationExtension(
         return functionsCache.getValue(classSymbol)?.get(callableId.callableName).orEmpty().map { it.symbol }
     }
 
-    override fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        val outerClassId = classId.outerClassId ?: return null
-        val outerClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(outerClassId) as? FirClassSymbol<*> ?: return null
-        if (!outerClassSymbol.isSuitableJavaClass()) return null
-        return builderClassCache.getValue(outerClassSymbol)?.symbol
+    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+        if (!owner.isSuitableJavaClass()) return null
+        return builderClassCache.getValue(owner)?.symbol
     }
 
     private fun createFunctions(classSymbol: FirClassSymbol<*>): Map<Name, List<FirJavaMethod>>? {