Support simple function inlining in ir

2017-06-22 14:55:59 +02:00
parent 54dc828c8e
commit 665a697093
7 changed files with 222 additions and 87 deletions
@@ -57,6 +57,8 @@ abstract class LambdaInfo(@JvmField val isCrossInline: Boolean) : LabelOwner {

    abstract fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner)

+    open val hasDispatchReceiver = true
+
    fun addAllParameters(remapper: FieldRemapper): Parameters {
        val builder = ParametersBuilder.initializeBuilderFrom(AsmTypes.OBJECT_TYPE, invokeMethod.descriptor, this)

@@ -175,15 +177,32 @@ class DefaultLambda(

 fun Type.boxReceiverForBoundReference() = AsmUtil.boxType(this)

+abstract class ExpressionLambda(protected val typeMapper: KotlinTypeMapper, isCrossInline: Boolean): LambdaInfo(isCrossInline) {

-abstract class ExpressionLambda(isCrossInline: Boolean): LambdaInfo(isCrossInline)
+    override fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner) {
+        val jvmMethodSignature = typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor)
+        val asmMethod = jvmMethodSignature.asmMethod
+        val methodNode = MethodNode(
+                API, AsmUtil.getMethodAsmFlags(invokeMethodDescriptor, OwnerKind.IMPLEMENTATION, sourceCompiler.state),
+                asmMethod.name, asmMethod.descriptor, null, null
+        )
+
+        node = wrapWithMaxLocalCalc(methodNode).let { adapter ->
+            val smap = sourceCompiler.generateLambdaBody(
+                    adapter, jvmMethodSignature, this
+            )
+            adapter.visitMaxs(-1, -1)
+            SMAPAndMethodNode(methodNode, smap)
+        }
+    }
+}

 class PsiExpressionLambda(
        expression: KtExpression,
-        private val typeMapper: KotlinTypeMapper,
+        typeMapper: KotlinTypeMapper,
        isCrossInline: Boolean,
        override val isBoundCallableReference: Boolean
-) : ExpressionLambda(isCrossInline) {
+) : ExpressionLambda(typeMapper, isCrossInline) {

    override val lambdaClassType: Type

@@ -270,21 +289,4 @@ class PsiExpressionLambda(

    val isPropertyReference: Boolean
        get() = propertyReferenceInfo != null
-
-    override fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner) {
-        val jvmMethodSignature = typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor)
-        val asmMethod = jvmMethodSignature.asmMethod
-        val methodNode = MethodNode(
-                API, AsmUtil.getMethodAsmFlags(invokeMethodDescriptor, OwnerKind.IMPLEMENTATION, sourceCompiler.state),
-                asmMethod.name, asmMethod.descriptor, null, null
-        )
-
-        node = wrapWithMaxLocalCalc(methodNode).let { adapter ->
-            val smap = sourceCompiler.generateLambdaBody(
-                    adapter, jvmMethodSignature, this
-            )
-            adapter.visitMaxs(-1, -1)
-            SMAPAndMethodNode(methodNode, smap)
-        }
-    }
 }
