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

2021-09-24 16:54:21 +03:00
parent 270962e176
commit f3a9d70eb6
22 changed files with 515 additions and 183 deletions
@@ -6,13 +6,11 @@
 package org.jetbrains.kotlin.fir.extensions

 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
-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.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
 import kotlin.reflect.KClass

 /*
@@ -29,11 +27,20 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredi

    final override val extensionType: KClass<out FirExtension> = FirDeclarationGenerationExtension::class

-    open fun generateClassLikeDeclaration(classId: ClassId, owner: FirClassSymbol<*>?): FirClassLikeSymbol<*>? = null
+    // Can be called on SUPERTYPES stage
+    open fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
+
+    // Can be called on STATUS stage
    open fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> = emptyList()
    open fun generateProperties(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirPropertySymbol> = emptyList()
+    open fun generateConstructors(callableId: CallableId): List<FirConstructorSymbol> = emptyList()
+
+    // 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()
+
    fun interface Factory : FirExtension.Factory<FirDeclarationGenerationExtension>
 }

@@ -9,13 +9,11 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.validate
 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.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -56,7 +54,7 @@ class FirExtensionDeclarationsSymbolProvider private constructor(

    private fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
        // TODO: what we should do if multiple extensions want to generate class with same classId?
-        return extensions.firstNotNullOfOrNull { it.generateClassLikeDeclaration(classId, owner = null) }?.also { it.fir.validate() }
+        return extensions.firstNotNullOfOrNull { it.generateClassLikeDeclaration(classId) }?.also { it.fir.validate() }
    }

    private fun generateTopLevelFunctions(callableId: CallableId): List<FirNamedFunctionSymbol> {
@@ -74,23 +72,24 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
    // ------------------------------------------ provider methods ------------------------------------------

    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
-        return classCache.getValue(classId, context = null)
+        return classCache.getValue(classId)
    }

    @FirSymbolProviderInternals
    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
-        destination += functionCache.getValue(CallableId(packageFqName, name), context = null)
-        destination += propertyCache.getValue(CallableId(packageFqName, name), context = null)
+        val callableId = CallableId(packageFqName, name)
+        destination += functionCache.getValue(callableId)
+        destination += propertyCache.getValue(callableId)
    }

    @FirSymbolProviderInternals
    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
-        destination += functionCache.getValue(CallableId(packageFqName, name), context = null)
+        destination += functionCache.getValue(CallableId(packageFqName, name))
    }

    @FirSymbolProviderInternals
    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
-        destination += propertyCache.getValue(CallableId(packageFqName, name), context = null)
+        destination += propertyCache.getValue(CallableId(packageFqName, name))
    }

    override fun getPackage(fqName: FqName): FqName? {
@@ -17,13 +17,15 @@ import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames

-class FirClassDeclaredMemberScope(
+abstract class FirClassDeclaredMemberScope : FirScope(), FirContainingNamesAwareScope
+
+class FirClassDeclaredMemberScopeImpl(
    val useSiteSession: FirSession,
    klass: FirClass,
    useLazyNestedClassifierScope: Boolean = false,
    existingNames: List<Name>? = null,
    symbolProvider: FirSymbolProvider? = null
-) : FirScope(), FirContainingNamesAwareScope {
+) : FirClassDeclaredMemberScope() {
    private val nestedClassifierScope: FirScope? = if (useLazyNestedClassifierScope) {
        lazyNestedClassifierScope(klass.symbol.classId, existingNames!!, symbolProvider!!)
    } else {
@@ -50,7 +50,11 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
        existingNames: List<Name>?,
        symbolProvider: FirSymbolProvider?
    ): FirClassDeclaredMemberScope {
-        return FirClassDeclaredMemberScope(useSiteSession, klass, useLazyNestedClassifierScope, existingNames, symbolProvider)
+        return if (klass.origin.generated) {
+            FirGeneratedClassDeclaredMemberScope(useSiteSession, klass)
+        } else {
+            FirClassDeclaredMemberScopeImpl(useSiteSession, klass, useLazyNestedClassifierScope, existingNames, symbolProvider)
+        }
    }

    fun nestedClassifierScope(klass: FirClass): FirNestedClassifierScope? {
@@ -58,7 +62,11 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
    }

    private fun createNestedClassifierScope(klass: FirClass): FirNestedClassifierScope? {
-        return FirNestedClassifierScope(klass, useSiteSession).takeUnless { it.isEmpty() }
+        return if (klass.origin.generated) {
+            FirGeneratedClassNestedClassifierScope(klass, useSiteSession)
+        } else {
+            FirNestedClassifierScopeImpl(klass, useSiteSession)
+        }.takeUnless { it.isEmpty() }
    }
 }

@@ -0,0 +1,154 @@
+/*
+ * 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.scopes.impl
+
+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.caches.getValue
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+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.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
+
+class FirGeneratedClassDeclaredMemberScope(
+    val useSiteSession: FirSession,
+    val firClass: FirClass
+) : FirClassDeclaredMemberScope() {
+    private val extension: FirDeclarationGenerationExtension = firClass.findGeneratedExtension(useSiteSession)
+    private val nestedClassifierScope: FirNestedClassifierScope? = useSiteSession.nestedClassifierScope(firClass)
+
+    private val firCachesFactory = useSiteSession.firCachesFactory
+
+    // ------------------------------------------ caches ------------------------------------------
+
+    private val functionCache: FirCache<Name, List<FirNamedFunctionSymbol>, Nothing?> = firCachesFactory.createCache { callableId, _ ->
+        generateMemberFunctions(callableId)
+    }
+
+    private val propertyCache: FirCache<Name, List<FirPropertySymbol>, Nothing?> = firCachesFactory.createCache { callableId, _ ->
+        generateMemberProperties(callableId)
+    }
+
+    private val constructorCache: FirLazyValue<List<FirConstructorSymbol>, Nothing?> = firCachesFactory.createLazyValue {
+        generateConstructors()
+    }
+
+    private val callableNamesCache: FirLazyValue<Set<Name>, Nothing?> = firCachesFactory.createLazyValue {
+        extension.getCallableNamesForGeneratedClass(firClass.symbol)
+    }
+
+    // ------------------------------------------ generators ------------------------------------------
+
+    private fun generateMemberFunctions(name: Name): List<FirNamedFunctionSymbol> {
+        return extension.generateFunctions(CallableId(firClass.classId, name), firClass.symbol)
+    }
+
+    private fun generateMemberProperties(name: Name): List<FirPropertySymbol> {
+        return extension.generateProperties(CallableId(firClass.classId, name), firClass.symbol)
+    }
+
+    private fun generateConstructors(): List<FirConstructorSymbol> {
+        val classId = firClass.symbol.classId
+        val callableId = if (classId.isNestedClass) {
+            CallableId(classId.parentClassId!!, classId.shortClassName)
+        } else {
+            CallableId(classId.asSingleFqName().parent(), classId.shortClassName)
+        }
+        return extension.generateConstructors(callableId)
+    }
+
+    // ------------------------------------------ scope methods ------------------------------------------
+
+    override fun getCallableNames(): Set<Name> {
+        return callableNamesCache.getValue()
+    }
+
+    override fun getClassifierNames(): Set<Name> {
+        return nestedClassifierScope?.getClassifierNames() ?: emptySet()
+    }
+
+    override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {
+        nestedClassifierScope?.processClassifiersByNameWithSubstitution(name, processor)
+    }
+
+    override fun processFunctionsByName(name: Name, processor: (FirNamedFunctionSymbol) -> Unit) {
+        if (name !in getCallableNames()) return
+        for (functionSymbol in functionCache.getValue(name)) {
+            processor(functionSymbol)
+        }
+    }
+
+    override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
+        if (name !in getCallableNames()) return
+        for (propertySymbol in propertyCache.getValue(name)) {
+            processor(propertySymbol)
+        }
+    }
+
+    override fun processDeclaredConstructors(processor: (FirConstructorSymbol) -> Unit) {
+        for (constructorSymbol in constructorCache.getValue()) {
+            processor(constructorSymbol)
+        }
+    }
+}
+
+class FirGeneratedClassNestedClassifierScope(
+    klass: FirClass,
+    useSiteSession: FirSession
+) : FirNestedClassifierScope(klass, useSiteSession) {
+    private val extension = klass.findGeneratedExtension(useSiteSession)
+
+    private val nestedClassifierCache: FirCache<Name, FirRegularClassSymbol?, Nothing?> =
+        useSiteSession.firCachesFactory.createCache { name, _ ->
+            generateNestedClassifier(name)
+        }
+
+    private val nestedClassifiersNames: FirLazyValue<Set<Name>, Nothing?> =
+        useSiteSession.firCachesFactory.createLazyValue {
+            extension.getNestedClassifiersNamesForGeneratedClass(klass.symbol)
+        }
+
+    private fun generateNestedClassifier(name: Name): FirRegularClassSymbol? {
+        if (name !in getClassifierNames()) return null
+        val generatedClass = useSiteSession.symbolProvider.getClassLikeSymbolByClassId(klass.classId.createNestedClassId(name))
+        require(generatedClass is FirRegularClassSymbol?) { "Only regular class are allowed as nested classes" }
+        return generatedClass
+    }
+
+    override fun getNestedClassSymbol(name: Name): FirRegularClassSymbol? {
+        return nestedClassifierCache.getValue(name)
+    }
+
+    override fun isEmpty(): Boolean {
+        return getClassifierNames().isEmpty()
+    }
+
+    override fun getClassifierNames(): Set<Name> {
+        return nestedClassifiersNames.getValue()
+    }
+}
+
+
+private fun FirClass.findGeneratedExtension(useSiteSession: FirSession): FirDeclarationGenerationExtension {
+    val origin = origin
+    require(origin is FirDeclarationOrigin.Plugin) {
+        "GeneratedClassDeclaredMemberScope can not be created for non-generated class: ${this.render()}"
+    }
+
+    return useSiteSession.extensionService.declarationGenerators.firstOrNull {
+        it.key == origin.key
+    } ?: error("Extension for ${origin.key} not found")
+}
@@ -0,0 +1,16 @@
+/*
+ * 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.scopes.impl
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.name.Name
+
@@ -21,7 +21,26 @@ import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.name.Name

-class FirNestedClassifierScope(val klass: FirClass, val useSiteSession: FirSession) : FirScope(), FirContainingNamesAwareScope {
+abstract class FirNestedClassifierScope(val klass: FirClass, val useSiteSession: FirSession) : FirScope(), FirContainingNamesAwareScope {
+    protected abstract fun getNestedClassSymbol(name: Name): FirRegularClassSymbol?
+
+    override fun processClassifiersByNameWithSubstitution(
+        name: Name,
+        processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit
+    ) {
+        val matchedClass = getNestedClassSymbol(name) ?: return
+        val substitution = klass.typeParameters.associate {
+            it.symbol to it.toConeType()
+        }
+        processor(matchedClass, ConeSubstitutorByMap(substitution, useSiteSession))
+    }
+
+    abstract fun isEmpty(): Boolean
+
+    override fun getCallableNames(): Set<Name> = emptySet()
+}
+
+class FirNestedClassifierScopeImpl(klass: FirClass, useSiteSession: FirSession) : FirNestedClassifierScope(klass, useSiteSession) {
    private val classIndex: Map<Name, FirRegularClassSymbol> = run {
        val result = mutableMapOf<Name, FirRegularClassSymbol>()
        for (declaration in klass.declarations) {
@@ -32,22 +51,13 @@ class FirNestedClassifierScope(val klass: FirClass, val useSiteSession: FirSessi
        result
    }

-    fun isEmpty() = classIndex.isEmpty()
-
-    override fun processClassifiersByNameWithSubstitution(
-        name: Name,
-        processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit
-    ) {
-        val matchedClass = classIndex[name] ?: return
-        val substitution = klass.typeParameters.associate {
-            it.symbol to it.toConeType()
-        }
-        processor(matchedClass, ConeSubstitutorByMap(substitution, useSiteSession))
+    override fun getNestedClassSymbol(name: Name): FirRegularClassSymbol? {
+        return classIndex[name]
    }

-    override fun getClassifierNames(): Set<Name> = classIndex.keys
+    override fun isEmpty(): Boolean = classIndex.isEmpty()

-    override fun getCallableNames(): Set<Name> = emptySet()
+    override fun getClassifierNames(): Set<Name> = classIndex.keys
 }

 fun FirTypeParameterRef.toConeType(): ConeKotlinType = symbol.toConeType()
@@ -17,3 +17,14 @@ inline fun <K : Any, V> FirCache<K, V, Nothing?>.getValue(key: K): V =
 operator fun <K : Any, V> FirCache<K, V, Nothing>.contains(key: K): Boolean {
    return getValueIfComputed(key) != null
 }
+
+class FirLazyValue<out V, in CONTEXT>(private val cache: FirCache<Unit, V, CONTEXT>) {
+    fun getValue(context: CONTEXT): V {
+        return cache.getValue(Unit, context)
+    }
+}
+
+@Suppress("NOTHING_TO_INLINE")
+inline fun <V> FirLazyValue<V, Nothing?>.getValue(): V {
+    return getValue(null)
+}
@@ -37,6 +37,10 @@ abstract class FirCachesFactory : FirSessionComponent {
        createValue: (K, CONTEXT) -> Pair<V, DATA>,
        postCompute: (K, V, DATA) -> Unit
    ): FirCache<K, V, CONTEXT>
+
+    fun <V, CONTEXT> createLazyValue(createValue: (CONTEXT) -> V): FirLazyValue<V, CONTEXT> {
+        return FirLazyValue(createCache { _, context -> createValue(context) })
+    }
 }

 val FirSession.firCachesFactory: FirCachesFactory by FirSession.sessionComponentAccessor()
@@ -5,7 +5,7 @@

 package org.jetbrains.kotlin.fir.declarations

-sealed class FirDeclarationOrigin(private val displayName: String? = null, val fromSupertypes: Boolean = false) {
+sealed class FirDeclarationOrigin(private val displayName: String? = null, val fromSupertypes: Boolean = false, val generated: Boolean = false) {
    object Source : FirDeclarationOrigin()
    object Library : FirDeclarationOrigin()
    object BuiltIns : FirDeclarationOrigin()
@@ -18,7 +18,7 @@ sealed class FirDeclarationOrigin(private val displayName: String? = null, val f
    object IntersectionOverride : FirDeclarationOrigin(fromSupertypes = true)
    object Delegated : FirDeclarationOrigin()

-    class Plugin(val key: FirPluginKey) : FirDeclarationOrigin(displayName = "Plugin[$key]")
+    class Plugin(val key: FirPluginKey) : FirDeclarationOrigin(displayName = "Plugin[$key]", generated = true)

    override fun toString(): String {
        return displayName ?: this::class.simpleName!!