JVM_IR: add $assertionsDisabled when an inlined function uses it

NOTE: jvmCrossinlineLambdaDeclarationSite.kt is muted because the
inliner does not remap references to an anonymous object's parent
class after regenerating it. Unlike the JVM backend, JVM_IR uses the
top level named class' assertion status for all inner classes. (The
test used to pass because the lambda in `inline fun call` read the
`$assertionsDisabled` field of `CrossinlineLambdaContainer`, which
was not reloaded after changing the assertion status of package `test`.)

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

@@ -200,9 +200,7 @@ private val jvmFilePhases =
         propertiesPhase then
         renameFieldsPhase then
         anonymousObjectSuperConstructorPhase then
-        assertionPhase then
         tailrecPhase then
-        returnableBlocksPhase then
 
         jvmInlineClassPhase then
 
@@ -215,6 +213,8 @@ private val jvmFilePhases =
 
         callableReferencePhase then
         singleAbstractMethodPhase then
+        assertionPhase then
+        returnableBlocksPhase then
         localDeclarationsPhase then
 
         addContinuationPhase then

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

@@ -8,7 +8,9 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import org.jetbrains.kotlin.backend.common.descriptors.WrappedClassDescriptor
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
+import org.jetbrains.kotlin.backend.jvm.lower.buildAssertionsDisabledField
 import org.jetbrains.kotlin.backend.jvm.lower.constantValue
+import org.jetbrains.kotlin.backend.jvm.lower.hasAssertionsDisabledField
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.inline.DefaultSourceMapper
@@ -20,7 +22,12 @@ import org.jetbrains.kotlin.codegen.serialization.JvmSerializerExtension
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
+import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrSetFieldImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.util.*
@@ -91,6 +98,9 @@ open class ClassCodegen protected constructor(
         }
     }
 
