IR: makes vararg friendly to inlining

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/VarargLowering.kt

@@ -1,15 +1,15 @@
 package org.jetbrains.kotlin.backend.konan.lower
 
 import org.jetbrains.kotlin.backend.common.FunctionLoweringPass
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.konan.Context
 import org.jetbrains.kotlin.backend.konan.KonanPlatform
+import org.jetbrains.kotlin.backend.konan.descriptors.synthesizedString
 import org.jetbrains.kotlin.backend.konan.ir.ir2string
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns.FQ_NAMES
 import org.jetbrains.kotlin.builtins.PrimitiveType
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
-import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
@@ -17,13 +17,13 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrVarargImpl
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.TypeUtils
 
 
 class VarargInjectionLowering internal constructor(val context: Context): FunctionLoweringPass {
@@ -32,117 +32,98 @@ class VarargInjectionLowering internal constructor(val context: Context): Functi
             lower(irFunction.descriptor, irFunction.body!!)
     }
 
-    fun noVarargArguments(descriptor: FunctionDescriptor) = descriptor.valueParameters.none { it.varargElementType != null }
-    val irContext = IrGeneratorContext(context.ir.irModule.irBuiltins)
-
     private fun lower(owner:FunctionDescriptor, irBody: IrBody) {
-        irBody.transformChildrenVoid(object: IrElementTransformerVoid(){
+        irBody.transformChildrenVoid(object: IrElementTransformerVoid() {
             override fun visitCall(expression: IrCall): IrExpression {
-                super.visitCall(expression)
-                val functionDescriptor = expression.descriptor as FunctionDescriptor
-                if (noVarargArguments(functionDescriptor))
-                    return expression
-                val varargToBlock = blockPerVararg(expression, owner)
-                val scope = owner.scope()
-                offset(expression, scope){
-                    val originalCall = irCall(
-                            descriptor = expression.descriptor,
-                            typeArguments = expression.descriptor.original.typeParameters.map { it to expression.getTypeArgument(it)!! }.toMap())
-                    originalCall.dispatchReceiver  = expression.dispatchReceiver
-                    originalCall.extensionReceiver = expression.extensionReceiver
-                    functionDescriptor.valueParameters.forEach {
-                        originalCall.putValueArgument(it.index, if (varargToBlock.containsKey(it)) varargToBlock[it] else expression.getValueArgument(it))
+                context.createIrBuilder(owner, expression.startOffset, expression.endOffset).apply {
+                    expression.descriptor.valueParameters.filter{it.varargElementType != null && expression.getValueArgument(it) == null}.forEach {
+                        expression.putValueArgument(it.index,
+                            IrVarargImpl(startOffset       = startOffset,
+                                         endOffset         = endOffset,
+                                         type              = it.type,
+                                         varargElementType = it.varargElementType!!)
+                        )
                     }
-                    return originalCall
+                }
+                super.visitCall(expression)
+                return expression
+            }
+
+            override fun visitVararg(expression: IrVararg): IrExpression {
+                val hasSpreadElement = hasSpreadElement(expression)
+                if (!hasSpreadElement && expression.elements.all { it is IrConst<*> && KotlinBuiltIns.isString(it.type) }) {
+                    log("skipped vararg expression because it's string array literal")
+                    return expression
+                }
+                val irBuilder = context.createIrBuilder(owner, expression.startOffset, expression.endOffset)
+                irBuilder.run {
+                    val type = expression.varargElementType
+                    log("$expression: array type:$type, is array of primitives ${!KotlinBuiltIns.isArray(expression.type)}")
+                    val arrayHandle = arrayType(expression.type)
+                    val arrayConstructor = arrayHandle.arrayDescriptor.constructors.find { it.valueParameters.size == 1 }!!
+                    val block = irBlock(arrayHandle.arrayDescriptor.defaultType)
+                    val arrayConstructorCall = irCall(
+                            descriptor = arrayConstructor,
+                            typeArguments = typeArgument(arrayConstructor, type))
+
+
+                    val vars = expression.elements.map {
+                        val initVar = scope.createTemporaryVariable(
+                                (it as? IrSpreadElement)?.expression ?: it as IrExpression,
+                                "elem".synthesizedString, true)
+                        block.statements.add(initVar)
+                        it to initVar
+                    }.toMap()
+                    arrayConstructorCall.putValueArgument(0, calculateArraySize(arrayHandle, hasSpreadElement, scope, expression, vars))
+                    val arrayTmpVariable = scope.createTemporaryVariable(arrayConstructorCall, "array".synthesizedString, true)
+                    val indexTmpVariable = scope.createTemporaryVariable(kIntZero, "index".synthesizedString, true)
+                    block.statements.add(arrayTmpVariable)
+                    if (hasSpreadElement) {
+                        block.statements.add(indexTmpVariable)
+                    }
+                    expression.elements.forEachIndexed { i, element ->
+                        irBuilder.startOffset = element.startOffset
+                        irBuilder.endOffset   = element.endOffset
+                        irBuilder.apply {
+                            log("element:$i> ${ir2string(element)}")
+                            val dst = vars[element]!!
+                            if (element !is IrSpreadElement) {
+                                val setArrayElementCall = irCall(
+                                        descriptor    = arrayHandle.setMethodDescriptor,
+                                        typeArguments = null
+                                )
+                                setArrayElementCall.dispatchReceiver = irGet(arrayTmpVariable.descriptor)
+                                setArrayElementCall.putValueArgument(0, if (hasSpreadElement) irGet(indexTmpVariable.descriptor) else irConstInt(i))
+                                setArrayElementCall.putValueArgument(1, irGet(dst.descriptor))
+                                block.statements.add(setArrayElementCall)
+                                if (hasSpreadElement) {
+                                    block.statements.add(incrementVariable(indexTmpVariable.descriptor, kIntOne))
+                                }
+                            } else {
+                                val arraySizeVariable = scope.createTemporaryVariable(irArraySize(arrayHandle, irGet(dst.descriptor)), "length".synthesizedString)
+                                block.statements.add(arraySizeVariable)
+                                val copyCall = irCall(arrayHandle.copyRangeToDescriptor, null).apply {
+                                    extensionReceiver = irGet(dst.descriptor)
+                                    putValueArgument(0, irGet(arrayTmpVariable.descriptor))  /* destination */
+                                    putValueArgument(1, kIntZero)                            /* fromIndex */
+                                    putValueArgument(2, irGet(arraySizeVariable.descriptor)) /* toIndex */
+                                    putValueArgument(3, irGet(indexTmpVariable.descriptor))  /* destinationIndex */
+                                }
+                                block.statements.add(copyCall)
+                                block.statements.add(incrementVariable(indexTmpVariable.descriptor,
+                                        irGet(arraySizeVariable.descriptor)))
+                                log("element:$i:spread element> ${ir2string(element.expression)}")
+                            }
+                        }
+                    }
+                    block.statements.add(irGet(arrayTmpVariable.descriptor))
+                    return block
                 }
             }
         })
     }
 
 
-    private fun blockPerVararg(expression: IrCall, owner: FunctionDescriptor): Map<ValueParameterDescriptor, IrBlock> {
-        val varargArgs = mutableMapOf<ValueParameterDescriptor, IrBlock>()
-        val scope = owner.scope()
-        val calleeDescriptor = expression.descriptor
-        calleeDescriptor.valueParameters
-                .filter{ it.varargElementType != null &&  expression.getValueArgument(it) is IrVararg?}
-                .forEach {
-            val type = it.varargElementType!!
-            val parameterExpression = expression.getValueArgument(it) as IrVararg?
-            offset(expression, scope) {
-                val hasSpreadElement = hasSpreadElement(parameterExpression)
-                if (!hasSpreadElement && parameterExpression?.elements?.all { it is IrConst<*> && KotlinBuiltIns.isString(it.type)}?:false) {
-                    log("skipped vararg expression because it's string array literal")
-                    return@forEach
-                }
-                log ("$it: array type:$type, is array of primitives ${KotlinBuiltIns.isArray(it.type)}")
-                val arrayHandle      = arrayType(it.type)
-                val arrayConstructor = arrayHandle.arrayDescriptor.constructors.find { it.valueParameters.size == 1 }!!
-                val block            = irBlock(arrayHandle.arrayDescriptor.defaultType)
-                val arrayConstructorCall = irCall(
-                        descriptor    = arrayConstructor,
-                        typeArguments = typeArgument(arrayConstructor, type))
-
-                if (parameterExpression == null) {
-                    arrayConstructorCall.putValueArgument(0, kIntZero)
-                    block.statements.add(arrayConstructorCall)
-                    varargArgs.put(it, block)
-                    return@forEach
-                }
-
-                val vars = parameterExpression.elements.map {
-                    val initVar = scope.createTemporaryVariable((it as? IrSpreadElement)?.expression ?: it as IrExpression, "__elem\$", true)
-                    block.statements.add(initVar)
-                    it to initVar
-                }.toMap()
-                arrayConstructorCall.putValueArgument(0,  calculateArraySize(arrayHandle, hasSpreadElement, scope, parameterExpression, vars))
-                val arrayTmpVariable = scope.createTemporaryVariable(arrayConstructorCall, "__array\$", true)
-                val indexTmpVariable = scope.createTemporaryVariable(kIntZero, "__index\$", true)
-                block.statements.add(arrayTmpVariable)
-                if (hasSpreadElement) {
-                    block.statements.add(indexTmpVariable)
-                }
-                parameterExpression.elements.forEachIndexed { i, element ->
-                    offset(parameterExpression, scope) {
-                        log("element:$i> ${ir2string(element)}")
-                        val dst = vars[element]!!
-                        if (element !is IrSpreadElement) {
-                            val setArrayElementCall = irCall(
-                                    descriptor    = arrayHandle.setMethodDescriptor,
-                                    typeArguments = null
-                            )
-                            setArrayElementCall.dispatchReceiver = irGet(arrayTmpVariable.descriptor)
-                            setArrayElementCall.putValueArgument(0, if (hasSpreadElement) irGet(indexTmpVariable.descriptor) else irConstInt(i))
-                            setArrayElementCall.putValueArgument(1, irGet(dst.descriptor))
-                            block.statements.add(setArrayElementCall)
-                            if (hasSpreadElement) {
-                                block.statements.add(incrementVariable(indexTmpVariable.descriptor, kIntOne))
-                            }
-                        } else {
-                            val arraySizeVariable = scope.createTemporaryVariable(irArraySize(arrayHandle, irGet(dst.descriptor)), "__length\$")
-                            block.statements.add(arraySizeVariable)
-                            val copyCall = irCall(arrayHandle.copyRangeToDescriptor, null).apply {
-                                extensionReceiver = irGet(dst.descriptor)
-                                putValueArgument(0, irGet(arrayTmpVariable.descriptor))  /* destination */
-                                putValueArgument(1, kIntZero)                            /* fromIndex */
-                                putValueArgument(2, irGet(arraySizeVariable.descriptor)) /* toIndex */
-                                putValueArgument(3, irGet(indexTmpVariable.descriptor))  /* destinationIndex */
-                            }
-                            block.statements.add(copyCall)
-                            block.statements.add(incrementVariable(indexTmpVariable.descriptor,
-                                    irGet(arraySizeVariable.descriptor)))
-                            log("element:$i:spread element> ${ir2string(element.expression)}")
-                        }
-                    }
-                }
-
-                block.statements.add(irGet(arrayTmpVariable.descriptor))
-                varargArgs.put(it, block)
-            }
-        }
-        return varargArgs
-    }
-
     private fun typeArgument(arrayConstructor: ClassConstructorDescriptor, type: KotlinType):Map<TypeParameterDescriptor, KotlinType>? {
         return if (!arrayConstructor.typeParameters.isEmpty())
             mapOf(arrayConstructor.typeParameters.first() to type)
@@ -168,22 +149,22 @@ class VarargInjectionLowering internal constructor(val context: Context): Functi
         else -> kArrayHandler
     }
 
-    private fun Offset.intPlus() = irCall(kIntPlusDescriptor, null)
-    private fun Offset.increment(expression: IrExpression, value: IrExpression): IrExpression {
+    private fun IrBuilderWithScope.intPlus() = irCall(kIntPlusDescriptor, null)
+    private fun IrBuilderWithScope.increment(expression: IrExpression, value: IrExpression): IrExpression {
         return intPlus().apply {
             dispatchReceiver = expression
             putValueArgument(0, value)
         }
     }
 
-    private fun Offset.incrementVariable(descriptor: VariableDescriptor, value: IrExpression): IrExpression {
+    private fun IrBuilderWithScope.incrementVariable(descriptor: VariableDescriptor, value: IrExpression): IrExpression {
         return irSetVar(descriptor, intPlus().apply {
             dispatchReceiver = irGet(descriptor)
             putValueArgument(0, value)
         })
     }
     private fun calculateArraySize(arrayHandle: ArrayHandle, hasSpreadElement: Boolean, scope:Scope, expression: IrVararg, vars: Map<IrVarargElement, IrVariable>): IrExpression? {
-        offset(expression, scope) {
+        context.createIrBuilder(scope.scopeOwner, expression.startOffset, expression.endOffset).run {
             if (!hasSpreadElement)
                 return irConstInt(expression.elements.size)
             val notSpreadElementCount = expression.elements.filter { it !is IrSpreadElement}.size
@@ -193,10 +174,9 @@ class VarargInjectionLowering internal constructor(val context: Context): Functi
                 increment(result, arraySize)
             }
         }
-
     }
 
-    private fun Offset.irArraySize(arrayHandle: ArrayHandle, expression: IrExpression): IrExpression {
+    private fun IrBuilderWithScope.irArraySize(arrayHandle: ArrayHandle, expression: IrExpression): IrExpression {
         val arraySize = irCall((arrayHandle.sizeDescriptor as PropertyDescriptor).getter as FunctionDescriptor, null).apply {
             dispatchReceiver = expression
         }
@@ -248,23 +228,10 @@ class VarargInjectionLowering internal constructor(val context: Context): Functi
     private fun methodDescriptor(descriptor:ClassDescriptor, methodName:String) = DescriptorUtils.getFunctionByName(descriptor.unsubstitutedMemberScope, Name.identifier(methodName))
     private fun arrayClassDescriptor(name:String) = kKotlinPackage.memberScope.getContributedClassifier(Name.identifier(name), NoLookupLocation.FROM_BACKEND) as ClassDescriptor
     private fun descriptor(descriptor: ClassDescriptor, name:String) = DescriptorUtils.getAllDescriptors(descriptor.unsubstitutedMemberScope).find{it.name.asString() == name}
-    inline internal fun <R> offset(ir:IrElement, scope:Scope, block: Offset.() -> R):R = offset(irContext, scope, ir.startOffset, ir.endOffset, block)
-}
-internal class Offset(context:IrGeneratorContext, scope: Scope, startOffset:Int, endOfset:Int) : IrBuilderWithScope(context, scope, startOffset, endOfset) {
-    val kIntZero = irConstInt(0)
-    val kIntOne  = irConstInt(1)
-    companion object {
-        internal inline fun <R> use(context:IrGeneratorContext, scope: Scope, startOffset: Int, endOfset: Int, block: Offset.() -> R):R = Offset(context, scope, startOffset, endOfset).block()
-    }
 }
 
-internal fun CallableDescriptor.scope() = Scope(this)
-
-private fun Offset.irConstInt(value: Int): IrConst<Int> = irConst<Int>(kind = IrConstKind.Int, value = value, type = KonanPlatform.builtIns.intType)
-private fun  <T> Offset.irConst(kind: IrConstKind<T>, value: T, type: KotlinType): IrConst<T> = IrConstImpl<T>(startOffset, endOffset,type, kind,value)
-private fun Offset.irBlock(type: KotlinType): IrBlock = IrBlockImpl(startOffset, endOffset, type)
-private fun Offset.irCall(descriptor: CallableDescriptor, typeArguments: Map<TypeParameterDescriptor, KotlinType>?): IrCall = IrCallImpl(startOffset, endOffset, descriptor, typeArguments)
-inline internal fun <R> offset(context:IrGeneratorContext, scope: Scope, startOffset: Int, endOffset: Int, block: Offset.() -> R):R {
-    @Suppress("NON_PUBLIC_CALL_FROM_PUBLIC_INLINE")
-    return Offset.use(context, scope, startOffset, endOffset, block)
-}
+private fun IrBuilderWithScope.irConstInt(value: Int): IrConst<Int> = IrConstImpl.int(startOffset, endOffset, context.builtIns.intType, value)
+private fun IrBuilderWithScope.irBlock(type: KotlinType): IrBlock = IrBlockImpl(startOffset, endOffset, type)
+private fun IrBuilderWithScope.irCall(descriptor: CallableDescriptor, typeArguments: Map<TypeParameterDescriptor, KotlinType>?): IrCall = IrCallImpl(startOffset, endOffset, descriptor, typeArguments)
+private val IrBuilderWithScope.kIntZero get() = irConstInt(0)
+private val IrBuilderWithScope.kIntOne get() = irConstInt(1)