JVM IR: Add a lowering to desugar varargs to arrays

2019-08-06 11:38:26 +02:00
parent a7b31ba42b
commit 8d9f2f4238
5 changed files with 369 additions and 1 deletions
@@ -120,6 +120,8 @@ private val jvmFilePhases =
        provisionalFunctionExpressionPhase then
        inventNamesForLocalClassesPhase then
        kCallableNamePropertyPhase then
+        annotationPhase then
+        varargPhase then
        arrayConstructorPhase then

        lateinitPhase then
@@ -132,7 +134,6 @@ private val jvmFilePhases =
        propertiesPhase then
        renameFieldsPhase then
        assertionPhase then
-        annotationPhase then
        tailrecPhase then

        jvmInlineClassPhase then
@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 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.classOrNull
 import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.types.isNullableAny
 import org.jetbrains.kotlin.ir.types.typeWith
@@ -344,4 +345,57 @@ class JvmSymbols(
                returnType = javaLangClass.typeWith()
            }
        }.symbol
+
+    val spreadBuilder = createClass(FqName("kotlin.jvm.internal.SpreadBuilder")) { klass ->
+        klass.addConstructor().apply {
+            addValueParameter("size", irBuiltIns.intType)
+        }
+
+        klass.addFunction("addSpread", irBuiltIns.unitType).apply {
+            addValueParameter("container", irBuiltIns.anyNType)
+        }
+
+        klass.addFunction("add", irBuiltIns.unitType).apply {
+            addValueParameter("element", irBuiltIns.anyNType)
+        }
+
+        klass.addFunction("size", irBuiltIns.intType)
+
+        klass.addFunction("toArray", irBuiltIns.arrayClass.typeWith(irBuiltIns.anyNType)).apply {
+            addValueParameter("a", irBuiltIns.arrayClass.typeWith(irBuiltIns.anyNType))
+        }
+    }
+
+    val primitiveSpreadBuilders = irBuiltIns.primitiveIrTypes.associateWith { irType ->
+        val name = irType.classOrNull!!.owner.name
+        createClass(FqName("kotlin.jvm.internal.${name}SpreadBuilder")) { klass ->
+            klass.addConstructor().apply {
+                addValueParameter("size", irBuiltIns.intType)
+            }
+
+            klass.addFunction("addSpread", irBuiltIns.unitType).apply {
+                // This is really a generic method in the superclass (PrimitiveSpreadBuilder).
+                // That is why the argument type is Object rather than the correct
+                // primitive array type.
+                addValueParameter("container", irBuiltIns.anyNType)
+            }
+
+            klass.addFunction("add", irBuiltIns.unitType).apply {
+                addValueParameter("element", irType)
+            }
+
+            klass.addFunction("toArray", irBuiltIns.primitiveArrayForType.getValue(irType).owner.typeWith())
+        }
+    }
+
+    val systemClass = createClass(FqName("java.lang.System")) { klass ->
+        klass.addFunction("arraycopy", irBuiltIns.unitType, isStatic = true).apply {
+            addValueParameter("src", irBuiltIns.anyNType)
+            addValueParameter("srcPos", irBuiltIns.intType)
+            addValueParameter("dest", irBuiltIns.anyNType)
+            addValueParameter("destPos", irBuiltIns.intType)
+            addValueParameter("length", irBuiltIns.intType)
+        }
+    }
+    val systemArraycopy = systemClass.functions.single { it.name.asString() == "arraycopy" }
 }
