Generate continuation class for named functions

TODO: Generate label and result for suspend lambdas

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.jvm.descriptors.JvmDeclarationFactory
 import org.jetbrains.kotlin.backend.jvm.descriptors.JvmSharedVariablesManager
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
+import org.jetbrains.kotlin.codegen.coroutines.coroutinesJvmInternalPackageFqName
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.coroutinesPackageFqName
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
@@ -54,6 +55,7 @@ class JvmBackendContext(
     val transformedSuspendFunctionsCache = mutableMapOf<IrSimpleFunction, IrSimpleFunction>()
 
     private val coroutinePackage = state.module.getPackage(state.languageVersionSettings.coroutinesPackageFqName())
+    private val coroutinesJvmInternalPackage = state.module.getPackage(state.languageVersionSettings.coroutinesJvmInternalPackageFqName())
 
     val continuationClass = symbolTable.referenceClass(
         coroutinePackage.memberScope.getContributedClassifier(
@@ -61,6 +63,12 @@ class JvmBackendContext(
         ) as ClassDescriptor
     )
 
+    val continuationImpl = symbolTable.referenceClass(
+        coroutinesJvmInternalPackage.memberScope.getContributedClassifier(
+            Name.identifier("ContinuationImpl"), NoLookupLocation.FROM_BACKEND
+        ) as ClassDescriptor
+    )
+
     internal fun getTopLevelClass(fqName: FqName): IrClassSymbol {
         val descriptor = state.module.getPackage(fqName.parent()).memberScope.getContributedClassifier(
             fqName.shortName(), NoLookupLocation.FROM_BACKEND

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

@@ -0,0 +1,37 @@
+/*
+ * 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.codegen
+
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
+import org.jetbrains.org.objectweb.asm.MethodVisitor
+
+fun wrapWithCoroutineTransformer(
+    irFunction: IrFunction,
+    classCodegen: ClassCodegen,
+    methodVisitor: MethodVisitor,
+    access: Int,
+    signature: JvmMethodGenericSignature
+): MethodVisitor {
+    return methodVisitor
+//    assert(irFunction.isSuspend)
+//    val state = classCodegen.state
+//    val languageVersionSettings = state.languageVersionSettings
+//    assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
+//    val continuationClass = createContinuationClassForNamedFunction()
+//    CoroutineTransformerMethodVisitor(
+//        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, null, null,
+//        obtainClassBuilderForCoroutineState = { TODO() },
+//        element = irFunction.symbol.descriptor.psiElement as KtElement,
+//        diagnostics = state.diagnostics,
+//        languageVersionSettings = languageVersionSettings,
+//        shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
+//        containingClassInternalName = classCodegen.visitor.thisName,
+//        isForNamedFunction = true,
+//        needDispatchReceiver = true,
+//        internalNameForDispatchReceiver = classCodegen.visitor.thisName
+//    )
+}

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

@@ -70,7 +70,16 @@ open class FunctionCodegen(
             generateAnnotationDefaultValueIfNeeded(methodVisitor)
         } else {
             val frameMap = createFrameMapWithReceivers(signature)
-            ExpressionCodegen(irFunction, frameMap, InstructionAdapter(methodVisitor), classCodegen, isInlineLambda).generate()
+            ExpressionCodegen(
+                irFunction,
+                frameMap,
+                InstructionAdapter(
+                    if (irFunction.isSuspend) wrapWithCoroutineTransformer(irFunction, classCodegen, methodVisitor, flags, signature)
+                    else methodVisitor
+                ),
+                classCodegen,
+                isInlineLambda
+            ).generate()
             methodVisitor.visitMaxs(-1, -1)
         }
         methodVisitor.visitEnd()

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

@@ -7,15 +7,27 @@ package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.descriptors.WrappedValueParameterDescriptor
+import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
+import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.codegen.coroutines.COROUTINE_LABEL_FIELD_NAME
+import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
+import org.jetbrains.kotlin.codegen.coroutines.dataFieldName
 import org.jetbrains.kotlin.codegen.coroutines.getOrCreateJvmSuspendFunctionView
+import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.builders.declarations.buildClass
+import org.jetbrains.kotlin.ir.builders.declarations.buildConstructor
+import org.jetbrains.kotlin.ir.builders.declarations.buildField
+import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl
@@ -29,8 +41,10 @@ import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.*
+import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.util.OperatorNameConventions
 import java.util.*
 
 // TODO: Move to common prefix
@@ -40,19 +54,140 @@ internal val addContinuationPhase = makeIrFilePhase(
     "Add continuation parameter to suspend functions and suspend calls"
 )
 
-private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
+private object CONTINUATION_PARAMETER : IrDeclarationOriginImpl("CONTINUATION_PARAMETER")
-    object CONTINUATION_PARAMETER : IrDeclarationOriginImpl("CONTINUATION_PARAMETER")
+private object CONTINUATION_CLASS : IrDeclarationOriginImpl("CONTINUATION_CLASS")
+private object CONTINUATION_CLASS_CONSTRUCTOR : IrDeclarationOriginImpl("CONTINUATION_CLASS_CONSTRUCTOR")
+private object CONTINUATION_CLASS_COMPLETION_PARAMETER : IrDeclarationOriginImpl("CONTINUATION_CLASS_COMPLETION_PARAMETER")
+private object CONTINUATION_CLASS_INVOKE_SUSPEND : IrDeclarationOriginImpl("CONTINUATION_CLASS_INVOKE_SUSPEND")
+private object CONTINUATION_CLASS_RESULT_FIELD : IrDeclarationOriginImpl("CONTINUATION_CLASS_RESULT_FIELD")
+private object CONTINUATION_CLASS_LABEL_FIELD : IrDeclarationOriginImpl("CONTINUATION_CLASS_LABEL_FIELD")
 
+private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
         val suspendLambdas = markSuspendLambdas(irFile)
-        addContinuationParameter(irFile)
+        val suspendFunctions = addContinuationParameter(irFile, suspendLambdas)
         transformSuspendCalls(irFile, suspendLambdas)
+        suspendFunctions.forEach {
+            generateContinuationClass(it)
+        }
+        // TODO: suspend lambdas are covered by CallableReferenceLowering, which is not ideal
     }
 
-    private fun addContinuationParameter(irFile: IrFile) {
+    // TODO: move to separate lowering and make the lowering common
+    private fun generateContinuationClass(irFunction: IrFunction) {
+        val continuationImpl = context.continuationImpl.owner
+        val continuationClass = buildClass {
+            // TODO: Use the same counter as CallableReferenceLowering
+            name = "${irFunction.name}\$1".synthesizedName
+            origin = CONTINUATION_CLASS
+            visibility = JavaVisibilities.PACKAGE_VISIBILITY
+        }.apply {
+            createImplicitParameterDeclarationWithWrappedDescriptor()
+            parent = irFunction.parent
+            superTypes.add(continuationImpl.defaultType)
+
+            (parent as IrDeclarationContainer).declarations.add(this)
+        }
+
+        val resultField = buildField {
+            origin = CONTINUATION_CLASS_RESULT_FIELD
+            name = Name.identifier(context.state.languageVersionSettings.dataFieldName())
+            type = context.irBuiltIns.anyType
+            isFinal = false
+            visibility = JavaVisibilities.PACKAGE_VISIBILITY
+        }.apply {
+            parent = continuationClass
+            continuationClass.declarations.add(this)
+        }
+
+        val labelField = buildField {
+            origin = CONTINUATION_CLASS_LABEL_FIELD
+            name = Name.identifier(COROUTINE_LABEL_FIELD_NAME)
+            type = context.irBuiltIns.intType
+            isFinal = false
+            visibility = JavaVisibilities.PACKAGE_VISIBILITY
+        }.apply {
+            parent = continuationClass
+            continuationClass.declarations.add(this)
+        }
+
+        // TODO: How the fuck do I disable parameter assertions generation????
+        buildConstructor {
+            isPrimary = true
+            origin = CONTINUATION_CLASS_CONSTRUCTOR
+            visibility = Visibilities.PUBLIC
+            returnType = continuationClass.defaultType
+        }.apply {
+            parent = continuationClass
+            continuationClass.declarations.add(this)
+
+            val completionParameterDescriptor = WrappedValueParameterDescriptor()
+            val completionParameterSymbol = IrValueParameterSymbolImpl(completionParameterDescriptor)
+            valueParameters.add(
+                0, IrValueParameterImpl(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                    CONTINUATION_CLASS_COMPLETION_PARAMETER,
+                    completionParameterSymbol,
+                    Name.identifier("completion"),
+                    0,
+                    context.continuationClass.owner.defaultType,
+                    null,
+                    isCrossinline = false, isNoinline = false
+                ).also { completionParameterDescriptor.bind(it) }
+            )
+
+            val superClassConstructor = continuationImpl.constructors.single { it.valueParameters.size == 1 }
+            body = context.createIrBuilder(symbol).irBlockBody {
+                +irDelegatingConstructorCall(superClassConstructor).also {
+                    it.putValueArgument(0, irGet(completionParameterSymbol.owner))
+                }
+            }
+        }
+
+        val thisReceiver = continuationClass.thisReceiver!!
+
+        val invokeSuspendName = Name.identifier(INVOKE_SUSPEND_METHOD_NAME)
+        val invokeSuspend = continuationImpl.functions.single { it.name == invokeSuspendName }.symbol
+        buildFun {
+            origin = CONTINUATION_CLASS_INVOKE_SUSPEND
+            name = invokeSuspendName
+            visibility = Visibilities.PUBLIC
+            isSuspend = false
+            returnType = context.irBuiltIns.anyType
+        }.apply {
+            parent = continuationClass
+            continuationClass.declarations.add(this)
+            overriddenSymbols.add(invokeSuspend)
+            dispatchReceiverParameter = thisReceiver.copyTo(this)
+
+            val result = invokeSuspend.owner.valueParameters[0].copyTo(this)
+            valueParameters.add(result)
+            body = context.createIrBuilder(symbol).irBlockBody {
+                +irSetField(irGet(dispatchReceiverParameter!!), resultField, irGet(result))
+                +irSetField(
+                    irGet(dispatchReceiverParameter!!), labelField,
+                    // TODO: Why the fuck label is boxed
+                    irCallOp(
+                        context.irBuiltIns.intClass.functions.single { it.owner.name == OperatorNameConventions.OR },
+                        context.irBuiltIns.intType,
+                        irGetField(irGet(dispatchReceiverParameter!!), labelField),
+                        irInt(1 shl 31)
+                    )
+                )
+                +irReturn(irCall(irFunction).also { it.putValueArgument(0, irGet(dispatchReceiverParameter!!)) })
+            }
+        }
+    }
+
+    private fun addContinuationParameter(irFile: IrFile, suspendLambdas: List<IrFunction>): Set<IrFunction> {
+        val views = hashSetOf<IrFunction>()
         fun tryAddContinuationParameter(element: IrElement): List<IrDeclaration>? {
             return if (element is IrSimpleFunction && element.isSuspend) {
-                listOf(element.getOrCreateView())
+                val view = element.getOrCreateView()
+                if (element !in suspendLambdas) {
+                    views.add(view)
+                }
+                listOf(view)
             } else null
         }
 
@@ -67,6 +202,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 declaration.transformDeclarationsFlat(::tryAddContinuationParameter)
             }
         })
+        return views
     }
 
     private fun markSuspendLambdas(irElement: IrElement): List<IrFunction> {
@@ -151,7 +287,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     private fun IrSimpleFunction.getOrCreateView(): IrSimpleFunction =
         context.transformedSuspendFunctionsCache.getOrPut(this) {
             IrFunctionImpl(
-                startOffset, endOffset, origin, getOrCreateJvmSuspendFunctionView(descriptor, context.state), returnType
+                startOffset, endOffset, origin, getOrCreateJvmSuspendFunctionView(descriptor, context.state), context.irBuiltIns.anyType
             ).also {
                 it.parent = parent
                 it.copyTypeParametersFrom(this)
@@ -195,4 +331,4 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         arguments = listOf(makeTypeProjection(returnType, Variance.INVARIANT))
     )
 
-}
+}