diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt index db5e1b070bc..53ae5603ad5 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt @@ -7,15 +7,16 @@ package org.jetbrains.kotlin.fir.extensions import com.google.common.collect.LinkedHashMultimap import com.google.common.collect.Multimap +import org.jetbrains.kotlin.descriptors.ClassKind import org.jetbrains.kotlin.fir.FirAnnotationContainer import org.jetbrains.kotlin.fir.FirSession import org.jetbrains.kotlin.fir.FirSessionComponent +import org.jetbrains.kotlin.fir.declarations.FirRegularClass import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall import org.jetbrains.kotlin.fir.resolve.firSymbolProvider import org.jetbrains.kotlin.fir.resolve.getSymbolByTypeRef import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol -import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner -import org.jetbrains.kotlin.fir.utils.ComponentTypeRegistry +import org.jetbrains.kotlin.fir.utils.* import org.jetbrains.kotlin.name.FqName import kotlin.properties.ReadOnlyProperty import kotlin.reflect.KClass @@ -58,6 +59,9 @@ class FirExtensionPointService( val extensionsWithAllInAnnotatedMode = createMultimap() for (extension in extensions) { _metaAnnotations += extension.metaAnnotations + for (metaAnnotation in extension.metaAnnotations) { + extensionsWithMetaAnnotations.put(metaAnnotation, extension) + } _annotations += extension.annotations when (extension.mode) { FirExtensionPoint.Mode.ANNOTATED_ELEMENT -> { @@ -83,6 +87,25 @@ class FirExtensionPointService( ) } + fun registerUserDefinedAnnotation(annotation: FirRegularClass) { + require(annotation.classKind == ClassKind.ANNOTATION_CLASS) + val fqName = annotation.symbol.classId.asSingleFqName() + _annotations += fqName + val extensions = extensionsWithMetaAnnotations[fqName] + if (extensions.isEmpty()) return + for (extension in extensions) { + val registeredExtensions = this[extension::class] + + @Suppress("UNCHECKED_CAST") + val map = when (extension.mode) { + FirExtensionPoint.Mode.ANNOTATED_ELEMENT -> registeredExtensions.extensionsWithAnnotatedMode + FirExtensionPoint.Mode.ALL_IN_ANNOTATED_ELEMENT -> registeredExtensions.extensionsWithAllInAnnotatedMode + FirExtensionPoint.Mode.ALL -> throw IllegalStateException("Extension with mode ALL can't be subscribed to meta annotation") + } as Multimap + map.put(fqName, extension) + } + } + val annotations: Set get() = _annotations private val _annotations: MutableSet = mutableSetOf() @@ -91,6 +114,8 @@ class FirExtensionPointService( get() = _metaAnnotations private val _metaAnnotations: MutableSet = mutableSetOf() + private val extensionsWithMetaAnnotations: Multimap = createMultimap() + class ExtensionsAccessor

( private val session: FirSession, private val extensions: FirRegisteredExtension

diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ComponentArray.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ComponentArray.kt index 450f204cb08..49bbe8b147a 100644 --- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ComponentArray.kt +++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ComponentArray.kt @@ -16,6 +16,10 @@ abstract class ComponentArrayOwner { protected fun registerComponent(tClass: KClass, value: V) { componentArray[(typeRegistry.getId(tClass))] = value } + + protected operator fun get(tClass: KClass): V { + return componentArray[typeRegistry.getId(tClass)] + } } @@ -60,7 +64,7 @@ class ComponentArray { data[index] = value } - operator fun get(index: Int): T? { + operator fun get(index: Int): T { @Suppress("UNCHECKED_CAST") return data.getOrNull(index) as T }