@@ -0,0 +1,171 @@
+/*
+ * 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.ir
+
+import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.unboxInlineClass
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionAccessExpression
+import org.jetbrains.kotlin.ir.expressions.IrGetValue
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.util.constructors
+import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.getPropertyGetter
+
+inline fun JvmIrBuilder.irArray(arrayType: IrType, block: IrArrayBuilder.() -> Unit): IrExpression =
+    IrArrayBuilder(this, arrayType).apply { block() }.build()
+
+fun JvmIrBuilder.irArrayOf(arrayType: IrType, elements: List<IrExpression> = listOf()): IrExpression =
+    irArray(arrayType) { elements.forEach { +it } }
+
+private class IrArrayElement(val expression: IrExpression, val isSpread: Boolean)
+
+class IrArrayBuilder(val builder: JvmIrBuilder, val arrayType: IrType) {
+    // We build unboxed arrays for inline classes (UIntArray, etc) by first building
+    // an unboxed array of the underlying primitive type, then coercing the result
+    // to the correct type.
+    val unwrappedArrayType = arrayType.unboxInlineClass()
+
+    // Check if the array type is an inline class wrapper (UIntArray, etc.)
+    val isUnboxedInlineClassArray
+        get() = arrayType !== unwrappedArrayType
+
+    // The unwrapped element type
+    val elementType = unwrappedArrayType.getArrayElementType(builder.context.irBuiltIns)
+
+    private val elements: MutableList<IrArrayElement> = mutableListOf()
+
+    private val hasSpread
+        get() = elements.any { it.isSpread }
+
+    operator fun IrExpression.unaryPlus() = add(this)
+    fun add(expression: IrExpression) = elements.add(IrArrayElement(expression, false))
+
+    fun addSpread(expression: IrExpression) = elements.add(IrArrayElement(expression, true))
+
+    fun build(): IrExpression {
+        val array = when {
+            elements.isEmpty() -> newArray(0)
+            !hasSpread -> buildSimpleArray()
+            elements.size == 1 -> copyArray(elements.single().expression)
+            else -> buildComplexArray()
+        }
+        return coerce(array, arrayType)
+    }
+
+    // Construct a new array of the specified size
+    private fun newArray(size: Int) = newArray(builder.irInt(size))
+
+    private fun newArray(size: IrExpression): IrExpression {
+        val arrayConstructor = if (unwrappedArrayType.isBoxedArray)
+            builder.backendContext.ir.symbols.arrayOfNulls
+        else
+            unwrappedArrayType.classOrNull!!.constructors.single { it.owner.valueParameters.size == 1 }
+
+        return builder.irCall(arrayConstructor, unwrappedArrayType).apply {
+            if (typeArgumentsCount != 0)
+                putTypeArgument(0, elementType)
+            putValueArgument(0, size)
+        }
+    }
+
+    // Build an array without spreads
+    private fun buildSimpleArray(): IrExpression =
+        builder.irBlock {
+            val result = irTemporary(newArray(elements.size))
+
+            val set = unwrappedArrayType.classOrNull!!.functions.single {
+                it.owner.name.asString() == "set"
+            }
+
+            for ((index, element) in elements.withIndex()) {
+                +irCall(set).apply {
+                    dispatchReceiver = irGet(result)
+                    putValueArgument(0, irInt(index))
+                    putValueArgument(1, coerce(element.expression, elementType))
+                }
+            }
+
+            +irGet(result)
+        }
+
+    // Copy a single spread expression, unless it refers to a newly constructed array.
+    private fun copyArray(spread: IrExpression): IrExpression {
+        if (spread is IrConstructorCall ||
+            (spread is IrFunctionAccessExpression && spread.symbol == builder.backendContext.ir.symbols.arrayOfNulls))
+            return spread
+
+        return builder.irBlock {
+            val spreadVar = if (spread is IrGetValue) spread.symbol.owner else irTemporary(spread)
+            val size = unwrappedArrayType.classOrNull!!.getPropertyGetter("size")!!
+            fun getSize() = irCall(size).apply { dispatchReceiver = irGet(spreadVar) }
+            val result = irTemporary(newArray(getSize()))
+            +irCall(builder.backendContext.ir.symbols.systemArraycopy).apply {
+                putValueArgument(0, irGet(spreadVar))
+                putValueArgument(1, irInt(0))
+                putValueArgument(2, irGet(result))
+                putValueArgument(3, irInt(0))
+                putValueArgument(4, getSize())
+            }
+            +irGet(result)
+        }
+    }
+
+    // Build an array containing spread expressions.
+    private fun buildComplexArray(): IrExpression {
+        val spreadBuilder = if (unwrappedArrayType.isBoxedArray)
+            builder.backendContext.ir.symbols.spreadBuilder
+        else
+            builder.backendContext.ir.symbols.primitiveSpreadBuilders.getValue(elementType)
+
+        val addElement = spreadBuilder.functions.single { it.name.asString() == "add" }
+        val addSpread = spreadBuilder.functions.single { it.name.asString() == "addSpread" }
+        val toArray = spreadBuilder.functions.single { it.name.asString() == "toArray" }
+
+        return builder.irBlock {
+            val spreadBuilderVar = irTemporary(irCallConstructor(spreadBuilder.constructors.single().symbol, listOf()).apply {
+                putValueArgument(0, irInt(elements.size))
+            })
+
+            for (element in elements) {
+                +irCall(if (element.isSpread) addSpread else addElement).apply {
+                    dispatchReceiver = irGet(spreadBuilderVar)
+                    putValueArgument(0, if (element.isSpread) element.expression else coerce(element.expression, elementType))
+                }
+            }
+
+            val toArrayCall = irCall(toArray).apply {
+                dispatchReceiver = irGet(spreadBuilderVar)
+                if (unwrappedArrayType.isBoxedArray) {
+                    val size = spreadBuilder.functions.single { it.name.asString() == "size" }
+                    putValueArgument(0, irCall(builder.backendContext.ir.symbols.arrayOfNulls, arrayType).apply {
+                        putTypeArgument(0, elementType)
+                        putValueArgument(0, irCall(size).apply {
+                            dispatchReceiver = irGet(spreadBuilderVar)
+                        })
+                    })
+                }
+            }
+
+            if (unwrappedArrayType.isBoxedArray)
+                +builder.irImplicitCast(toArrayCall, unwrappedArrayType)
+            else
+                +toArrayCall
+        }
+    }
+
+    // Coerce expression to irType if we are working with an inline class array type
+    private fun coerce(expression: IrExpression, irType: IrType): IrExpression =
+        if (isUnboxedInlineClassArray)
+            builder.irCall(builder.backendContext.ir.symbols.unsafeCoerceIntrinsicSymbol, irType).apply {
+                putTypeArgument(0, expression.type)
+                putTypeArgument(1, irType)
+                putValueArgument(0, expression)
+            }
+        else expression
+}
@@ -5,11 +5,17 @@

 package org.jetbrains.kotlin.backend.jvm.ir

+import org.jetbrains.kotlin.backend.common.lower.IrLoweringContext
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.codegen.isJvmInterface
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
+import org.jetbrains.kotlin.ir.builders.Scope
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
@@ -91,3 +97,22 @@ fun IrType.getArrayElementType(irBuiltIns: IrBuiltIns): IrType =
        ((this as IrSimpleType).arguments.single() as IrTypeProjection).type
    else
        irBuiltIns.primitiveArrayElementTypes.getValue(this.classOrNull!!)
+
+// An IR builder with a reference to the JvmBackendContext
+class JvmIrBuilder(
+    val backendContext: JvmBackendContext,
+    val symbol: IrSymbol,
+    startOffset: Int = UNDEFINED_OFFSET,
+    endOffset: Int = UNDEFINED_OFFSET
+) : IrBuilderWithScope(
+    IrLoweringContext(backendContext),
+    Scope(symbol),
+    startOffset,
+    endOffset
+)
+
+fun JvmBackendContext.createJvmIrBuilder(
+    symbol: IrSymbol,
+    startOffset: Int = UNDEFINED_OFFSET,
+    endOffset: Int = UNDEFINED_OFFSET
+) = JvmIrBuilder(this, symbol, startOffset, endOffset)
@@ -0,0 +1,117 @@
+/*
+ * 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.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.ir.IrArrayBuilder
+import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
+import org.jetbrains.kotlin.backend.jvm.ir.irArray
+import org.jetbrains.kotlin.backend.jvm.ir.irArrayOf
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.PrimitiveType
+import org.jetbrains.kotlin.builtins.UnsignedType
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrPackageFragment
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.util.*
+
+val varargPhase = makeIrFilePhase(
+    ::VarargLowering,
+    name = "VarargLowering",
+    description = "Replace varargs with array arguments and lower arrayOf calls"
+)
+
+private class VarargLowering(val context: JvmBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
+    override fun lower(irFile: IrFile) = irFile.transformChildrenVoid()
+
+    // Ignore annotations
+    override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
+        val constructor = expression.symbol.owner
+        if (constructor.constructedClass.isAnnotationClass)
+            return expression
+        return super.visitConstructorCall(expression)
+    }
+
+    override fun visitFunctionAccess(expression: IrFunctionAccessExpression): IrExpression {
+        expression.transformChildrenVoid()
+        val function = expression.symbol
+
+        // Replace empty varargs with empty arrays
+        for (i in 0 until expression.valueArgumentsCount) {
+            if (expression.getValueArgument(i) != null)
+                continue
+
+            val parameter = function.owner.valueParameters[i]
+            if (parameter.varargElementType != null && !parameter.hasDefaultValue()) {
+                // Compute the correct type for the array argument.
+                val arrayType = parameter.type.substitute(expression.typeSubstitutionMap)
+                expression.putValueArgument(i, createBuilder().irArrayOf(arrayType))
+            }
+        }
+
+        // Lower `arrayOf` calls. When `isArrayOf` returns true we know that the function has exactly one
+        // vararg parameter. Meanwhile, the code above ensures that the corresponding argument is not null.
+        return if (function.isArrayOf) expression.getValueArgument(0)!! else expression
+    }
+
+    override fun visitVararg(expression: IrVararg): IrExpression =
+        createBuilder(expression.startOffset, expression.endOffset).irArray(expression.type) { addVararg(expression) }
+
+    private fun IrArrayBuilder.addVararg(expression: IrVararg) {
+        loop@ for (element in expression.elements) {
+            when (element) {
+                is IrExpression -> +element.transform(this@VarargLowering, null)
+                is IrSpreadElement -> {
+                    val spread = element.expression
+                    if (spread is IrFunctionAccessExpression && spread.symbol.isArrayOf) {
+                        // Skip empty arrays and don't copy immediately created arrays
+                        val argument = spread.getValueArgument(0) ?: continue@loop
+                        if (argument is IrVararg) {
+                            addVararg(argument)
+                            continue@loop
+                        }
+                    }
+                    addSpread(spread.transform(this@VarargLowering, null))
+                }
+                else -> error("Unexpected IrVarargElement subclass: $element")
+            }
+        }
+    }
+
+    private fun createBuilder(startOffset: Int = UNDEFINED_OFFSET, endOffset: Int = UNDEFINED_OFFSET) =
+        context.createJvmIrBuilder(currentScope!!.scope.scopeOwnerSymbol, startOffset, endOffset)
+
+    private val IrFunctionSymbol.isArrayOf
+        get() = this == context.ir.symbols.arrayOf || owner.isPrimitiveArrayOf
+
+    companion object {
+        private val PRIMITIVE_ARRAY_OF_NAMES: Set<String> =
+            (PrimitiveType.values().map { type -> type.name } + UnsignedType.values().map { type -> type.typeName.asString() })
+                .map { name -> name.toLowerCase() + "ArrayOf" }.toSet()
+
+        private val IrFunction.isPrimitiveArrayOf: Boolean
+            get() {
+                val parent = when (val directParent = parent) {
+                    is IrClass -> directParent.getPackageFragment() ?: return false
+                    is IrPackageFragment -> directParent
+                    else -> return false
+                }
+                return parent.fqName == KotlinBuiltIns.BUILT_INS_PACKAGE_FQ_NAME &&
+                        name.asString() in PRIMITIVE_ARRAY_OF_NAMES &&
+                        extensionReceiverParameter == null &&
+                        dispatchReceiverParameter == null &&
+                        valueParameters.size == 1 &&
+                        valueParameters[0].isVararg
+            }
+    }
+}