JVM IR: Handle -Xassertion modes

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

@@ -41,6 +41,7 @@ interface JvmLoweredDeclarationOrigin : IrDeclarationOrigin {
     object GENERATED_SAM_IMPLEMENTATION : IrDeclarationOriginImpl("GENERATED_SAM_IMPLEMENTATION", isSynthetic = true)
     object ENUM_MAPPINGS_FOR_WHEN : IrDeclarationOriginImpl("ENUM_MAPPINGS_FOR_WHEN", isSynthetic = true)
     object SYNTHETIC_INLINE_CLASS_MEMBER : IrDeclarationOriginImpl("SYNTHETIC_INLINE_CLASS_MEMBER", isSynthetic = true)
+    object GENERATED_ASSERTION_ENABLED_FIELD : IrDeclarationOriginImpl("GENERATED_ASSERTION_ENABLED_FIELD", isSynthetic = true)
 }
 
 interface JvmLoweredStatementOrigin : IrStatementOrigin {

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

@@ -127,6 +127,7 @@ private val jvmFilePhases =
         propertiesToFieldsPhase then
         propertiesPhase then
         renameFieldsPhase then
+        assertionPhase then
         annotationPhase then
         tailrecPhase then
 

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

@@ -23,8 +23,10 @@ import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrExternalPackageFragmentSymbolImpl
 import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.types.isNullableAny
 import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
+import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -114,6 +116,15 @@ class JvmSymbols(
     val javaLangClass: IrClassSymbol =
         if (firMode) createClass(FqName("java.lang.Class")) {}.symbol else context.getTopLevelClass(FqName("java.lang.Class"))
 
+    val javaLangAssertionError: IrClassSymbol =
+        context.getTopLevelClass(FqName("java.lang.AssertionError"))
+
+    val assertionErrorConstructor by lazy {
+        context.ir.symbols.javaLangAssertionError.constructors.single {
+            it.owner.valueParameters.size == 1 && it.owner.valueParameters[0].type.isNullableAny()
+        }
+    }
+
     val lambdaClass: IrClassSymbol = createClass(FqName("kotlin.jvm.internal.Lambda")) { klass ->
         klass.addConstructor().apply {
             addValueParameter("arity", irBuiltIns.intType)
@@ -270,14 +281,21 @@ class JvmSymbols(
         }
     }.symbol
 
+    private fun IrClassSymbol.functionByName(name: String): IrSimpleFunctionSymbol =
+        functions.single { it.owner.name.asString() == name }
+
     val getOrCreateKotlinPackage: IrSimpleFunctionSymbol =
-        reflection.functions.single { it.owner.name.asString() == "getOrCreateKotlinPackage" }
+        reflection.functionByName("getOrCreateKotlinPackage")
 
     val getOrCreateKotlinClass: IrSimpleFunctionSymbol =
-        reflection.functions.single { it.owner.name.asString() == "getOrCreateKotlinClass" }
+        reflection.functionByName("getOrCreateKotlinClass")
 
     val getOrCreateKotlinClasses: IrSimpleFunctionSymbol =
-        reflection.functions.single { it.owner.name.asString() == "getOrCreateKotlinClasses" }
+        reflection.functionByName("getOrCreateKotlinClasses")
+
+    val desiredAssertionStatus: IrSimpleFunctionSymbol by lazy {
+        javaLangClass.functionByName("desiredAssertionStatus")
+    }
 
     private val unsafeCoerceIntrinsic =
         buildFun {

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

@@ -0,0 +1,158 @@
+/*
+ * Copyright 2010-2019 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.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.ir.asSimpleLambda
+import org.jetbrains.kotlin.backend.common.ir.inline
+import org.jetbrains.kotlin.backend.common.lower.*
+import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.codegen.ASSERTIONS_DISABLED_FIELD_NAME
+import org.jetbrains.kotlin.config.JVMAssertionsMode
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.builders.declarations.buildField
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.expressions.IrCall
+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.functions
+import org.jetbrains.kotlin.ir.util.getPackageFragment
+import org.jetbrains.kotlin.ir.util.patchDeclarationParents
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+internal val assertionPhase = makeIrFilePhase(
+    ::AssertionLowering,
+    name = "Assertion",
+    description = "Lower assert calls depending on the assertions mode"
+)
+
+private class AssertionLowering(private val context: JvmBackendContext) :
+    FileLoweringPass,
+    IrElementTransformer<AssertionLowering.ClassInfo?>
+{
+    // Keeps track of the $assertionsDisabled field, which we generate lazily for classes containing
+    // assertions when compiled with -Xassertions=jvm.
+    private class ClassInfo(val irClass: IrClass, val topLevelClass: IrClass, var assertionsDisabledField: IrField? = null)
+
+    override fun lower(irFile: IrFile) {
+        // In legacy mode we treat assertions as inline function calls
+        if (context.state.assertionsMode != JVMAssertionsMode.LEGACY)
+            irFile.transformChildren(this, null)
+    }
+
+    override fun visitClass(declaration: IrClass, data: ClassInfo?): IrStatement {
+        val info = ClassInfo(declaration, data?.topLevelClass ?: declaration)
+
+        super.visitClass(declaration, info)
+
+        // Note that it's necessary to add this member at the beginning of the class, before all user-visible
+        // initializers, which may contain assertions. At the same time, assertions are supposed to be enabled
+        // for code which runs before class initialization. This is the reason why this field records whether
+        // assertions are disabled rather than enabled. During initialization, $assertionsDisabled is going
+        // to be false, meaning that assertions are checked.
+        info.assertionsDisabledField?.let {
+            declaration.declarations.add(0, it)
+
+            // Some parents of local declarations are not updated during ad-hoc inlining
+            // TODO: Remove when generic inliner is used
+            declaration.patchDeclarationParents(declaration.parent)
+        }
+
+        return declaration
+    }
+
+    override fun visitCall(expression: IrCall, data: ClassInfo?): IrElement {
+        val function = expression.symbol.owner
+        if (!function.isAssert)
+            return super.visitCall(expression, data)
+
+        val mode = context.state.assertionsMode
+        if (mode == JVMAssertionsMode.ALWAYS_DISABLE)
+            return IrCompositeImpl(expression.startOffset, expression.endOffset, context.irBuiltIns.unitType)
+
+        context.createIrBuilder(expression.symbol).run {
+            at(expression)
+            val assertCondition = expression.getValueArgument(0)!!
+            val lambdaArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
+
+            return if (mode == JVMAssertionsMode.ALWAYS_ENABLE) {
+                checkAssertion(assertCondition, lambdaArgument)
+            } else {
+                require(mode == JVMAssertionsMode.JVM && data != null)
+                irIfThen(
+                    irNot(getAssertionDisabled(this, data)),
+                    checkAssertion(assertCondition, lambdaArgument)
+                )
+            }
+        }
+    }
+
+    private fun IrBuilderWithScope.checkAssertion(assertCondition: IrExpression, lambdaArgument: IrExpression?) =
+        irBlock {
+            val lambda = lambdaArgument?.asSimpleLambda()
+            val invokeVar = if (lambda == null && lambdaArgument != null) irTemporary(lambdaArgument) else null
+
+            val constructor = this@AssertionLowering.context.ir.symbols.assertionErrorConstructor
+            val throwError = irThrow(irCall(constructor).apply {
+                putValueArgument(
+                    0,
+                    when {
+                        lambda != null -> lambda.inline()
+                        lambdaArgument != null -> {
+                            val invoke =
+                                lambdaArgument.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
+                            irCallOp(invoke.symbol, invoke.returnType, irGet(invokeVar!!))
+                        }
+                        else -> irString("Assertion failed")
+                    }
+                )
+            })
+            +irIfThen(irNot(assertCondition), throwError)
+        }
+
+    fun getAssertionDisabled(irBuilder: IrBuilderWithScope, data: ClassInfo): IrExpression {
+        if (data.assertionsDisabledField == null)
+            data.assertionsDisabledField = buildAssertionsDisabledField(data.irClass, 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
+}

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

@@ -451,29 +451,33 @@ internal class CallableReferenceLowering(private val context: JvmBackendContext)
     }
 
     companion object {
-        internal fun IrBuilderWithScope.calculateOwner(irContainer: IrDeclarationParent, context: JvmBackendContext): IrExpression {
+        private fun IrBuilderWithScope.kClassReference(classType: IrType) =
-            val symbols = context.ir.symbols
+            IrClassReferenceImpl(
-
+                startOffset,
-            val isContainerPackage =
+                endOffset,
-                (irContainer as? IrClass)?.origin == IrDeclarationOrigin.FILE_CLASS || irContainer is IrPackageFragment
+                context.irBuiltIns.kClassClass.typeWith(),
-
+                context.irBuiltIns.kClassClass,
-            val clazzRef = IrClassReferenceImpl(
+                classType
-                UNDEFINED_OFFSET,
-                UNDEFINED_OFFSET,
-                symbols.javaLangClass.typeWith(),
-                symbols.javaLangClass,
-                // For built-in members (i.e. top level `toString`) we don't know any meaningful container, so we're generating Any.
-                // The non-IR backend generates equally meaningless "kotlin/KotlinPackage" in this case (see KT-17151).
-                (irContainer as? IrClass)?.defaultType ?: context.irBuiltIns.anyNType
             )
 
-            if (!isContainerPackage) return clazzRef
+        private fun IrBuilderWithScope.kClassToJavaClass(kClassReference: IrExpression, context: JvmBackendContext) =
-
+            irGet(context.ir.symbols.javaLangClass.typeWith(), null, context.ir.symbols.kClassJava.owner.getter!!.symbol).apply {
-            val jClass = irGet(symbols.javaLangClass.typeWith(), null, symbols.kClassJava.owner.getter!!.symbol).apply {
+                extensionReceiver = kClassReference
-                extensionReceiver = clazzRef
             }
-            return irCall(symbols.getOrCreateKotlinPackage).apply {
+
-                putValueArgument(0, jClass)
+        internal fun IrBuilderWithScope.javaClassReference(classType: IrType, context: JvmBackendContext) =
+            kClassToJavaClass(kClassReference(classType), context)
+
+        internal fun IrBuilderWithScope.calculateOwner(irContainer: IrDeclarationParent, context: JvmBackendContext): IrExpression {
+            // For built-in members (i.e. top level `toString`) we don't know any meaningful container, so we're generating Any.
+            // The non-IR backend generates equally meaningless "kotlin/KotlinPackage" in this case (see KT-17151).
+            val kClass = kClassReference((irContainer as? IrClass)?.defaultType ?: context.irBuiltIns.anyNType)
+
+            if ((irContainer as? IrClass)?.origin != IrDeclarationOrigin.FILE_CLASS && irContainer !is IrPackageFragment)
+                return kClass
+
+            return irCall(context.ir.symbols.getOrCreateKotlinPackage).apply {
+                putValueArgument(0, kClassToJavaClass(kClass, context))
                 // Note that this name is not used in reflection. There should be the name of the referenced declaration's
                 // module instead, but there's no nice API to obtain that name here yet
                 // TODO: write the referenced declaration's module name and use it in reflection