[FIR] Add type resolution service into SerializationFirSupertypesExtension

This service allows resolve custom user types from plugins, which allows
  to support specifying type arguments of generated supertypes basing
  on arguments passed to annotations

build.gradle.kts

@@ -371,9 +371,11 @@ val projectsWithEnabledContextReceivers by extra {
     listOf(
         ":core:descriptors.jvm",
         ":compiler:frontend.common",
+        ":compiler:fir:resolve",
         ":compiler:fir:fir2ir",
         ":kotlin-lombok-compiler-plugin.k1",
         ":kotlinx-serialization-compiler-plugin.k2",
+        ":plugins:fir-plugin-prototype"
     )
 }
 

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

@@ -8,6 +8,7 @@ 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.FirUserTypeRef
 import kotlin.reflect.KClass
 
 abstract class FirSupertypeGenerationExtension(session: FirSession) : FirExtension(session) {
@@ -22,12 +23,21 @@ abstract class FirSupertypeGenerationExtension(session: FirSession) : FirExtensi
 
     abstract fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean
 
+    context(TypeResolveServiceContainer)
+    @Suppress("IncorrectFormatting") // KTIJ-22227
     abstract fun computeAdditionalSupertypes(
         classLikeDeclaration: FirClassLikeDeclaration,
         resolvedSupertypes: List<FirResolvedTypeRef>
     ): List<FirResolvedTypeRef>
 
     fun interface Factory : FirExtension.Factory<FirSupertypeGenerationExtension>
+
+    class TypeResolveServiceContainer(val typeResolver: TypeResolveService)
+
+    abstract class TypeResolveService {
+        abstract fun resolveUserType(type: FirUserTypeRef): FirResolvedTypeRef
+    }
 }
 
 val FirExtensionService.supertypeGenerators: List<FirSupertypeGenerationExtension> by FirExtensionService.registeredExtensions()
+

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.superConeTypes
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.supertypeGenerators
 import org.jetbrains.kotlin.fir.resolve.*
@@ -421,22 +422,39 @@ open class FirSupertypeResolverVisitor(
                         superTypeRef
                 }
             }.also {
-                addSupertypesFromExtensions(classLikeDeclaration, it)
+                addSupertypesFromExtensions(classLikeDeclaration, it, transformer, scopeDeclaration)
             }
         }
     }
 
     private fun <T> List<T>.createCopy(): List<T> = ArrayList(this)
 
