[FIR] Add extension for adding new supertypes

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

@@ -17,7 +17,8 @@ abstract class FirExtensionRegistrar {
         val AVAILABLE_EXTENSIONS = listOf(
             FirStatusTransformerExtension::class,
             FirDeclarationGenerationExtension::class,
-            AbstractFirAdditionalCheckersExtension::class
+            AbstractFirAdditionalCheckersExtension::class,
+            FirSupertypeGenerationExtension::class
         )
     }
 
@@ -41,6 +42,11 @@ abstract class FirExtensionRegistrar {
                 AbstractFirAdditionalCheckersExtension.Factory { this.invoke(it) }
             )
         }
+
+        @JvmName("plusSupertypeGenerationExtension")
+        operator fun ((FirSession) -> FirSupertypeGenerationExtension).unaryPlus() {
+            registerExtension(FirSupertypeGenerationExtension::class, FirSupertypeGenerationExtension.Factory { this.invoke(it) })
+        }
     }
 
     @OptIn(PluginServicesInitialization::class)

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2010-2020 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.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import kotlin.reflect.KClass
+
+abstract class FirSupertypeGenerationExtension(session: FirSession) : FirPredicateBasedExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("SupertypeGenerator")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    override val extensionType: KClass<out FirExtension> = FirSupertypeGenerationExtension::class
+
+    abstract fun computeAdditionalSupertypes(
+        classLikeDeclaration: FirClassLikeDeclaration<*>,
+        resolvedSupertypes: List<FirTypeRef>
+    ): List<FirResolvedTypeRef>
+
+    fun interface Factory : FirExtension.Factory<FirSupertypeGenerationExtension>
+}
+
+val FirExtensionService.supertypeGenerators: List<FirSupertypeGenerationExtension> by FirExtensionService.registeredExtensions()

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

@@ -10,6 +10,9 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.extensions.supertypeGenerators
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.providers.getNestedClassifierScope
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.LocalClassesNavigationInfo
@@ -167,6 +170,8 @@ private class FirSupertypeResolverVisitor(
     private val scopeForLocalClass: FirImmutableCompositeScope? = null,
     private val localClassesNavigationInfo: LocalClassesNavigationInfo? = null
 ) : FirDefaultVisitorVoid() {
+    private val supertypeGenerationExtensions = session.extensionService.supertypeGenerators
+
     override fun visitElement(element: FirElement) {}
 
     private fun prepareFileScope(file: FirFile): FirImmutableCompositeScope {
@@ -278,7 +283,7 @@ private class FirSupertypeResolverVisitor(
         }
 
         return resolveSpecificClassLikeSupertypes(classLikeDeclaration) { transformer ->
-            supertypeRefs.map {
+            supertypeRefs.mapTo(mutableListOf()) {
                 val superTypeRef = transformer.transformTypeRef(it, null).single
 
                 if (superTypeRef.coneTypeSafe<ConeTypeParameterType>() != null)
@@ -288,6 +293,18 @@ private class FirSupertypeResolverVisitor(
                     )
                 else
                     superTypeRef
+            }.also {
+                addSupertypesFromExtensions(classLikeDeclaration, it)
+            }
+        }
+    }
+
+    private fun addSupertypesFromExtensions(klass: FirClassLikeDeclaration<*>, supertypeRefs: MutableList<FirTypeRef>) {
+        if (supertypeGenerationExtensions.isEmpty()) return
+        val provider = session.predicateBasedProvider
+        for (extension in supertypeGenerationExtensions) {
+            if (provider.matches(extension.predicate, klass)) {
+                supertypeRefs += extension.computeAdditionalSupertypes(klass, supertypeRefs)
             }
         }
     }