@@ -124,8 +124,10 @@ class ParametersBuilder private constructor() {
                objectType: Type, descriptor: String, inlineLambda: LambdaInfo? = null
        ): ParametersBuilder {
            val builder = newBuilder()
-            //skipped this for inlined lambda cause it will be removed
-            builder.addThis(objectType, inlineLambda != null).lambda = inlineLambda
+            if (inlineLambda?.hasDispatchReceiver != false) {
+                //skipped this for inlined lambda cause it will be removed
+                builder.addThis(objectType, inlineLambda != null).lambda = inlineLambda
+            }

            for (type in Type.getArgumentTypes(descriptor)) {
                builder.addNextParameter(type, false)
@@ -218,59 +218,65 @@ class ExpressionCodegen(
        if (callable is IrIntrinsicFunction) {
            return callable.invoke(mv, this, data)
        } else {
-            val callGenerator = getOrCreateCallGenerator(expression, expression.descriptor)
+            return generateCall(expression, callable, data)
+        }
+    }

-            val receiver = expression.dispatchReceiver
-            receiver?.apply {
-                //gen(receiver, callable.dispatchReceiverType!!, data)
-                callGenerator.genValueAndPut(null, this, callable.dispatchReceiverType!!, -1, this@ExpressionCodegen, data)
-            }
+    fun generateCall(expression: IrMemberAccessExpression, callable: Callable, data: BlockInfo): StackValue {
+        val callGenerator = getOrCreateCallGenerator(expression, expression.descriptor)

-            expression.extensionReceiver?.apply {
-                //gen(this, callable.extensionReceiverType!!, data)
-                callGenerator.genValueAndPut(null, this, callable.extensionReceiverType!!, -1, this@ExpressionCodegen, data)
-            }
+        val receiver = expression.dispatchReceiver
+        receiver?.apply {
+            //gen(receiver, callable.dispatchReceiverType!!, data)
+            callGenerator.genValueAndPut(null, this, callable.dispatchReceiverType!!, -1, this@ExpressionCodegen, data)
+        }

-            val defaultMask = DefaultCallArgs(callable.valueParameterTypes.size)
-            expression.descriptor.valueParameters.forEachIndexed { i, parameterDescriptor ->
-                val arg = expression.getValueArgument(i)
-                val parameterType = callable.valueParameterTypes[i]
-                when {
-                    arg != null -> {
-                        callGenerator.genValueAndPut(parameterDescriptor, arg, parameterType, i, this@ExpressionCodegen, data)
-                    }
-                    parameterDescriptor.hasDefaultValue() -> {
-                        callGenerator.putValueIfNeeded(parameterType, StackValue.createDefaultValue(parameterType), ValueKind.DEFAULT_PARAMETER, i, this@ExpressionCodegen)
-                        defaultMask.mark(i)
-                    }
-                    else -> {
-                        assert(parameterDescriptor.varargElementType != null)
-                        //empty vararg
+        expression.extensionReceiver?.apply {
+            //gen(this, callable.extensionReceiverType!!, data)
+            callGenerator.genValueAndPut(null, this, callable.extensionReceiverType!!, -1, this@ExpressionCodegen, data)
+        }

-                        callGenerator.putValueIfNeeded(
-                                parameterType,
-                                StackValue.operation(parameterType) {
-                                    it.aconst(0)
-                                    it.newarray(correctElementType(parameterType))
-                                },
-                                ValueKind.GENERAL_VARARG, i, this@ExpressionCodegen)
-                    }
+        callGenerator.beforeValueParametersStart()
+        val defaultMask = DefaultCallArgs(callable.valueParameterTypes.size)
+        expression.descriptor.valueParameters.forEachIndexed { i, parameterDescriptor ->
+            val arg = expression.getValueArgument(i)
+            val parameterType = callable.valueParameterTypes[i]
+            when {
+                arg != null -> {
+                    callGenerator.genValueAndPut(parameterDescriptor, arg, parameterType, i, this@ExpressionCodegen, data)
+                }
+                parameterDescriptor.hasDefaultValue() -> {
+                    callGenerator.putValueIfNeeded(parameterType, StackValue.createDefaultValue(parameterType), ValueKind.DEFAULT_PARAMETER, i, this@ExpressionCodegen)
+                    defaultMask.mark(i)
+                }
+                else -> {
+                    assert(parameterDescriptor.varargElementType != null)
+                    //empty vararg
+
+                    callGenerator.putValueIfNeeded(
+                            parameterType,
+                            StackValue.operation(parameterType) {
+                                it.aconst(0)
+                                it.newarray(correctElementType(parameterType))
+                            },
+                            ValueKind.GENERAL_VARARG, i, this@ExpressionCodegen)
                }
            }
-
-            callGenerator.genCall(
-                    callable,
-                    defaultMask.generateOnStackIfNeeded(callGenerator, expression.descriptor is ConstructorDescriptor, this),
-                    this
-            )
-
-            val returnType = expression.descriptor.returnType
-            if (returnType != null && KotlinBuiltIns.isNothing(returnType)) {
-                mv.aconst(null)
-                mv.athrow()
-            }
-            return StackValue.onStack(callable.returnType)
        }
+
+        callGenerator.genCall(
+                callable,
+                defaultMask.generateOnStackIfNeeded(callGenerator, expression.descriptor is ConstructorDescriptor, this),
+                this,
+                expression
+        )
+
+        val returnType = expression.descriptor.returnType
+        if (returnType != null && KotlinBuiltIns.isNothing(returnType)) {
+            mv.aconst(null)
+            mv.athrow()
+        }
+        return StackValue.onStack(callable.returnType)
    }

    override fun visitInstanceInitializerCall(expression: IrInstanceInitializerCall, data: BlockInfo): StackValue {
@@ -957,12 +963,12 @@ class ExpressionCodegen(
            TODO()
        }
        else {
-            IrInlineCodegen(this, state, original, typeParameterMappings!!, IrSourceCompilerForInline(state, element))
+            IrInlineCodegen(this, state, original, typeParameterMappings!!, IrSourceCompilerForInline(state, element, this))
        }
    }

    internal fun getOrCreateCallGenerator(memberAccessExpression: IrMemberAccessExpression, descriptor: CallableDescriptor): IrCallGenerator {
-        val typeArguments = descriptor.typeParameters.keysToMap { memberAccessExpression.getTypeArgumentOrDefault(it) }
+        val typeArguments = descriptor.original.typeParameters.keysToMap { memberAccessExpression.getTypeArgumentOrDefault(it) }

        val mappings = TypeParameterMappings()
        for (entry in typeArguments.entries) {
@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import org.jetbrains.kotlin.backend.jvm.descriptors.JvmDescriptorWithExtraFlags
 import org.jetbrains.kotlin.backend.jvm.lower.InitializersLowering
 import org.jetbrains.kotlin.codegen.*
-import org.jetbrains.kotlin.codegen.AsmUtil.isStaticMethod
 import org.jetbrains.kotlin.codegen.FunctionCodegen
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
@@ -29,6 +28,7 @@ import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
+import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.resolve.source.getPsi
@@ -36,7 +36,7 @@ import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter

-class FunctionCodegen(private val irFunction: IrFunction, private val classCodegen: ClassCodegen) {
+open class FunctionCodegen(private val irFunction: IrFunction, private val classCodegen: ClassCodegen) {

    val state = classCodegen.state

@@ -66,10 +66,7 @@ class FunctionCodegen(private val irFunction: IrFunction, private val classCodeg
            //reset abstract flag
            flags = flags.xor(Opcodes.ACC_ABSTRACT)
        }
-        val methodVisitor = classCodegen.visitor.newMethod(irFunction.OtherOrigin,
-                                                           flags,
-                                                           signature.asmMethod.name, signature.asmMethod.descriptor,
-                                                           signature.genericsSignature, null/*TODO support exception*/)
+        val methodVisitor = createMethod(flags, signature)

        FunctionCodegen.generateMethodAnnotations(descriptor, signature.asmMethod, methodVisitor, classCodegen, state.typeMapper)
        FunctionCodegen.generateParameterAnnotations(descriptor, methodVisitor, signature, classCodegen, state)
@@ -83,6 +80,13 @@ class FunctionCodegen(private val irFunction: IrFunction, private val classCodeg
        ExpressionCodegen(irFunction, frameMap, InstructionAdapter(methodVisitor), classCodegen).generate()
    }

+    open protected fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {
+        return classCodegen.visitor.newMethod(irFunction.OtherOrigin,
+                                              flags,
+                                              signature.asmMethod.name, signature.asmMethod.descriptor,
+                                              signature.genericsSignature, null/*TODO support exception*/)
+    }
+
    private fun generateAnnotationDefaultValueIfNeeded(methodVisitor: MethodVisitor) {
        val directMember = JvmCodegenUtil.getDirectMember(descriptor)
        if (DescriptorUtils.isAnnotationClass(directMember.containingDeclaration)) {
@@ -22,11 +22,12 @@ import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.ValueKind
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.org.objectweb.asm.Type

 interface IrCallGenerator {

-    fun genCall(callableMethod: Callable, callDefault: Boolean, codegen: ExpressionCodegen) {
+    fun genCall(callableMethod: Callable, callDefault: Boolean, codegen: ExpressionCodegen, expression: IrMemberAccessExpression) {
        if (!callDefault) {
            callableMethod.genInvokeInstruction(codegen.mv)
        } else {
@@ -34,6 +35,9 @@ interface IrCallGenerator {
        }
    }

+    fun beforeValueParametersStart() {
+
+    }

    fun genValueAndPut(
            valueParameterDescriptor: ValueParameterDescriptor?,
@@ -16,15 +16,20 @@

 package org.jetbrains.kotlin.backend.jvm.codegen

-import org.jetbrains.kotlin.codegen.*
+import org.jetbrains.kotlin.codegen.Callable
+import org.jetbrains.kotlin.codegen.StackValue
+import org.jetbrains.kotlin.codegen.ValueKind
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
-import org.jetbrains.kotlin.ir.expressions.IrBlock
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.util.getArguments
+import org.jetbrains.kotlin.utils.keysToMap
 import org.jetbrains.org.objectweb.asm.Type
+import org.jetbrains.org.objectweb.asm.commons.Method

 class IrInlineCodegen(
        codegen: ExpressionCodegen,
@@ -35,13 +40,90 @@ class IrInlineCodegen(
 ) : InlineCodegen<ExpressionCodegen>(codegen, state, function, typeParameterMappings, sourceCompiler), IrCallGenerator {

    override fun putClosureParametersOnStack(next: LambdaInfo, functionReferenceReceiver: StackValue?) {
-        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        val lambdaInfo = next as IrExpressionLambda
+
+        val argumentTypes = lambdaInfo.loweredMethod.argumentTypes
+        lambdaInfo.reference.getArguments().forEachIndexed { index, (descriptor, ir) ->
+            putCapturedValueOnStack(ir, argumentTypes[index], index)
+        }
    }

    override fun genValueAndPut(valueParameterDescriptor: ValueParameterDescriptor?, argumentExpression: IrExpression, parameterType: Type, parameterIndex: Int, codegen: ExpressionCodegen, blockInfo: BlockInfo) {
+        //if (isInliningParameter(argumentExpression, valueParameterDescriptor)) {
        if (argumentExpression is IrBlock && argumentExpression.origin == IrStatementOrigin.LAMBDA) {
-            //TODO
+            val irReference: IrFunctionReference = argumentExpression.statements.filterIsInstance<IrFunctionReference>().single()
+            rememberClosure(irReference, parameterType, valueParameterDescriptor!!) as IrExpressionLambda
+        }
+        else {
+            putValueOnStack(argumentExpression, parameterType, valueParameterDescriptor?.index ?: -1)
        }
-        super.genValueAndPut(valueParameterDescriptor, argumentExpression, parameterType, parameterIndex, codegen, blockInfo)
    }
+
+    override fun putValueIfNeeded(parameterType: Type, value: StackValue, kind: ValueKind, parameterIndex: Int, codegen: ExpressionCodegen) {
+        putArgumentOrCapturedToLocalVal(value.type, value, -1, parameterIndex, ValueKind.CAPTURED)
+    }
+
+    fun putCapturedValueOnStack(argumentExpression: IrExpression, valueType: Type, capturedParamindex: Int) {
+        val onStack = codegen.gen(argumentExpression, valueType, BlockInfo.create())
+        putArgumentOrCapturedToLocalVal(onStack.type, onStack, capturedParamindex, capturedParamindex, ValueKind.CAPTURED)
+    }
+
+    fun putValueOnStack(argumentExpression: IrExpression, valueType: Type, paramIndex: Int) {
+        val onStack = codegen.gen(argumentExpression, valueType, BlockInfo.create())
+        putArgumentOrCapturedToLocalVal(onStack.type, onStack, -1, paramIndex, ValueKind.CAPTURED)
+    }
+
+    override fun beforeValueParametersStart() {
+        invocationParamBuilder.markValueParametersStart()
+    }
+
+    override fun genCall(callableMethod: Callable, callDefault: Boolean, codegen: ExpressionCodegen, expression: IrMemberAccessExpression) {
+        val typeArguments = expression.descriptor.typeParameters.keysToMap { expression.getTypeArgumentOrDefault(it) }
+        performInline(typeArguments, callDefault, codegen)
+    }
+
+    private fun rememberClosure(irReference: IrFunctionReference, type: Type, parameter: ValueParameterDescriptor): LambdaInfo {
+        //assert(InlineUtil.isInlinableParameterExpression(ktLambda)) { "Couldn't find inline expression in ${expression.text}" }
+
+        val expression = irReference.symbol.owner as IrFunction
+        return IrExpressionLambda(
+                irReference, expression, typeMapper, parameter.isCrossinline, false/*TODO*/
+        ).also { lambda ->
+            val closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.index)
+            closureInfo.lambda = lambda
+            expressionMap.put(closureInfo.index, lambda)
+        }
+    }
+}
+
+class IrExpressionLambda(
+        val reference: IrFunctionReference,
+        val function: IrFunction,
+        typeMapper: KotlinTypeMapper,
+        isCrossInline: Boolean,
+        override val isBoundCallableReference: Boolean
+) : ExpressionLambda(typeMapper, isCrossInline) {
+
+    override fun isMyLabel(name: String): Boolean {
+        //TODO("not implemented")
+        return false
+    }
+
+    override val lambdaClassType: Type
+        get() = Type.getObjectType("test123")
+
+    override val invokeMethod: Method
+        get() = typeMapper.mapAsmMethod(function.descriptor)
+
+    val loweredMethod: Method
+        get() = typeMapper.mapAsmMethod(function.descriptor)
+
+    override val invokeMethodDescriptor: FunctionDescriptor
+        get() = function.descriptor
+
+    override val capturedVars: List<CapturedParamDesc>
+        get() = emptyList() //cause closure conversion
+
+    override val hasDispatchReceiver: Boolean
+        get() = false
 }
@@ -25,15 +25,20 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
+import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
+import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.commons.Method
 import org.jetbrains.org.objectweb.asm.tree.MethodNode

 class IrSourceCompilerForInline(
        override val state: GenerationState,
-        override val callElement: IrMemberAccessExpression
+        override val callElement: IrMemberAccessExpression,
+        private val codegen: ExpressionCodegen
        ): SourceCompilerForInline {


@@ -58,11 +63,41 @@ class IrSourceCompilerForInline(
        get() = DefaultSourceMapper(SourceInfo("TODO", "TODO", 100))

    override fun generateLambdaBody(adapter: MethodVisitor, jvmMethodSignature: JvmMethodSignature, lambdaInfo: ExpressionLambda): SMAP {
-        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        lambdaInfo as? IrExpressionLambda ?: error("Expecting ir lambda, but $lambdaInfo")
+
+        val functionCodegen = object : FunctionCodegen(lambdaInfo.function, codegen.classCodegen) {
+            override fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {
+                //TODO: to avoid smap assertion
+                adapter.visitLineNumber(1, Label())
+                return adapter
+            }
+        }
+        functionCodegen.generate()
+
+        return SMAP(/*TODO*/listOf(FileMapping("TODO", "TODO").also { it.id = 1; it.addRangeMapping(RangeMapping(1, 1, 1)) }))
    }

    override fun doCreateMethodNodeFromSource(callableDescriptor: FunctionDescriptor, jvmSignature: JvmMethodSignature, callDefault: Boolean, asmMethod: Method): SMAPAndMethodNode {
-        TODO("not implemented")
+        assert(callableDescriptor == callElement.descriptor.original)
+        val owner = (callElement as IrCall).symbol.owner as IrFunction
+        //ExpressionCodegen()
+        var node: MethodNode? = null
+        var maxCalcAdapter: MethodVisitor? = null
+        val functionCodegen = object : FunctionCodegen(owner, codegen.classCodegen) {
+            override fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {
+                node = MethodNode(API,
+                                  flags,
+                                  signature.asmMethod.name, signature.asmMethod.descriptor,
+                                  signature.genericsSignature, null)
+                maxCalcAdapter = wrapWithMaxLocalCalc(node!!)
+                return maxCalcAdapter!!
+            }
+        }
+        functionCodegen.generate()
+        maxCalcAdapter!!.visitMaxs(-1, -1)
+        maxCalcAdapter!!.visitEnd()
+
+        return SMAPAndMethodNode(node!!, SMAP(/*TODO*/listOf(FileMapping.SKIP)))
    }

    override fun generateAndInsertFinallyBlocks(intoNode: MethodNode, insertPoints: List<MethodInliner.PointForExternalFinallyBlocks>, offsetForFinallyLocalVar: Int) {