-    private fun addSupertypesFromExtensions(klass: FirClassLikeDeclaration, supertypeRefs: MutableList<FirResolvedTypeRef>) {
+    private fun addSupertypesFromExtensions(
+        klass: FirClassLikeDeclaration,
+        supertypeRefs: MutableList<FirResolvedTypeRef>,
+        typeResolveTransformer: FirTransformer<ScopeClassDeclaration>,
+        scopeDeclaration: ScopeClassDeclaration
+    ) {
         if (supertypeGenerationExtensions.isEmpty()) return
-        for (extension in supertypeGenerationExtensions) {
-            if (extension.needTransformSupertypes(klass)) {
-                supertypeRefs += extension.computeAdditionalSupertypes(klass, supertypeRefs)
+        val typeResolveService = TypeResolveServiceForPlugins(typeResolveTransformer, scopeDeclaration)
+        with(FirSupertypeGenerationExtension.TypeResolveServiceContainer(typeResolveService)) {
+            for (extension in supertypeGenerationExtensions) {
+                if (extension.needTransformSupertypes(klass)) {
+                    supertypeRefs += extension.computeAdditionalSupertypes(klass, supertypeRefs)
+                }
             }
         }
     }
 
+    private class TypeResolveServiceForPlugins(
+        val typeResolveTransformer: FirTransformer<ScopeClassDeclaration>,
+        val scopeDeclaration: ScopeClassDeclaration
+    ) : FirSupertypeGenerationExtension.TypeResolveService() {
+        override fun resolveUserType(type: FirUserTypeRef): FirResolvedTypeRef {
+            return type.transform(typeResolveTransformer, scopeDeclaration)
+        }
+    }
+
     override fun visitTypeAlias(typeAlias: FirTypeAlias, data: Any?) {
         // TODO: this if is a temporary hack for built-in types (because we can't load file for them)
         if (typeAlias.expandedTypeRef is FirResolvedTypeRef) {

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

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.fir.plugin
 
+import kotlin.reflect.KClass
+
 annotation class AllOpen
 
 annotation class DummyFunction
@@ -27,3 +29,5 @@ annotation class Negative
 enum class Visibility {
     Public, Internal, Private, Protected
 }
+
+annotation class SupertypeWithTypeArgument(val kClass: KClass<*>)

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

@@ -19,6 +19,7 @@ class FirPluginPrototypeExtensionRegistrar : FirExtensionRegistrar() {
         +::AllOpenStatusTransformer
         +::AllPublicVisibilityTransformer
         +::SomeAdditionalSupertypeGenerator
+        +::SupertypeWithArgumentGenerator
         +::PluginAdditionalCheckers
         +::FirNumberSignAttributeExtension
         +::AlgebraReceiverInjector

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

@@ -30,6 +30,8 @@ class SomeAdditionalSupertypeGenerator(session: FirSession) : FirSupertypeGenera
 
     }
 
+    context(TypeResolveServiceContainer)
+    @Suppress("IncorrectFormatting") // KTIJ-22227
     override fun computeAdditionalSupertypes(
         classLikeDeclaration: FirClassLikeDeclaration,
         resolvedSupertypes: List<FirResolvedTypeRef>

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

@@ -0,0 +1,78 @@
+/*
+ * 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.plugin
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.buildUserTypeFromQualifierParts
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.annotated
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.references.impl.FirSimpleNamedReference
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.constructClassLikeType
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+/*
+ * Adds MyInterface supertype for all classes annotated with @MyInterfaceSupertype
+ */
+class SupertypeWithArgumentGenerator(session: FirSession) : FirSupertypeGenerationExtension(session) {
+    companion object {
+        private val supertypeClassId = ClassId(FqName("foo"), Name.identifier("InterfaceWithArgument"))
+        private val annotationClassId = ClassId.topLevel("SupertypeWithTypeArgument".fqn())
+        private val PREDICATE: DeclarationPredicate = annotated(annotationClassId.asSingleFqName())
+
+    }
+
+    context(TypeResolveServiceContainer)
+    @Suppress("IncorrectFormatting") // KTIJ-22227
+    override fun computeAdditionalSupertypes(
+        classLikeDeclaration: FirClassLikeDeclaration,
+        resolvedSupertypes: List<FirResolvedTypeRef>
+    ): List<FirResolvedTypeRef> {
+        if (resolvedSupertypes.any { it.type.classId == supertypeClassId }) return emptyList()
+
+        val annotation = classLikeDeclaration.getAnnotationByClassId(annotationClassId) ?: return emptyList()
+        val getClassArgument = (annotation as? FirAnnotationCall)?.argument as? FirGetClassCall ?: return emptyList()
+
+        val typeToResolve = buildUserTypeFromQualifierParts(isMarkedNullable = false) {
+            fun visitQualifiers(expression: FirExpression) {
+                if (expression !is FirPropertyAccessExpression) return
+                expression.explicitReceiver?.let { visitQualifiers(it) }
+                expression.qualifierName?.let { part(it) }
+            }
+            visitQualifiers(getClassArgument.argument)
+        }
+
+        val resolvedArgument = typeResolver.resolveUserType(typeToResolve).type
+
+        return listOf(
+            buildResolvedTypeRef {
+                type = supertypeClassId.constructClassLikeType(arrayOf(resolvedArgument), isNullable = false)
+            }
+        )
+    }
+
+    private val FirPropertyAccessExpression.qualifierName: Name?
+        get() = (calleeReference as? FirSimpleNamedReference)?.name
+
+    override fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean {
+        return session.predicateBasedProvider.matches(PREDICATE, declaration)
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
+}

plugins/fir-plugin-prototype/testData/diagnostics/supertypes/supertypeWithArgument.fir.txt

@@ -0,0 +1,75 @@
+FILE: supertypeWithArgument.kt
+    package foo
+
+    public abstract interface InterfaceWithArgument<T> : R|kotlin/Any| {
+        public open fun generate(): R|T| {
+            ^generate Null(null)!!
+        }
+
+    }
+    @R|org/jetbrains/kotlin/fir/plugin/SupertypeWithTypeArgument|(kClass = <getClass>(Q|kotlin/String|)) public final class TopLevelType : R|kotlin/Any|, R|foo/InterfaceWithArgument<kotlin/String>| {
+        public constructor(): R|foo/TopLevelType| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final class Nested : R|kotlin/Any| {
+        public constructor(): R|foo/Nested| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public abstract class Base : R|kotlin/Any| {
+        public constructor(): R|foo/Base| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final class Nested : R|kotlin/Any| {
+            public constructor(): R|foo/Base.Nested| {
+                super<R|kotlin/Any|>()
+            }
+
+        }
+
+    }
+    public final class Derived : R|foo/Base| {
+        public constructor(): R|foo/Derived| {
+            super<R|foo/Base|>()
+        }
+
+        @R|org/jetbrains/kotlin/fir/plugin/SupertypeWithTypeArgument|(kClass = <getClass>(Q|foo/Base.Nested|)) public final class TypeFromSupertype : R|kotlin/Any|, R|foo/InterfaceWithArgument<foo/Base.Nested>| {
+            public constructor(): R|foo/Derived.TypeFromSupertype| {
+                super<R|kotlin/Any|>()
+            }
+
+        }
+
+        @R|org/jetbrains/kotlin/fir/plugin/SupertypeWithTypeArgument|(kClass = <getClass>(Q|foo/Nested|)) public final class QualifiedType : R|kotlin/Any|, R|foo/InterfaceWithArgument<foo/Nested>| {
+            public constructor(): R|foo/Derived.QualifiedType| {
+                super<R|kotlin/Any|>()
+            }
+
+        }
+
+    }
+    @R|org/jetbrains/kotlin/fir/plugin/SupertypeWithTypeArgument|(kClass = <getClass>(<Unresolved name: UnresolvedClass>#)) public final class UnresolvedType : R|kotlin/Any|, R|foo/InterfaceWithArgument<ERROR CLASS: Symbol not found for UnresolvedClass>| {
+        public constructor(): R|foo/UnresolvedType| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final fun takeString(x: R|kotlin/String|): R|kotlin/Unit| {
+    }
+    public final fun takeTopLevelNested(x: R|foo/Base.Nested|): R|kotlin/Unit| {
+    }
+    public final fun takeInnerNested(x: R|foo/Nested|): R|kotlin/Unit| {
+    }
+    public final fun test_1(x: R|foo/TopLevelType|): R|kotlin/Unit| {
+        R|foo/takeString|(R|<local>/x|.R|SubstitutionOverride<foo/TopLevelType.generate: R|kotlin/String|>|())
+    }
+    public final fun test_2(x: R|foo/Derived.TypeFromSupertype|): R|kotlin/Unit| {
+        R|foo/takeTopLevelNested|(R|<local>/x|.R|SubstitutionOverride<foo/Derived.TypeFromSupertype.generate: R|foo/Base.Nested|>|())
+    }
+    public final fun test_3(x: R|foo/Derived.QualifiedType|): R|kotlin/Unit| {
+        R|foo/takeInnerNested|(R|<local>/x|.R|SubstitutionOverride<foo/Derived.QualifiedType.generate: R|foo/Nested|>|())
+    }

plugins/fir-plugin-prototype/testData/diagnostics/supertypes/supertypeWithArgument.kt

@@ -0,0 +1,43 @@
+package foo
+
+import org.jetbrains.kotlin.fir.plugin.SupertypeWithTypeArgument
+
+interface InterfaceWithArgument<T> {
+    fun generate(): T = null!!
+}
+
+@SupertypeWithTypeArgument(String::class)
+class TopLevelType
+
+class Nested
+
+abstract class Base {
+    class Nested
+}
+
+class Derived : Base() {
+    @SupertypeWithTypeArgument(Nested::class)
+    class TypeFromSupertype
+
+    @SupertypeWithTypeArgument(foo.Nested::class)
+    class QualifiedType
+}
+
+@SupertypeWithTypeArgument(<!ANNOTATION_ARGUMENT_MUST_BE_CONST!><!UNRESOLVED_REFERENCE!>UnresolvedClass<!>::class<!>)
+class UnresolvedType
+
+fun takeString(x: String) {}
+fun takeTopLevelNested(x: Base.Nested) {}
+fun takeInnerNested(x: Nested) {}
+
+fun test_1(x: TopLevelType) {
+    takeString(x.generate())
+}
+
+fun test_2(x: Derived.TypeFromSupertype) {
+    takeTopLevelNested(x.generate())
+}
+
+fun test_3(x: Derived.QualifiedType) {
+    takeInnerNested(x.generate())
+}

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

@@ -138,5 +138,11 @@ public class FirPluginDiagnosticTestGenerated extends AbstractFirPluginDiagnosti
         public void testSimple() throws Exception {
             runTest("plugins/fir-plugin-prototype/testData/diagnostics/supertypes/simple.kt");
         }
+
+        @Test
+        @TestMetadata("supertypeWithArgument.kt")
+        public void testSupertypeWithArgument() throws Exception {
+            runTest("plugins/fir-plugin-prototype/testData/diagnostics/supertypes/supertypeWithArgument.kt");
+        }
     }
 }

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

@@ -26,6 +26,8 @@ class SerializationFirSupertypesExtension(session: FirSession) : FirSupertypeGen
         register(serializerFor)
     }
 
+    context(TypeResolveServiceContainer)
+    @Suppress("IncorrectFormatting") // KTIJ-22227
     override fun computeAdditionalSupertypes(
         classLikeDeclaration: FirClassLikeDeclaration,
         resolvedSupertypes: List<FirResolvedTypeRef>
@@ -39,4 +41,4 @@ class SerializationFirSupertypesExtension(session: FirSession) : FirSupertypeGen
             TODO("Support @Serializer(for=...) supertype generation")
         } else emptyList()
     }
-}
+}