+    private var classInitializer: IrSimpleFunction? = null
+    private var generatingClInit: Boolean = false
+
     fun generate(): ReifiedTypeParametersUsages {
         if (withinInline) {
             getOrCreateSourceMapper() //initialize default mapping that would be later written in class file
@@ -128,8 +138,16 @@ 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) {
-            generateDeclaration(declaration)
+            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
@@ -148,6 +166,41 @@ open class ClassCodegen protected constructor(
         return reifiedTypeParametersUsages
     }
 
+    private var hasAssertField = irClass.hasAssertionsDisabledField(context)
+
+    fun generateAssertFieldIfNeeded(): IrExpression? {
+        if (hasAssertField)
+            return null
+        hasAssertField = true
+        val topLevelClass = generateSequence(this) { it.parentClassCodegen }.last().irClass
+        val field = irClass.buildAssertionsDisabledField(context, topLevelClass)
+        generateField(field)
+        // Normally, `InitializersLowering` would move the initializer to <clinit>, but
+        // it's obviously too late for that.
+        val init = IrSetFieldImpl(
+            field.startOffset, field.endOffset, field.symbol, null,
+            field.initializer!!.expression, context.irBuiltIns.unitType
+        )
+        if (classInitializer == null) {
+            classInitializer = buildFun {
+                name = Name.special("<clinit>")
+                returnType = context.irBuiltIns.unitType
+            }.apply {
+                parent = irClass
+                body = IrBlockBodyImpl(startOffset, endOffset)
+            }
+            // Do not add it to `irClass.declarations` to avoid a concurrent modification error.
+        } else if (generatingClInit) {
+            // Not only `classInitializer` is non-null, we're in fact generating it right now.
+            // Attempting to do `body.statements.add` will cause a concurrent modification error,
+            // so the currently active ExpressionCodegen needs to be asked to generate this
+            // initializer directly.
+            return init
+        }
+        (classInitializer!!.body as IrBlockBody).statements.add(0, init)
+        return null
+    }
+
     private fun generateKotlinMetadataAnnotation() {
         val localDelegatedProperties = (irClass.attributeOwnerId as? IrClass)?.let(context.localDelegatedProperties::get)
         if (localDelegatedProperties != null && localDelegatedProperties.isNotEmpty()) {

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

@@ -41,7 +41,13 @@ class IrInlineCodegen(
     codegen, state, function, methodOwner, signature, typeParameterMappings, sourceCompiler, reifiedTypeInliner
 ), IrCallGenerator {
     override fun generateAssertFieldIfNeeded(info: RootInliningContext) {
-        // TODO: JVM assertions are not implemented yet in IR backend
+        if (info.generateAssertField && (sourceCompiler as IrSourceCompilerForInline).isPrimaryCopy) {
+            codegen.classCodegen.generateAssertFieldIfNeeded()?.let {
+                // Generating <clinit> right now, so no longer can insert the initializer into it.
+                // Instead, ask ExpressionCodegen to generate the code for it directly.
+                it.accept(codegen, BlockInfo()).discard()
+            }
+        }
     }
 
     override fun putClosureParametersOnStack(next: LambdaInfo, functionReferenceReceiver: StackValue?) {

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

@@ -164,6 +164,9 @@ class IrSourceCompilerForInline(
         return setOf(name)
     }
 
+    internal val isPrimaryCopy: Boolean
+        get() = codegen.classCodegen !is FakeClassCodegen
+
     private class FakeClassCodegen(irFunction: IrFunction, codegen: ClassCodegen) :
         ClassCodegen(irFunction.parent as IrClass, codegen.context) {
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AssertionLowering.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.types.typeWith
+import org.jetbrains.kotlin.ir.util.fields
 import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.getPackageFragment
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
@@ -38,7 +39,10 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 internal val assertionPhase = makeIrFilePhase(
     ::AssertionLowering,
     name = "Assertion",
-    description = "Lower assert calls depending on the assertions mode"
+    description = "Lower assert calls depending on the assertions mode",
+    // Necessary to place the `$assertionsDisabled` field into the reference's class, not the
+    // class that contains it.
+    prerequisite = setOf(callableReferencePhase)
 )
 
 private class AssertionLowering(private val context: JvmBackendContext) :
@@ -125,34 +129,37 @@ private class AssertionLowering(private val context: JvmBackendContext) :
             +irIfThen(irNot(assertCondition), throwError)
         }
 
-    fun getAssertionDisabled(irBuilder: IrBuilderWithScope, data: ClassInfo): IrExpression {
+    private fun getAssertionDisabled(irBuilder: IrBuilderWithScope, data: ClassInfo): IrExpression {
         if (data.assertionsDisabledField == null)
-            data.assertionsDisabledField = buildAssertionsDisabledField(data.irClass, data.topLevelClass)
+            data.assertionsDisabledField = data.irClass.buildAssertionsDisabledField(context, data.topLevelClass)
         return irBuilder.irGetField(null, data.assertionsDisabledField!!)
     }
 
-    private fun buildAssertionsDisabledField(irClass: IrClass, topLevelClass: IrClass) =
-        buildField {
-            name = Name.identifier(ASSERTIONS_DISABLED_FIELD_NAME)
-            origin = JvmLoweredDeclarationOrigin.GENERATED_ASSERTION_ENABLED_FIELD
-            type = context.irBuiltIns.booleanType
-            isFinal = true
-            isStatic = true
-        }.apply {
-            parent = irClass
-            initializer = context.createIrBuilder(irClass.symbol).run {
-                at(this@apply)
-                irExprBody(irNot(irCall(this@AssertionLowering.context.ir.symbols.desiredAssertionStatus).apply {
-                    dispatchReceiver = getJavaClass(topLevelClass)
-                }))
-            }
-        }
-
-    private fun IrBuilderWithScope.getJavaClass(irClass: IrClass) =
-        with(CallableReferenceLowering) {
-            javaClassReference(irClass.typeWith(), this@AssertionLowering.context)
-        }
-
     private val IrFunction.isAssert: Boolean
         get() = name.asString() == "assert" && getPackageFragment()?.fqName == KotlinBuiltIns.BUILT_INS_PACKAGE_FQ_NAME
 }
+
+private fun IrBuilderWithScope.getJavaClass(backendContext: JvmBackendContext, irClass: IrClass) =
+    with(CallableReferenceLowering) {
+        javaClassReference(irClass.typeWith(), backendContext)
+    }
+
+fun IrClass.buildAssertionsDisabledField(backendContext: JvmBackendContext, topLevelClass: IrClass) =
+    buildField {
+        name = Name.identifier(ASSERTIONS_DISABLED_FIELD_NAME)
+        origin = JvmLoweredDeclarationOrigin.GENERATED_ASSERTION_ENABLED_FIELD
+        type = backendContext.irBuiltIns.booleanType
+        isFinal = true
+        isStatic = true
+    }.apply {
+        parent = this@buildAssertionsDisabledField
+        initializer = backendContext.createIrBuilder(this@buildAssertionsDisabledField.symbol).run {
+            at(this@apply)
+            irExprBody(irNot(irCall(backendContext.ir.symbols.desiredAssertionStatus).apply {
+                dispatchReceiver = getJavaClass(backendContext, topLevelClass)
+            }))
+        }
+    }
+
+fun IrClass.hasAssertionsDisabledField(context: JvmBackendContext) =
+    fields.any { it.name.asString() == ASSERTIONS_DISABLED_FIELD_NAME && it.type == context.irBuiltIns.booleanType && it.isStatic }