JVM IR: introduce ClassGeneratorExtension

This is a low-level extension point for Kotlin/JVM, which is supposed to
be used instead of ClassBuilderInterceptorExtension.

 #KT-54758 Fixed
 #KT-56814 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/state/GenerationState.kt

@@ -397,7 +397,7 @@ class GenerationState private constructor(
                         ?.let { destination -> SignatureDumpingBuilderFactory(it, File(destination)) } ?: it
                 }
             )
-            .wrapWith(ClassBuilderInterceptorExtension.getInstances(project)) { classBuilderFactory, extension ->
+            .wrapWith(loadClassBuilderInterceptors()) { classBuilderFactory, extension ->
                 extension.interceptClassBuilderFactory(classBuilderFactory, originalFrontendBindingContext, diagnostics)
             }
 
@@ -405,6 +405,16 @@ class GenerationState private constructor(
         this.factory = ClassFileFactory(this, interceptedBuilderFactory, finalizers)
     }
 
+    @Suppress("UNCHECKED_CAST")
+    private fun loadClassBuilderInterceptors(): List<ClassBuilderInterceptorExtension> {
+        // Using Class.forName here because we're in the old JVM backend, and we need to load extensions declared in the JVM IR backend.
+        val adapted = Class.forName("org.jetbrains.kotlin.backend.jvm.extensions.ClassBuilderExtensionAdapter")
+            .getDeclaredMethod("getExtensions", Project::class.java)
+            .invoke(null, project) as List<ClassBuilderInterceptorExtension>
+
+        return ClassBuilderInterceptorExtension.getInstances(project) + adapted
+    }
+
     fun beforeCompile() {
         markUsed()
     }

compiler/cli/cli-common/resources/META-INF/extensions/compiler.xml

@@ -41,6 +41,9 @@
         <extensionPoint qualifiedName="org.jetbrains.kotlin.classBuilderFactoryInterceptorExtension"
                         interface="org.jetbrains.kotlin.codegen.extensions.ClassBuilderInterceptorExtension"
                         area="IDEA_PROJECT" dynamic="true"/>
+        <extensionPoint qualifiedName="org.jetbrains.kotlin.classGeneratorExtension"
+                        interface="org.jetbrains.kotlin.backend.jvm.extensions.ClassGeneratorExtension"
+                        area="IDEA_PROJECT" dynamic="true"/>
         <extensionPoint qualifiedName="org.jetbrains.kotlin.packageFragmentProviderExtension"
                         interface="org.jetbrains.kotlin.resolve.jvm.extensions.PackageFragmentProviderExtension"
                         area="IDEA_PROJECT" dynamic="true"/>

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt

@@ -47,6 +47,7 @@ import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.jvm.extensions.ClassGeneratorExtension
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.CliModuleVisibilityManagerImpl
 import org.jetbrains.kotlin.cli.common.CompilerSystemProperties
@@ -637,6 +638,7 @@ class KotlinCoreEnvironment private constructor(
             SyntheticResolveExtension.registerExtensionPoint(project)
             SyntheticJavaResolveExtension.registerExtensionPoint(project)
             ClassBuilderInterceptorExtension.registerExtensionPoint(project)
+            ClassGeneratorExtension.registerExtensionPoint(project)
             ClassFileFactoryFinalizerExtension.registerExtensionPoint(project)
             AnalysisHandlerExtension.registerExtensionPoint(project)
             PackageFragmentProviderExtension.registerExtensionPoint(project)

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt

@@ -5,13 +5,11 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.backend.common.lower.ANNOTATION_IMPLEMENTATION
-import org.jetbrains.kotlin.backend.common.psi.PsiSourceManager
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.MultifileFacadeFileEntry
-import org.jetbrains.kotlin.backend.jvm.extensions.JvmIrDeclarationOrigin
+import org.jetbrains.kotlin.backend.jvm.extensions.descriptorOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.*
 import org.jetbrains.kotlin.backend.jvm.mapping.IrTypeMapper
 import org.jetbrains.kotlin.backend.jvm.mapping.MethodSignatureMapper
@@ -57,9 +55,7 @@ import org.jetbrains.kotlin.name.JvmNames.VOLATILE_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtPropertyDelegate
 import org.jetbrains.kotlin.resolve.jvm.checkers.JvmSimpleNameBacktickChecker
-import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOriginKind
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.MemberKind
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.RawSignature
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmClassSignature
@@ -527,41 +523,6 @@ class ClassCodegen private constructor(
         }
     }
 
-    private val IrDeclaration.descriptorOrigin: JvmIrDeclarationOrigin
-        get() {
-            val psiElement = findPsiElementForDeclarationOrigin()
-            return when {
-                origin == IrDeclarationOrigin.FILE_CLASS ->
-                    JvmIrDeclarationOrigin(JvmDeclarationOriginKind.PACKAGE_PART, psiElement, this)
-                (this is IrSimpleFunction && isSuspend && isEffectivelyInlineOnly()) ||
-                        origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE ||
-                        origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE ->
-                    JvmIrDeclarationOrigin(JvmDeclarationOriginKind.INLINE_VERSION_OF_SUSPEND_FUN, psiElement, this)
-                else -> JvmIrDeclarationOrigin(JvmDeclarationOriginKind.OTHER, psiElement, this)
-            }
-        }
-
-    private fun IrDeclaration.findPsiElementForDeclarationOrigin(): PsiElement? {
-        // For synthetic $annotations methods for properties, use the PSI for the property or the constructor parameter.
-        // It's used in KAPT stub generation to sort the properties correctly based on their source position (see KT-44130).
-        if (this is IrFunction && name.asString().endsWith("\$annotations")) {
-            val metadata = metadata as? DescriptorMetadataSource.Property
-            if (metadata != null) {
-                return metadata.descriptor.psiElement
-            }
-        }
-
-        val element = PsiSourceManager.findPsiElement(this)
-
-        // Offsets for accessors and field of delegated property in IR point to the 'by' keyword, so the closest PSI element is the
-        // KtPropertyDelegate (`by ...` expression). However, old JVM backend passed the PSI element of the property instead.
-        // This is important for example in case of KAPT stub generation in the "correct error types" mode, which tries to find the
-        // PSI element for each declaration with unresolved types and tries to heuristically "resolve" those unresolved types to
-        // generate them into the Java stub. In case of delegated property accessors, it should look for the property declaration,
-        // since the type can only be provided there, and not in the `by ...` expression.
-        return if (element is KtPropertyDelegate) element.parent else element
-    }
-
     private fun storeSerializedIr(serializedIr: ByteArray) {
         val av = visitor.newAnnotation(JvmAnnotationNames.SERIALIZED_IR_DESC, true)
         val partsVisitor = av.visitArray(JvmAnnotationNames.SERIALIZED_IR_BYTES_FIELD_NAME)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/extensions/ClassBuilderExtensionAdapter.kt

@@ -0,0 +1,145 @@
+/*
+ * Copyright 2010-2023 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.backend.jvm.extensions
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.backend.jvm.ir.psiElement
+import org.jetbrains.kotlin.codegen.ClassBuilder
+import org.jetbrains.kotlin.codegen.ClassBuilderFactory
+import org.jetbrains.kotlin.codegen.DelegatingClassBuilder
+import org.jetbrains.kotlin.codegen.DelegatingClassBuilderFactory
+import org.jetbrains.kotlin.codegen.extensions.ClassBuilderInterceptorExtension
+import org.jetbrains.kotlin.diagnostics.DiagnosticSink
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
+import org.jetbrains.org.objectweb.asm.AnnotationVisitor
+import org.jetbrains.org.objectweb.asm.FieldVisitor
+import org.jetbrains.org.objectweb.asm.MethodVisitor
+import org.jetbrains.org.objectweb.asm.RecordComponentVisitor
+
+// Loads an converts deprecated ClassBuilderInterceptorExtension implementations to the new ClassGeneratorExtension EP.
+@Suppress("unused") // Used reflectively in GenerationState.
+internal object ClassBuilderExtensionAdapter {
+    @JvmStatic
+    fun getExtensions(project: Project): List<ClassBuilderInterceptorExtension> =
+        ClassGeneratorExtension.getInstances(project).map(::ExtensionAdapter)
+}
+
+private class ExtensionAdapter(private val extension: ClassGeneratorExtension) : ClassBuilderInterceptorExtension {
+    override fun interceptClassBuilderFactory(
+        interceptedFactory: ClassBuilderFactory,
+        bindingContext: BindingContext,
+        diagnostics: DiagnosticSink,
+    ): ClassBuilderFactory = object : DelegatingClassBuilderFactory(interceptedFactory) {
+        override fun newClassBuilder(origin: JvmDeclarationOrigin): DelegatingClassBuilder {
+            val classBuilder = interceptedFactory.newClassBuilder(origin)
+            val irClass = origin.unwrapOrigin<IrClass>()
+            return DelegatingClassBuilderAdapter(
+                extension.generateClass(
+                    ClassGeneratorAdapter(irClass, classBuilder),
+                    irClass
+                ),
+                classBuilder
+            )
+        }
+    }
+}
+
+private class ClassGeneratorAdapter(val irClass: IrClass?, val builder: ClassBuilder) : ClassGenerator {
+    override fun defineClass(
+        version: Int, access: Int, name: String, signature: String?, superName: String, interfaces: Array<out String>
+    ) {
+        builder.defineClass(irClass?.psiElement, version, access, name, signature, superName, interfaces)
+    }
+
+    override fun newField(
+        declaration: IrField?, access: Int, name: String, desc: String, signature: String?, value: Any?
+    ): FieldVisitor =
+        builder.newField(declaration.wrapToOrigin(), access, name, desc, signature, value)
+
+    override fun newMethod(
+        declaration: IrFunction?, access: Int, name: String, desc: String, signature: String?, exceptions: Array<out String>?
+    ): MethodVisitor =
+        builder.newMethod(declaration.wrapToOrigin(), access, name, desc, signature, exceptions)
+
+    override fun newRecordComponent(name: String, desc: String, signature: String?): RecordComponentVisitor =
+        builder.newRecordComponent(name, desc, signature)
+
+    override fun visitAnnotation(desc: String, visible: Boolean): AnnotationVisitor =
+        builder.newAnnotation(desc, visible)
+
+    override fun visitInnerClass(name: String, outerName: String?, innerName: String?, access: Int) {
+        builder.visitInnerClass(name, outerName, innerName, access)
+    }
+
+    override fun visitEnclosingMethod(owner: String, name: String?, desc: String?) {
+        builder.visitOuterClass(owner, name, desc)
+    }
+
+    override fun visitSource(name: String, debug: String?) {
+        builder.visitSource(name, debug)
+    }
+
+    override fun done(generateSmapCopyToAnnotation: Boolean) {
+        builder.done(generateSmapCopyToAnnotation)
+    }
+}
+
+private class DelegatingClassBuilderAdapter(
+    private val generator: ClassGenerator,
+    private val originalClassBuilder: ClassBuilder,
+) : DelegatingClassBuilder() {
+    override fun getDelegate(): ClassBuilder = originalClassBuilder
+
+    override fun defineClass(
+        origin: PsiElement?, version: Int, access: Int, name: String, signature: String?, superName: String, interfaces: Array<out String>
+    ) {
+        generator.defineClass(version, access, name, signature, superName, interfaces)
+    }
+
+    override fun newField(
+        origin: JvmDeclarationOrigin, access: Int, name: String, desc: String, signature: String?, value: Any?
+    ): FieldVisitor =
+        generator.newField(origin.unwrapOrigin(), access, name, desc, signature, value)
+
+    override fun newMethod(
+        origin: JvmDeclarationOrigin, access: Int, name: String, desc: String, signature: String?, exceptions: Array<out String>?
+    ): MethodVisitor =
+        generator.newMethod(origin.unwrapOrigin(), access, name, desc, signature, exceptions)
+
+    override fun newRecordComponent(name: String, desc: String, signature: String?): RecordComponentVisitor =
+        generator.newRecordComponent(name, desc, signature)
+
+    override fun newAnnotation(desc: String, visible: Boolean): AnnotationVisitor =
+        generator.visitAnnotation(desc, visible)
+
+    override fun visitOuterClass(owner: String, name: String?, desc: String?) {
+        generator.visitEnclosingMethod(owner, name, desc)
+    }
+
+    override fun visitInnerClass(name: String, outerName: String?, innerName: String?, access: Int) {
+        generator.visitInnerClass(name, outerName, innerName, access)
+    }
+
+    override fun visitSource(name: String, debug: String?) {
+        generator.visitSource(name, debug)
+    }
+
+    override fun done(generateSmapCopyToAnnotation: Boolean) {
+        generator.done(generateSmapCopyToAnnotation)
+    }
+}
+
+private inline fun <reified T : IrDeclaration> JvmDeclarationOrigin.unwrapOrigin(): T? =
+    (this as? JvmIrDeclarationOrigin)?.declaration as? T
+
+private fun IrDeclaration?.wrapToOrigin(): JvmDeclarationOrigin =
+    this?.descriptorOrigin ?: JvmDeclarationOrigin.NO_ORIGIN

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/extensions/ClassGeneratorExtension.kt

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2010-2023 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.backend.jvm.extensions
+
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.org.objectweb.asm.*
+
+/**
+ * An extension to the Kotlin/JVM compiler backend which allows to change how IR is generated into the class files.
+ * It's preferable for compiler plugins to use [IrGenerationExtension] to implement IR-based logic. This extension point is more low-level.
+ */
+interface ClassGeneratorExtension {
+    companion object : ProjectExtensionDescriptor<ClassGeneratorExtension>(
+        "org.jetbrains.kotlin.classGeneratorExtension", ClassGeneratorExtension::class.java
+    )
+
+    /**
+     * Override this method to decorate the [generator] that is used in the compiler backend to generate IR to bytecode.
+     * [Interface delegation](https://kotlinlang.org/docs/delegation.html) can be used to avoid implementing each member manually.
+     *
+     * @param generator the generator used to generate the original class
+     * @param declaration the IR representation of the generated class, or `null` if this class has no IR representation
+     *   (for example, if it's an anonymous object copied during inlining bytecode)
+     */
+    fun generateClass(generator: ClassGenerator, declaration: IrClass?): ClassGenerator
+}
+
+/**
+ * Similarly to ASM's [ClassWriter], provides methods that are used to generate parts of the class.
+ * [newField] and [newMethod] accept an IR element, which the compiler plugin can use to implement its custom logic.
+ */
+interface ClassGenerator {
+    fun defineClass(version: Int, access: Int, name: String, signature: String?, superName: String, interfaces: Array<out String>)
+
+    fun newField(
+        declaration: IrField?, access: Int, name: String, desc: String, signature: String?, value: Any?
+    ): FieldVisitor
+
+    fun newMethod(
+        declaration: IrFunction?, access: Int, name: String, desc: String, signature: String?, exceptions: Array<out String>?
+    ): MethodVisitor
+
+    fun newRecordComponent(name: String, desc: String, signature: String?): RecordComponentVisitor
+
+    fun visitAnnotation(desc: String, visible: Boolean): AnnotationVisitor
+
+    fun visitInnerClass(name: String, outerName: String?, innerName: String?, access: Int)
+
+    fun visitEnclosingMethod(owner: String, name: String?, desc: String?)
+
+    fun visitSource(name: String, debug: String?)
+
+    fun done(generateSmapCopyToAnnotation: Boolean)
+}

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/extensions/JvmIrDeclarationOrigin.kt

@@ -6,8 +6,13 @@
 package org.jetbrains.kotlin.backend.jvm.extensions
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.backend.common.psi.PsiSourceManager
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.backend.jvm.ir.isEffectivelyInlineOnly
+import org.jetbrains.kotlin.backend.jvm.ir.psiElement
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.toIrBasedDescriptor
+import org.jetbrains.kotlin.psi.KtPropertyDelegate
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOriginKind
 
@@ -16,3 +21,38 @@ class JvmIrDeclarationOrigin(
     element: PsiElement?,
     val declaration: IrDeclaration?,
 ) : JvmDeclarationOrigin(originKind, element, declaration?.toIrBasedDescriptor(), null)
+
+val IrDeclaration.descriptorOrigin: JvmIrDeclarationOrigin
+    get() {
+        val psiElement = findPsiElementForDeclarationOrigin()
+        return when {
+            origin == IrDeclarationOrigin.FILE_CLASS ->
+                JvmIrDeclarationOrigin(JvmDeclarationOriginKind.PACKAGE_PART, psiElement, this)
+            (this is IrSimpleFunction && isSuspend && isEffectivelyInlineOnly()) ||
+                    origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE ||
+                    origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE ->
+                JvmIrDeclarationOrigin(JvmDeclarationOriginKind.INLINE_VERSION_OF_SUSPEND_FUN, psiElement, this)
+            else -> JvmIrDeclarationOrigin(JvmDeclarationOriginKind.OTHER, psiElement, this)
+        }
+    }
+
+private fun IrDeclaration.findPsiElementForDeclarationOrigin(): PsiElement? {
+    // For synthetic $annotations methods for properties, use the PSI for the property or the constructor parameter.
+    // It's used in KAPT stub generation to sort the properties correctly based on their source position (see KT-44130).
+    if (this is IrFunction && name.asString().endsWith("\$annotations")) {
+        val metadata = metadata as? DescriptorMetadataSource.Property
+        if (metadata != null) {
+            return metadata.descriptor.psiElement
+        }
+    }
+
+    val element = PsiSourceManager.findPsiElement(this)
+
+    // Offsets for accessors and field of delegated property in IR point to the 'by' keyword, so the closest PSI element is the
+    // KtPropertyDelegate (`by ...` expression). However, old JVM backend passed the PSI element of the property instead.
+    // This is important for example in case of KAPT stub generation in the "correct error types" mode, which tries to find the
+    // PSI element for each declaration with unresolved types and tries to heuristically "resolve" those unresolved types to
+    // generate them into the Java stub. In case of delegated property accessors, it should look for the property declaration,
+    // since the type can only be provided there, and not in the `by ...` expression.
+    return if (element is KtPropertyDelegate) element.parent else element
+}