JVM_IR: refactor ClassCodegen a bit

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
 import org.jetbrains.kotlin.backend.jvm.lower.CollectionStubComputer
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.InlineClassAbi
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.MemoizedInlineClassReplacements
-import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionOriginal
 import org.jetbrains.kotlin.codegen.inline.NameGenerator
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
@@ -40,7 +39,6 @@ import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi2ir.PsiErrorBuilder
 import org.jetbrains.kotlin.psi2ir.PsiSourceManager
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
@@ -97,7 +95,7 @@ class JvmBackendContext(
     //      an inline function in the same module. Thus, if two inline functions happen to have the same name
     //      and call a third inline function that has an anonymous object, the one which is called last
     //      will overwrite the other's regenerated copy. (Or don't recompile the inline function for every call.)
-    internal val regeneratedObjectNameGenerators = mutableMapOf<Pair<IrClass, Name>, NameGenerator>()
+    internal val regeneratedObjectNameGenerators = mutableMapOf<Pair<IrClass, String>, NameGenerator>()
 
     internal val localDelegatedProperties = mutableMapOf<IrClass, List<IrLocalDelegatedPropertySymbol>>()
 

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

@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.codegen.inline.ReifiedTypeParametersUsages
 import org.jetbrains.kotlin.codegen.inline.SourceMapper
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializerExtension
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
@@ -52,43 +51,28 @@ open class ClassCodegen protected constructor(
     internal val irClass: IrClass,
     val context: JvmBackendContext,
     private val parentClassCodegen: ClassCodegen? = null,
-    private val parentFunction: IrFunction? = null,
+    private val parentFunction: IrFunction? = null
-    private val withinInline: Boolean = false
 ) : InnerClassConsumer {
     private val innerClasses = mutableListOf<IrClass>()
+    private val withinInline: Boolean = parentClassCodegen?.withinInline == true || parentFunction?.isInline == true
 
-    val state = context.state
+    private val state get() = context.state
-
+    private val typeMapper get() = context.typeMapper
-    val typeMapper = context.typeMapper
-    val methodSignatureMapper = context.methodSignatureMapper
 
     val type: Type = typeMapper.mapClass(irClass)
 
-    val visitor: ClassBuilder = createClassBuilder()
-
     val reifiedTypeParametersUsages = ReifiedTypeParametersUsages()
 
     private val jvmSignatureClashDetector = JvmSignatureClashDetector(irClass, type, context)
-    private lateinit var classOrigin: JvmDeclarationOrigin
 
-    open fun createClassBuilder(): ClassBuilder {
+    private val classOrigin = run {
         // The descriptor associated with an IrClass is never modified in lowerings, so it
         // doesn't reflect the state of the lowered class. To make the diagnostics work we
         // pass in a wrapped descriptor instead.
         // TODO: Migrate class builders away from descriptors
-        val descriptor = WrappedClassDescriptor()
+        val descriptor = WrappedClassDescriptor().apply { bind(irClass) }
-        descriptor.bind(irClass)
-        classOrigin = getClassOrigin(descriptor)
-        return state.factory.newVisitor(
-            classOrigin,
-            type,
-            irClass.fileParent.loadSourceFilesInfo()
-        )
-    }
-
-    private fun getClassOrigin(descriptor: ClassDescriptor): JvmDeclarationOrigin {
         val psiElement = context.psiSourceManager.findPsiElement(irClass)
-        return when (irClass.origin) {
+        when (irClass.origin) {
             IrDeclarationOrigin.FILE_CLASS ->
                 JvmDeclarationOrigin(JvmDeclarationOriginKind.PACKAGE_PART, psiElement, descriptor)
             else ->
@@ -96,8 +80,21 @@ open class ClassCodegen protected constructor(
         }
     }
 
+    val visitor: ClassBuilder = createClassBuilder()
+
+    open fun createClassBuilder(): ClassBuilder {
+        return state.factory.newVisitor(classOrigin, type, irClass.fileParent.loadSourceFilesInfo())
+    }
+
     private var sourceMapper: DefaultSourceMapper? = null
 
+    fun getOrCreateSourceMapper(): DefaultSourceMapper {
+        if (sourceMapper == null) {
+            sourceMapper = context.getSourceMapper(irClass)
+        }
+        return sourceMapper!!
+    }
+
     private val serializerExtension = JvmSerializerExtension(visitor.serializationBindings, state, typeMapper)
     private val serializer: DescriptorSerializer? =
         when (val metadata = irClass.metadata) {
@@ -108,22 +105,22 @@ open class ClassCodegen protected constructor(
         }
 
     fun getRegeneratedObjectNameGenerator(function: IrFunction): NameGenerator {
-        val name = if (function.name.isSpecial) Name.identifier("special") else function.name
+        val name = if (function.name.isSpecial) "special" else function.name.asString()
         return context.regeneratedObjectNameGenerators.getOrPut(irClass to name) {
             NameGenerator("${type.internalName}\$$name\$\$inlined")
         }
     }
 
-    private var classInitializer: IrSimpleFunction? = null
+    private var hasAssertField = irClass.hasAssertionsDisabledField(context)
-    private var generatingClInit: Boolean = false
+    private var classInitializer = irClass.functions.singleOrNull { it.name.asString() == "<clinit>" }
+    private var generatingClInit = false
 
     fun generate(): ReifiedTypeParametersUsages {
         if (withinInline) {
             getOrCreateSourceMapper() //initialize default mapping that would be later written in class file
         }
-        val superClassInfo = irClass.getSuperClassInfo(typeMapper)
-        val signature = getSignature(irClass, type, superClassInfo, typeMapper)
 
+        val signature = getSignature(irClass, type, irClass.getSuperClassInfo(typeMapper), typeMapper)
         // Ensure that the backend only produces class names that would be valid in the frontend for JVM.
         if (context.state.classBuilderMode.generateBodies && signature.hasInvalidName()) {
             throw IllegalStateException("Generating class with invalid name '${type.className}': ${irClass.dump()}")
@@ -138,21 +135,37 @@ open class ClassCodegen protected constructor(
             signature.superclassName,
             signature.interfaces.toTypedArray()
         )
+
+        for (declaration in irClass.declarations) {
+            when (declaration) {
+                is IrClass, classInitializer -> Unit // see below
+                is IrField -> generateField(declaration)
+                is IrFunction -> generateMethod(declaration)
+                else -> throw AssertionError("unexpected class member $declaration at codegen")
+            }
+        }
+
+        // Delay generation of <clinit> until the end because inline function calls
+        // might need to generate the `$assertionsDisabled` field initializer.
+        classInitializer?.let {
+            generatingClInit = true
+            generateMethod(it)
+        }
+
+        // Generate nested classes at the end, to ensure that when the companion's metadata is serialized
+        // everything moved to the outer class has already been recorded in `globalSerializationBindings`.
+        for (declaration in irClass.declarations) {
+            if (declaration is IrClass) {
+                ClassCodegen(declaration, context, this).generate()
+            }
+        }
+
         object : AnnotationCodegen(this@ClassCodegen, context) {
             override fun visitAnnotation(descr: String?, visible: Boolean): AnnotationVisitor {
                 return visitor.visitor.visitAnnotation(descr, visible)
             }
-        }.genAnnotations(
+        }.genAnnotations(irClass, null, null)
-            irClass,
+        generateKotlinMetadataAnnotation()
-            null,
-            null
-        )
-
-        val nestedClasses = irClass.declarations.mapNotNull { declaration ->
-            if (declaration is IrClass) {
-                ClassCodegen(declaration, context, this, withinInline = withinInline)
-            } else null
-        }
 
         val fileEntry = context.psiSourceManager.getFileEntry(irClass.fileParent)
         if (fileEntry != null) {
@@ -161,33 +174,22 @@ open class ClassCodegen protected constructor(
             visitor.visitSource(shortName, null)
         }
 
-        // Delay generation of <clinit> until the end because inline function calls
-        // might need to generate the `$assertionsDisabled` field initializer.
-        classInitializer = irClass.functions.singleOrNull { it.name.asString() == "<clinit>" }
-        for (declaration in irClass.declarations) {
-            if (declaration != classInitializer)
-                generateDeclaration(declaration)
-        }
-        classInitializer?.let {
-            generatingClInit = true
-            generateMethod(it)
-        }
-
-        // Generate nested classes at the end, to ensure that codegen for companion object will have the necessary JVM signatures in its
-        // trace for properties moved to the outer class
-        for (codegen in nestedClasses) {
-            codegen.generate()
-        }
-
-        generateKotlinMetadataAnnotation()
-
         if (irClass.origin != JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
             done()
         }
         return reifiedTypeParametersUsages
     }
 
-    private var hasAssertField = irClass.hasAssertionsDisabledField(context)
+    fun done() {
+        generateInnerAndOuterClasses()
+
+        sourceMapper?.let {
+            SourceMapper.flushToClassBuilder(it, visitor)
+        }
+
+        visitor.done()
+        jvmSignatureClashDetector.reportErrors(classOrigin)
+    }
 
     fun generateAssertFieldIfNeeded(): IrExpression? {
         if (hasAssertField)
@@ -291,19 +293,6 @@ open class ClassCodegen protected constructor(
         }
     }
 
-    fun done() {
-        writeInnerClasses()
-        writeOuterClassAndEnclosingMethod()
-
-        sourceMapper?.let {
-            SourceMapper.flushToClassBuilder(it, visitor)
-        }
-
-        jvmSignatureClashDetector.reportErrors(classOrigin)
-
-        visitor.done()
-    }
-
     private fun IrFile.loadSourceFilesInfo(): List<File> {
         val entry = fileEntry
         if (entry is MultifileFacadeFileEntry) {
@@ -321,29 +310,8 @@ open class ClassCodegen protected constructor(
             name.splitToSequence('/').any { identifier -> identifier.any { it in JvmSimpleNameBacktickChecker.INVALID_CHARS } }
     }
 
-    private fun generateDeclaration(declaration: IrDeclaration) {
-        when (declaration) {
-            is IrField ->
-                generateField(declaration)
-            is IrFunction -> {
-                generateMethod(declaration)
-            }
-            is IrAnonymousInitializer -> {
-                // skip
-            }
-            is IrClass -> {
-                // Nested classes are generated separately
-            }
-            else -> throw RuntimeException("Unsupported declaration $declaration")
-        }
-    }
-
-    fun generateLocalClass(klass: IrClass, parentFunction: IrFunction): ReifiedTypeParametersUsages {
-        return createLocalClassCodegen(klass, parentFunction).generate()
-    }
-
     fun createLocalClassCodegen(klass: IrClass, parentFunction: IrFunction): ClassCodegen =
-        ClassCodegen(klass, context, this, parentFunction, withinInline = withinInline || parentFunction.isInline)
+        ClassCodegen(klass, context, this, parentFunction)
 
     private fun generateField(field: IrField) {
         if (field.origin == IrDeclarationOrigin.FAKE_OVERRIDE) return
@@ -351,7 +319,7 @@ open class ClassCodegen protected constructor(
         val fieldType = typeMapper.mapType(field)
         val fieldSignature =
             if (field.origin == IrDeclarationOrigin.PROPERTY_DELEGATE) null
-            else methodSignatureMapper.mapFieldSignature(field)
+            else context.methodSignatureMapper.mapFieldSignature(field)
         val fieldName = field.name.asString()
         val fv = visitor.newField(
             field.OtherOrigin, field.flags, fieldName, fieldType.descriptor,
@@ -406,45 +374,17 @@ open class ClassCodegen protected constructor(
         }
     }
 
-    private fun writeInnerClasses() {
+    private fun generateInnerAndOuterClasses() {
         // JVMS7 (4.7.6): a nested class or interface member will have InnerClasses information
         // for each enclosing class and for each immediate member
-        val classForInnerClassRecord = getClassForInnerClassRecord()
+        parentClassCodegen?.innerClasses?.add(irClass)
-        if (classForInnerClassRecord != null) {
-            parentClassCodegen?.innerClasses?.add(classForInnerClassRecord)
-
-            var codegen: ClassCodegen? = this
-            while (codegen != null) {
-                val outerClass = codegen.getClassForInnerClassRecord()
-                if (outerClass != null) {
-                    innerClasses.add(outerClass)
-                }
-                codegen = codegen.parentClassCodegen
-            }
-        }
-
         for (innerClass in innerClasses) {
             writeInnerClass(innerClass, typeMapper, context, visitor)
         }
-    }
+        for (codegen in generateSequence(this) { it.parentClassCodegen }.takeWhile { it.parentClassCodegen != null }) {
-
+            writeInnerClass(codegen.irClass, typeMapper, context, visitor)
-    private fun getClassForInnerClassRecord(): IrClass? {
-        return if (parentClassCodegen != null) irClass else null
-    }
-
-    // It's necessary for proper recovering of classId by plain string JVM descriptor when loading annotations
-    // See FileBasedKotlinClass.convertAnnotationVisitor
-    override fun addInnerClassInfoFromAnnotation(innerClass: IrClass) {
-        var current: IrDeclaration? = innerClass
-        while (current != null && !current.isTopLevelDeclaration) {
-            if (current is IrClass) {
-                innerClasses.add(current)
-            }
-            current = current.parent as? IrDeclaration
         }
-    }
 
-    private fun writeOuterClassAndEnclosingMethod() {
         // JVMS7 (4.7.7): A class must have an EnclosingMethod attribute if and only if
         // it is a local class or an anonymous class.
         //
@@ -452,22 +392,22 @@ open class ClassCodegen protected constructor(
         // the current class. If the current class is immediately enclosed by a method
         // or constructor, the name and type of the function is recorded as well.
         if (parentClassCodegen != null) {
-            val outerClassName = parentClassCodegen.type.internalName
             val enclosingFunction = context.customEnclosingFunction[irClass.attributeOwnerId] ?: parentFunction
-            if (enclosingFunction != null) {
+            if (enclosingFunction != null || irClass.isAnonymousObject) {
-                val method = methodSignatureMapper.mapAsmMethod(enclosingFunction)
+                val method = enclosingFunction?.let(context.methodSignatureMapper::mapAsmMethod)
-                visitor.visitOuterClass(outerClassName, method.name, method.descriptor)
+                visitor.visitOuterClass(parentClassCodegen.type.internalName, method?.name, method?.descriptor)
-            } else if (irClass.isAnonymousObject) {
-                visitor.visitOuterClass(outerClassName, null, null)
             }
         }
     }
 
-    fun getOrCreateSourceMapper(): DefaultSourceMapper {
+    override fun addInnerClassInfoFromAnnotation(innerClass: IrClass) {
-        if (sourceMapper == null) {
+        // It's necessary for proper recovering of classId by plain string JVM descriptor when loading annotations
-            sourceMapper = context.getSourceMapper(irClass)
+        // See FileBasedKotlinClass.convertAnnotationVisitor
+        generateSequence<IrDeclaration>(innerClass) { it.parent as? IrDeclaration }.takeWhile { !it.isTopLevelDeclaration }.forEach {
+            if (it is IrClass) {
+                writeInnerClass(it, typeMapper, context, visitor)
+            }
         }
-        return sourceMapper!!
     }
 }
 

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

@@ -42,7 +42,7 @@ internal fun generateStateMachine(
     obtainContinuationClassBuilder: () -> ClassBuilder,
     element: KtElement
 ): MethodVisitor {
-    val state = classCodegen.state
+    val state = classCodegen.context.state
     val languageVersionSettings = state.languageVersionSettings
     assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
     return CoroutineTransformerMethodVisitor(

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

@@ -119,7 +119,7 @@ class ExpressionCodegen(
     val typeMapper = context.typeMapper
     val methodSignatureMapper = context.methodSignatureMapper
 
-    val state = classCodegen.state
+    val state = context.state
 
     private val fileEntry = classCodegen.context.psiSourceManager.getFileEntry(irFunction.fileParent)
 
@@ -636,9 +636,8 @@ class ExpressionCodegen(
 
     override fun visitClass(declaration: IrClass, data: BlockInfo): PromisedValue {
         if (declaration.origin != JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
-            classCodegen.generateLocalClass(declaration, generateSequence(this) { it.inlinedInto }.last().irFunction).also {
+            closureReifiedMarkers[declaration] =
-                closureReifiedMarkers[declaration] = it
+                classCodegen.createLocalClassCodegen(declaration, generateSequence(this) { it.inlinedInto }.last().irFunction).generate()
-            }
         }
         return unitValue
     }

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

@@ -41,10 +41,9 @@ open class FunctionCodegen(
     private val classCodegen: ClassCodegen,
     private val inlinedInto: ExpressionCodegen? = null
 ) {
-    val context = classCodegen.context
+    private val context = classCodegen.context
-    val state = classCodegen.state
 
-    val continuationClassCodegen = lazy {
+    private val continuationClassCodegen = lazy {
         classCodegen.createLocalClassCodegen(irFunction.continuationClass(), irFunction).also { it.generate() }
     }
 
@@ -56,13 +55,13 @@ open class FunctionCodegen(
         }
 
     private fun doGenerate(smapOverride: DefaultSourceMapper?): JvmMethodGenericSignature {
-        val signature = classCodegen.methodSignatureMapper.mapSignatureWithGeneric(irFunction)
+        val signature = context.methodSignatureMapper.mapSignatureWithGeneric(irFunction)
 
         val flags = calculateMethodFlags(irFunction.isStatic)
         var methodVisitor = createMethod(flags, signature)
 
-        if (state.generateParametersMetadata && flags.and(Opcodes.ACC_SYNTHETIC) == 0) {
+        if (context.state.generateParametersMetadata && flags.and(Opcodes.ACC_SYNTHETIC) == 0) {
-            generateParameterNames(irFunction, methodVisitor, signature, state)
+            generateParameterNames(irFunction, methodVisitor, signature, context.state)
         }
 
         if (irFunction.origin != IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER) {
@@ -88,7 +87,7 @@ open class FunctionCodegen(
             }
         }
 
-        if (!state.classBuilderMode.generateBodies || flags.and(Opcodes.ACC_ABSTRACT) != 0 || irFunction.isExternal) {
+        if (!context.state.classBuilderMode.generateBodies || flags.and(Opcodes.ACC_ABSTRACT) != 0 || irFunction.isExternal) {
             generateAnnotationDefaultValueIfNeeded(methodVisitor)
         } else {
             val frameMap = createFrameMapWithReceivers()
@@ -211,14 +210,14 @@ open class FunctionCodegen(
         )
 
     private fun getThrownExceptions(function: IrFunction): List<String>? {
-        if (state.languageVersionSettings.supportsFeature(LanguageFeature.DoNotGenerateThrowsForDelegatedKotlinMembers) &&
+        if (context.state.languageVersionSettings.supportsFeature(LanguageFeature.DoNotGenerateThrowsForDelegatedKotlinMembers) &&
             function.origin == IrDeclarationOrigin.DELEGATED_MEMBER
         ) return null
 
         // @Throws(vararg exceptionClasses: KClass<out Throwable>)
         val exceptionClasses = function.getAnnotation(FqName("kotlin.jvm.Throws"))?.getValueArgument(0) ?: return null
         return (exceptionClasses as IrVararg).elements.map { exceptionClass ->
-            classCodegen.typeMapper.mapType((exceptionClass as IrClassReference).classType).internalName
+            context.typeMapper.mapType((exceptionClass as IrClassReference).classType).internalName
         }
     }
 
@@ -251,7 +250,7 @@ open class FunctionCodegen(
     }
 
     private fun IrFrameMap.enterDispatchReceiver(parameter: IrValueParameter) {
-        val type = classCodegen.typeMapper.mapTypeAsDeclaration(parameter.type)
+        val type = context.typeMapper.mapTypeAsDeclaration(parameter.type)
         enter(parameter, type)
     }
 
@@ -264,10 +263,10 @@ open class FunctionCodegen(
             frameMap.enterDispatchReceiver(irFunction.dispatchReceiverParameter!!)
         }
         irFunction.extensionReceiverParameter?.let {
-            frameMap.enter(it, classCodegen.typeMapper.mapType(it))
+            frameMap.enter(it, context.typeMapper.mapType(it))
         }
         for (parameter in irFunction.valueParameters) {
-            frameMap.enter(parameter, classCodegen.typeMapper.mapType(parameter.type))
+            frameMap.enter(parameter, context.typeMapper.mapType(parameter.type))
         }
 
         return frameMap