[FIR] Add ability to generate members and nested classifiers for existing classes

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirDeclaredMemberScope.kt

@@ -5,19 +5,18 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.scopes
 
-import org.jetbrains.kotlin.fir.scopes.impl.FirClassDeclaredMemberScope
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
 import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
-import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
 import org.jetbrains.kotlin.analysis.api.scopes.KtDeclaredMemberScope
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
+import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 
 internal class KtFirDeclaredMemberScope(
     override val owner: KtSymbolWithMembers,
-    override val firScope: FirClassDeclaredMemberScope,
+    override val firScope: FirContainingNamesAwareScope,
     token: ValidityToken,
     builder: KtSymbolByFirBuilder
-) : KtFirDelegatingScope<FirClassDeclaredMemberScope>(builder, token),
+) : KtFirDelegatingScope<FirContainingNamesAwareScope>(builder, token),
     KtDeclaredMemberScope,
-    ValidityTokenOwner
+    ValidityTokenOwner

compiler/fir/cones/src/org/jetbrains/kotlin/fir/utils/ArrayMapOwner.kt

@@ -41,10 +41,13 @@ abstract class AbstractArrayMapOwner<K : Any, V : Any> : Iterable<V> {
 
 class ArrayMapAccessor<K : Any, V : Any, T : V>(
     key: KClass<out K>,
-    id: Int
+    id: Int,
+    val default: T? = null
 ) : AbstractArrayMapOwner.AbstractArrayMapAccessor<K, V, T>(key, id), ReadOnlyProperty<AbstractArrayMapOwner<K, V>, V> {
     override fun getValue(thisRef: AbstractArrayMapOwner<K, V>, property: KProperty<*>): T {
-        return extractValue(thisRef) ?: error("No '$key'($id) in array owner: $thisRef")
+        return extractValue(thisRef)
+            ?: default
+            ?: error("No '$key'($id) in array owner: $thisRef")
     }
 }
 
@@ -62,8 +65,8 @@ abstract class TypeRegistry<K : Any, V : Any> {
     private val idCounter = AtomicInteger(0)
 
 
-    fun <T : V, KK : K> generateAccessor(kClass: KClass<KK>): ArrayMapAccessor<K, V, T> {
-        return ArrayMapAccessor(kClass, getId(kClass))
+    fun <T : V, KK : K> generateAccessor(kClass: KClass<KK>, default: T? = null): ArrayMapAccessor<K, V, T> {
+        return ArrayMapAccessor(kClass, getId(kClass), default)
     }
 
     fun <T : V, KK : K> generateNullableAccessor(kClass: KClass<KK>): NullableArrayMapAccessor<K, V, T> {

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

@@ -38,8 +38,8 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredi
     // Can be called on IMPORTS stage
     open fun hasPackage(packageFqName: FqName): Boolean = false
 
-    open fun getCallableNamesForGeneratedClass(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
-    open fun getNestedClassifiersNamesForGeneratedClass(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
+    open fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
+    open fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
 
     fun interface Factory : FirExtension.Factory<FirDeclarationGenerationExtension>
 }

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

@@ -20,11 +20,13 @@ annotation class PluginServicesInitialization
 class FirExtensionService(val session: FirSession) : ComponentArrayOwner<FirExtension, List<FirExtension>>(), FirSessionComponent {
     companion object : TypeRegistry<FirExtension, List<FirExtension>>() {
         inline fun <reified P : FirExtension, V : List<P>> registeredExtensions(): ArrayMapAccessor<FirExtension, List<FirExtension>, V> {
-            return generateAccessor(P::class)
+            @Suppress("UNCHECKED_CAST")
+            return generateAccessor(P::class, default = emptyList<P>() as V)
         }
 
         fun <P : FirExtension, V : List<P>> registeredExtensions(kClass: KClass<P>): ArrayMapAccessor<FirExtension, List<FirExtension>, V> {
-            return generateAccessor(kClass)
+            @Suppress("UNCHECKED_CAST")
+            return generateAccessor(kClass, default = emptyList<P>() as V)
         }
     }
 

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

@@ -12,16 +12,20 @@ import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.extensions.declarationGenerators
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.resolve.declaredMemberScopeProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.FirCompositeScope
+import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 @ThreadSafeMutableState
 class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessionComponent {
-    private val declaredMemberCache: FirCache<FirClass, FirClassDeclaredMemberScope, DeclaredMemberScopeContext> =
+    private val declaredMemberCache: FirCache<FirClass, FirContainingNamesAwareScope, DeclaredMemberScopeContext> =
         useSiteSession.firCachesFactory.createCache { klass, context ->
             createDeclaredMemberScope(klass, context.useLazyNestedClassifierScope, context.existingNames, context.symbolProvider)
         }
@@ -29,12 +33,14 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
     private val nestedClassifierCache: FirCache<FirClass, FirNestedClassifierScope?, Nothing?> =
         useSiteSession.firCachesFactory.createCache { klass, _ -> createNestedClassifierScope(klass) }
 
+    private val extensions by lazy(LazyThreadSafetyMode.PUBLICATION) { useSiteSession.extensionService.declarationGenerators }
+
     fun declaredMemberScope(
         klass: FirClass,
         useLazyNestedClassifierScope: Boolean,
         existingNames: List<Name>?,
         symbolProvider: FirSymbolProvider?
-    ): FirClassDeclaredMemberScope {
+    ): FirContainingNamesAwareScope {
         return declaredMemberCache.getValue(klass, DeclaredMemberScopeContext(useLazyNestedClassifierScope, existingNames, symbolProvider))
     }
 
@@ -49,11 +55,27 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
         useLazyNestedClassifierScope: Boolean,
         existingNames: List<Name>?,
         symbolProvider: FirSymbolProvider?
-    ): FirClassDeclaredMemberScope {
-        return if (klass.origin.generated) {
-            FirGeneratedClassDeclaredMemberScope(useSiteSession, klass)
-        } else {
-            FirClassDeclaredMemberScopeImpl(useSiteSession, klass, useLazyNestedClassifierScope, existingNames, symbolProvider)
+    ): FirContainingNamesAwareScope {
+        return when {
+            klass.origin.generated -> {
+                FirGeneratedClassDeclaredMemberScope(useSiteSession, klass, needNestedClassifierScope = true)
+            }
+            else -> {
+                val baseScope = FirClassDeclaredMemberScopeImpl(
+                    useSiteSession,
+                    klass,
+                    useLazyNestedClassifierScope,
+                    existingNames,
+                    symbolProvider
+                )
+                if (extensions.any { useSiteSession.predicateBasedProvider.matches(it.predicate, klass) }) {
+                    FirCompositeScope(
+                        listOf(baseScope, FirGeneratedClassDeclaredMemberScope(useSiteSession, klass, needNestedClassifierScope = false))
+                    )
+                } else {
+                    baseScope
+                }
+            }
         }
     }
 
@@ -65,17 +87,26 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
         return if (klass.origin.generated) {
             FirGeneratedClassNestedClassifierScope(klass, useSiteSession)
         } else {
-            FirNestedClassifierScopeImpl(klass, useSiteSession)
+            val baseScope = FirNestedClassifierScopeImpl(klass, useSiteSession)
+            if (extensions.any { useSiteSession.predicateBasedProvider.matches(it.predicate, klass) }) {
+                FirCompositeNestedClassifierScope(
+                    listOf(baseScope, FirGeneratedClassNestedClassifierScope(klass, useSiteSession)),
+                    klass,
+                    useSiteSession
+                )
+            } else {
+                baseScope
+            }
         }.takeUnless { it.isEmpty() }
     }
 }
 
-fun FirSession.declaredMemberScope(klass: FirClass): FirClassDeclaredMemberScope {
+fun FirSession.declaredMemberScope(klass: FirClass): FirContainingNamesAwareScope {
     return declaredMemberScopeProvider
         .declaredMemberScope(klass, useLazyNestedClassifierScope = false, existingNames = null, symbolProvider = null)
 }
 
-fun FirSession.declaredMemberScope(klass: FirClassSymbol<*>): FirClassDeclaredMemberScope {
+fun FirSession.declaredMemberScope(klass: FirClassSymbol<*>): FirContainingNamesAwareScope {
     return declaredMemberScope(klass.fir)
 }
 
@@ -83,7 +114,7 @@ fun FirSession.declaredMemberScopeWithLazyNestedScope(
     klass: FirClass,
     existingNames: List<Name>,
     symbolProvider: FirSymbolProvider
-): FirScope {
+): FirContainingNamesAwareScope {
     return declaredMemberScopeProvider
         .declaredMemberScope(klass, useLazyNestedClassifierScope = true, existingNames = existingNames, symbolProvider = symbolProvider)
 }

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

@@ -16,19 +16,23 @@ import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.declarationGenerators
 import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class FirGeneratedClassDeclaredMemberScope(
     val useSiteSession: FirSession,
-    val firClass: FirClass
+    val firClass: FirClass,
+    needNestedClassifierScope: Boolean
 ) : FirClassDeclaredMemberScope() {
-    private val extension: FirDeclarationGenerationExtension = firClass.findGeneratedExtension(useSiteSession)
-    private val nestedClassifierScope: FirNestedClassifierScope? = useSiteSession.nestedClassifierScope(firClass)
+    private val extensions: List<FirDeclarationGenerationExtension> = firClass.findGeneratedExtensions(useSiteSession)
+    private val nestedClassifierScope: FirNestedClassifierScope? = runIf(needNestedClassifierScope) {
+        useSiteSession.nestedClassifierScope(firClass)
+    }
 
     private val firCachesFactory = useSiteSession.firCachesFactory
 
@@ -47,17 +51,17 @@ class FirGeneratedClassDeclaredMemberScope(
     }
 
     private val callableNamesCache: FirLazyValue<Set<Name>, Nothing?> = firCachesFactory.createLazyValue {
-        extension.getCallableNamesForGeneratedClass(firClass.symbol)
+        extensions.flatMapTo(mutableSetOf()) { it.getCallableNamesForClass(firClass.symbol) }
     }
 
     // ------------------------------------------ generators ------------------------------------------
 
     private fun generateMemberFunctions(name: Name): List<FirNamedFunctionSymbol> {
-        return extension.generateFunctions(CallableId(firClass.classId, name), firClass.symbol)
+        return extensions.flatMap { it.generateFunctions(CallableId(firClass.classId, name), firClass.symbol) }
     }
 
     private fun generateMemberProperties(name: Name): List<FirPropertySymbol> {
-        return extension.generateProperties(CallableId(firClass.classId, name), firClass.symbol)
+        return extensions.flatMap { it.generateProperties(CallableId(firClass.classId, name), firClass.symbol) }
     }
 
     private fun generateConstructors(): List<FirConstructorSymbol> {
@@ -67,7 +71,7 @@ class FirGeneratedClassDeclaredMemberScope(
         } else {
             CallableId(classId.asSingleFqName().parent(), classId.shortClassName)
         }
-        return extension.generateConstructors(callableId)
+        return extensions.flatMap { it.generateConstructors(callableId) }
     }
 
     // ------------------------------------------ scope methods ------------------------------------------
@@ -109,7 +113,7 @@ class FirGeneratedClassNestedClassifierScope(
     klass: FirClass,
     useSiteSession: FirSession
 ) : FirNestedClassifierScope(klass, useSiteSession) {
-    private val extension = klass.findGeneratedExtension(useSiteSession)
+    private val extensions = klass.findGeneratedExtensions(useSiteSession)
 
     private val nestedClassifierCache: FirCache<Name, FirRegularClassSymbol?, Nothing?> =
         useSiteSession.firCachesFactory.createCache { name, _ ->
@@ -118,7 +122,7 @@ class FirGeneratedClassNestedClassifierScope(
 
     private val nestedClassifiersNames: FirLazyValue<Set<Name>, Nothing?> =
         useSiteSession.firCachesFactory.createLazyValue {
-            extension.getNestedClassifiersNamesForGeneratedClass(klass.symbol)
+            extensions.flatMapTo(mutableSetOf()) { it.getNestedClassifiersNames(klass.symbol) }
         }
 
     private fun generateNestedClassifier(name: Name): FirRegularClassSymbol? {
@@ -142,13 +146,15 @@ class FirGeneratedClassNestedClassifierScope(
 }
 
 
-private fun FirClass.findGeneratedExtension(useSiteSession: FirSession): FirDeclarationGenerationExtension {
+private fun FirClass.findGeneratedExtensions(useSiteSession: FirSession): List<FirDeclarationGenerationExtension> {
     val origin = origin
-    require(origin is FirDeclarationOrigin.Plugin) {
-        "GeneratedClassDeclaredMemberScope can not be created for non-generated class: ${this.render()}"
+    val declarationGenerators = useSiteSession.extensionService.declarationGenerators
+    return if (origin is FirDeclarationOrigin.Plugin) {
+        declarationGenerators.filter { it.key == origin.key }.also {
+            require(it.isNotEmpty()) { "Extension for ${origin.key} not found" }
+        }
+    } else {
+        val predicateBasedProvider = useSiteSession.predicateBasedProvider
+        declarationGenerators.filter { predicateBasedProvider.matches(it.predicate, this) }
     }
-
-    return useSiteSession.extensionService.declarationGenerators.firstOrNull {
-        it.key == origin.key
-    } ?: error("Extension for ${origin.key} not found")
 }

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

@@ -59,6 +59,28 @@ class FirNestedClassifierScopeImpl(klass: FirClass, useSiteSession: FirSession)
     override fun getClassifierNames(): Set<Name> = classIndex.keys
 }
 
+class FirCompositeNestedClassifierScope(
+    val scopes: List<FirNestedClassifierScope>,
+    klass: FirClass,
+    useSiteSession: FirSession
+) : FirNestedClassifierScope(klass, useSiteSession) {
+    override fun getNestedClassSymbol(name: Name): FirRegularClassSymbol? {
+        error("Should not be called")
+    }
+
+    override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {
+        scopes.forEach { it.processClassifiersByNameWithSubstitution(name, processor) }
+    }
+
+    override fun isEmpty(): Boolean {
+        return scopes.all { it.isEmpty() }
+    }
+
+    override fun getClassifierNames(): Set<Name> {
+        return scopes.flatMapTo(mutableSetOf()) { it.getClassifierNames() }
+    }
+}
+
 fun FirTypeParameterRef.toConeType(): ConeKotlinType = symbol.toConeType()
 
 fun FirTypeParameterSymbol.toConeType(): ConeKotlinType = ConeTypeParameterTypeImpl(ConeTypeParameterLookupTag(this), isNullable = false)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/scopes/FirCompositeScope.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.scopes
 
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.name.Name
@@ -37,6 +38,20 @@ class FirCompositeScope(val scopes: Iterable<FirScope>) : FirContainingNamesAwar
         }
     }
 
+    private inline fun <T> processComposite(
+        process: FirScope.((T) -> Unit) -> Unit,
+        noinline processor: (T) -> Unit
+    ) {
+        val unique = mutableSetOf<T>()
+        for (scope in scopes) {
+            scope.process {
+                if (unique.add(it)) {
+                    processor(it)
+                }
+            }
+        }
+    }
+
     override fun processFunctionsByName(name: Name, processor: (FirNamedFunctionSymbol) -> Unit) {
         return processComposite(FirScope::processFunctionsByName, name, processor)
     }
@@ -45,6 +60,10 @@ class FirCompositeScope(val scopes: Iterable<FirScope>) : FirContainingNamesAwar
         return processComposite(FirScope::processPropertiesByName, name, processor)
     }
 
+    override fun processDeclaredConstructors(processor: (FirConstructorSymbol) -> Unit) {
+        processComposite(FirScope::processDeclaredConstructors, processor)
+    }
+
     override fun getCallableNames(): Set<Name> {
         return scopes.flatMapTo(hashSetOf()) { it.getContainingCallableNamesIfPresent() }
     }

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

@@ -13,13 +13,11 @@ class FirAllOpenComponentRegistrar : FirExtensionRegistrar() {
         +::AllOpenStatusTransformer
         +::AllOpenVisibilityTransformer
         +::AllOpenSupertypeGenerator
+        +::AllOpenAdditionalCheckers
 
         // Declaration generators
-//        +::AllOpenMemberGenerator
-//        +::AllOpenNestedClassGenerator
-        +::AllOpenAdditionalCheckers
         +::AllOpenTopLevelDeclarationsGenerator
         +::AllOpenClassGenerator
-//        +::AllOpenRecursiveNestedClassGenerator
+        +::AllOpenMembersGenerator
     }
 }

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

@@ -88,24 +88,7 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
             else -> return emptyList()
         }
 
-        val constructor = buildPrimaryConstructor {
-            moduleData = session.moduleData
-            origin = key.origin
-            returnTypeRef = buildResolvedTypeRef {
-                type = ConeClassLikeTypeImpl(
-                    ConeClassLikeLookupTagImpl(classId),
-                    emptyArray(),
-                    isNullable = false
-                )
-            }
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            symbol = FirConstructorSymbol(callableId)
-        }
-        return listOf(constructor.symbol)
+        return listOf(buildConstructor(classId, callableId, isInner = false).symbol)
     }
 
     private val CallableId.isGeneratedConstructor: Boolean
@@ -136,25 +119,7 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
         require(owner is FirRegularClassSymbol)
         val matchedClassId = owner.fir.matchedClass ?: return emptyList()
         val matchedClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(matchedClassId) ?: return emptyList()
-        val function = buildSimpleFunction {
-            moduleData = session.moduleData
-            origin = key.origin
-            status = FirResolvedDeclarationStatusImpl(
-                Visibilities.Public,
-                Modality.FINAL,
-                EffectiveVisibility.Public
-            )
-            returnTypeRef = buildResolvedTypeRef {
-                type = ConeClassLikeTypeImpl(
-                    matchedClassSymbol.toLookupTag(),
-                    emptyArray(),
-                    isNullable = false
-                )
-            }
-            name = MATERIALIZE_NAME
-            symbol = FirNamedFunctionSymbol(callableId)
-        }
-        return listOf(function.symbol)
+        return listOf(buildMaterializeFunction(matchedClassSymbol, callableId).symbol)
     }
 
     private fun buildClass(classId: ClassId): FirRegularClass {
@@ -170,7 +135,7 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
         }
     }
 
-    override fun getCallableNamesForGeneratedClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
         return if (classSymbol.classId in classIdsForMatchedClasses) {
             setOf(MATERIALIZE_NAME)
         } else {
@@ -178,7 +143,7 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
         }
     }
 
-    override fun getNestedClassifiersNamesForGeneratedClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
         return if (classSymbol.classId == GENERATED_CLASS_ID) {
             return classIdsForMatchedClasses.keys.mapTo(mutableSetOf()) { it.shortClassName }
         } else {

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

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.plugin.generators
+
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.descriptors.EffectiveVisibility
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
+import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.has
+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.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
+
+/*
+ * For each class annotated with @C generates
+ *  - member fun materialize(): ClassName
+ *  - nested class Nested with default constructor
+ */
+class AllOpenMembersGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
+    companion object {
+        private val MATERIALIZE_NAME = Name.identifier("materialize")
+        private val NESTED_NAME = Name.identifier("Nested")
+    }
+
+    private val predicateBasedProvider = session.predicateBasedProvider
+    private val matchedClasses by lazy {
+        predicateBasedProvider.getSymbolsByPredicate(predicate).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
+    }
+
+    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+        if (callableId.callableName != MATERIALIZE_NAME) return emptyList()
+        val classId = callableId.classId ?: return emptyList()
+        val matchedClassSymbol = matchedClasses.firstOrNull { it.classId == classId } ?: return emptyList()
+        return listOf(buildMaterializeFunction(matchedClassSymbol, callableId).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
+        return buildRegularClass {
+            moduleData = session.moduleData
+            origin = key.origin
+            classKind = ClassKind.CLASS
+            scopeProvider = session.kotlinScopeProvider
+            status = FirResolvedDeclarationStatusImpl(Visibilities.Public, Modality.FINAL, EffectiveVisibility.Public)
+            name = classId.shortClassName
+            symbol = FirRegularClassSymbol(classId)
+            superTypeRefs += session.builtinTypes.anyType
+        }.symbol
+    }
+
+    override fun generateConstructors(callableId: CallableId): List<FirConstructorSymbol> {
+        val classId = callableId.classId ?: return emptyList()
+        if (callableId.callableName != NESTED_NAME) return emptyList()
+        if (matchedClasses.none { it.classId == classId }) return emptyList()
+        return listOf(buildConstructor(classId.createNestedClassId(NESTED_NAME), callableId, isInner = false).symbol)
+    }
+
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+        return if (classSymbol in matchedClasses) setOf(MATERIALIZE_NAME) else emptySet()
+    }
+
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
+        return if (classSymbol in matchedClasses) setOf(NESTED_NAME) else emptySet()
+    }
+
+    object Key : FirPluginKey() {
+        override fun toString(): String {
+            return "AllOpenMembersGeneratorKey"
+        }
+    }
+
+    override val key: FirPluginKey
+        get() = Key
+    override val predicate: DeclarationPredicate
+        get() = has("C".fqn())
+}

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

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.plugin.generators
+
+import org.jetbrains.kotlin.descriptors.EffectiveVisibility
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.containingClassForStaticMemberAttr
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirConstructor
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.builder.buildPrimaryConstructor
+import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.resolve.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.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+
+fun FirDeclarationGenerationExtension.buildMaterializeFunction(
+    matchedClassSymbol: FirClassLikeSymbol<*>,
+    callableId: CallableId
+): FirSimpleFunction {
+    return buildSimpleFunction {
+        moduleData = session.moduleData
+        origin = key.origin
+        status = FirResolvedDeclarationStatusImpl(
+            Visibilities.Public,
+            Modality.FINAL,
+            EffectiveVisibility.Public
+        )
+        returnTypeRef = buildResolvedTypeRef {
+            type = ConeClassLikeTypeImpl(
+                matchedClassSymbol.toLookupTag(),
+                emptyArray(),
+                isNullable = false
+            )
+        }
+        name = callableId.callableName
+        symbol = FirNamedFunctionSymbol(callableId)
+    }
+}
+
+@OptIn(SymbolInternals::class)
+fun FirDeclarationGenerationExtension.buildConstructor(classId: ClassId, callableId: CallableId, isInner: Boolean): FirConstructor {
+    val lookupTag = ConeClassLikeLookupTagImpl(classId)
+    return buildPrimaryConstructor {
+        moduleData = session.moduleData
+        origin = key.origin
+        returnTypeRef = buildResolvedTypeRef {
+            type = ConeClassLikeTypeImpl(
+                lookupTag,
+                emptyArray(),
+                isNullable = false
+            )
+        }
+        status = FirResolvedDeclarationStatusImpl(
+            Visibilities.Public,
+            Modality.FINAL,
+            EffectiveVisibility.Public
+        )
+        symbol = FirConstructorSymbol(callableId)
+        if (isInner) {
+            dispatchReceiverType = callableId.classId?.let {
+                val firClass = session.symbolProvider.getClassLikeSymbolByClassId(it)?.fir as? FirClass
+                firClass?.defaultType()
+            }
+        }
+    }.also {
+        it.containingClassForStaticMemberAttr = lookupTag
+    }
+}

plugins/fir/fir-plugin-prototype/testData/memberGen/classWithGeneratedMembersAndNestedClass.fir.txt

@@ -0,0 +1,28 @@
+FILE: classWithGeneratedMembersAndNestedClass.kt
+    @R|org/jetbrains/kotlin/fir/plugin/C|() public final class Foo : R|kotlin/Any| {
+        public constructor(): R|Foo| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final class MyNested : R|kotlin/Any| {
+            public constructor(): R|Foo.MyNested| {
+                super<R|kotlin/Any|>()
+            }
+
+        }
+
+    }
+    public final class Bar : R|kotlin/Any| {
+        public constructor(): R|Bar| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final fun test_1(foo: R|Foo|): R|kotlin/Unit| {
+        lval foo2: R|Foo| = R|<local>/foo|.R|/Foo.materialize|()
+        lval nested: R|Foo.Nested| = Q|Foo|.R|/Foo.Nested|()
+    }
+    public final fun test_2(bar: R|Bar|): R|kotlin/Unit| {
+        lval foo2: R|Bar| = R|<local>/bar|.<Unresolved name: materialize>#()
+        lval nested: <ERROR TYPE REF: Unresolved name: Nested> = Q|Bar|.<Unresolved name: Nested>#()
+    }

plugins/fir/fir-plugin-prototype/testData/memberGen/classWithGeneratedMembersAndNestedClass.kt

@@ -0,0 +1,19 @@
+import org.jetbrains.kotlin.fir.plugin.C
+
+@C
+class Foo {
+    class MyNested
+}
+
+class Bar
+
+fun test_1(foo: Foo) {
+    val foo2: Foo = foo.materialize()
+    val nested = Foo.Nested()
+}
+
+// should be errors
+fun test_2(bar: Bar) {
+    val foo2: Bar = bar.<!UNRESOLVED_REFERENCE!>materialize<!>()
+    val nested = Bar.<!UNRESOLVED_REFERENCE!>Nested<!>()
+}

plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelClass.fir.txt

@@ -1,14 +0,0 @@
-FILE: topLevelClass.kt
-    @R|org/jetbrains/kotlin/fir/plugin/A|() public final class SomeClass : R|kotlin/Any| {
-        public constructor(): R|SomeClass| {
-            super<R|kotlin/Any|>()
-        }
-
-    }
-    public final fun test(): R|kotlin/Unit| {
-        Q|TopLevelSomeClass|.<Unresolved name: hello>#()
-    }
-    public final object TopLevelSomeClass {
-        public final fun hello(): R|kotlin/Int|
-
-    }

plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelClass.kt

@@ -1,8 +0,0 @@
-import org.jetbrains.kotlin.fir.plugin.A
-
-@A
-class SomeClass
-
-fun test() {
-    TopLevelSomeClass.hello()
-}

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

@@ -49,6 +49,12 @@ public class FirAllOpenDiagnosticTestGenerated extends AbstractFirAllOpenDiagnos
             KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("plugins/fir/fir-plugin-prototype/testData/memberGen"), Pattern.compile("^(.+)\\.kt$"), null, true);
         }
 
+        @Test
+        @TestMetadata("classWithGeneratedMembersAndNestedClass.kt")
+        public void testClassWithGeneratedMembersAndNestedClass() throws Exception {
+            runTest("plugins/fir/fir-plugin-prototype/testData/memberGen/classWithGeneratedMembersAndNestedClass.kt");
+        }
+
         @Test
         @TestMetadata("generatedClassWithMembersAndNestedClasses.kt")
         public void testGeneratedClassWithMembersAndNestedClasses() throws Exception {
@@ -60,12 +66,6 @@ public class FirAllOpenDiagnosticTestGenerated extends AbstractFirAllOpenDiagnos
         public void testTopLevelCallables() throws Exception {
             runTest("plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelCallables.kt");
         }
-
-        @Test
-        @TestMetadata("topLevelClass.kt")
-        public void testTopLevelClass() throws Exception {
-            runTest("plugins/fir/fir-plugin-prototype/testData/memberGen/topLevelClass.kt");
-        }
     }
 
     @Nested