JVM_IR: refactor FunctionCodegen a bit

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
 import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
 import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor
+import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
@@ -39,17 +40,20 @@ internal fun generateStateMachine(
     methodVisitor: MethodVisitor,
     access: Int,
     signature: JvmMethodGenericSignature,
-    obtainContinuationClassBuilder: () -> ClassBuilder,
+    obtainContinuationClassBuilder: () -> ClassBuilder
-    element: KtElement
 ): MethodVisitor {
     val state = classCodegen.context.state
     val languageVersionSettings = state.languageVersionSettings
     assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
+    val element = if (irFunction.isSuspend)
+        irFunction.symbol.descriptor.psiElement ?: classCodegen.irClass.descriptor.psiElement
+    else
+        classCodegen.context.suspendLambdaToOriginalFunctionMap[classCodegen.irClass.attributeOwnerId]!!.symbol.descriptor.psiElement
     return CoroutineTransformerMethodVisitor(
         methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, signature.genericsSignature, null,
         obtainClassBuilderForCoroutineState = obtainContinuationClassBuilder,
-        reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element, state, it) },
+        reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element as KtElement, state, it) },
-        lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0,
+        lineNumber = element?.let { CodegenUtil.getLineNumberForElement(it, false) } ?: 0,
         sourceFile = classCodegen.irClass.file.name,
         languageVersionSettings = languageVersionSettings,
         shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
@@ -71,6 +75,26 @@ internal fun IrFunction.anyOfOverriddenFunctionsReturnsNonUnit(): Boolean {
     } == true
 }
 
+internal fun IrFunction.suspendForInlineToOriginal(): IrSimpleFunction? {
+    if (origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
+        origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+    ) return null
+    return parentAsClass.declarations.find {
+        it is IrSimpleFunction && it.attributeOwnerId == (this as IrSimpleFunction).attributeOwnerId &&
+                it.name.asString() + FOR_INLINE_SUFFIX == name.asString()
+    } as IrSimpleFunction?
+}
+
+internal fun IrFunction.alwaysNeedsContinuation(): Boolean =
+    this is IrSimpleFunction && hasContinuation() && parentAsClass.declarations.any {
+        it is IrSimpleFunction && it.attributeOwnerId == attributeOwnerId &&
+                it.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+    }
+
+internal fun IrFunction.continuationClass(): IrClass? =
+    (body as? IrBlockBody)?.statements?.find { it is IrClass && it.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS }
+            as IrClass?
+
 internal fun IrFunction.continuationParameter(): IrValueParameter? = when {
     isInvokeSuspendOfLambda() || isInvokeSuspendForInlineOfLambda() -> dispatchReceiverParameter
     else -> valueParameters.singleOrNull { it.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS }

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

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.Companion.putNeedClassReificationMarker
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.OperationKind.AS
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.OperationKind.SAFE_AS
-import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.codegen.intrinsics.TypeIntrinsics
 import org.jetbrains.kotlin.codegen.pseudoInsns.fakeAlwaysFalseIfeq
 import org.jetbrains.kotlin.codegen.pseudoInsns.fixStackAndJump
@@ -52,7 +51,6 @@ import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeSystemCommonBackendContext
 import org.jetbrains.kotlin.types.model.TypeParameterMarker
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.keysToMap
@@ -222,10 +220,8 @@ class ExpressionCodegen(
 
     private fun generateFakeContinuationConstructorIfNeeded() {
         if (irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE) return
-        val continuationClass = classCodegen.irClass.functions.find {
+        val continuationClass = irFunction.suspendForInlineToOriginal()?.continuationClass()
-            it.attributeOwnerId == (irFunction as? IrSimpleFunction)?.attributeOwnerId &&
+            ?: error("could not find continuation for ${irFunction.render()}")
-                    it.name.asString() == irFunction.name.asString().removeSuffix(FOR_INLINE_SUFFIX)
-        }?.body?.statements?.firstIsInstance<IrClass>() ?: error("could not find continuation for ${irFunction.render()}")
         val continuationType = typeMapper.mapClass(continuationClass)
         val continuationIndex = frameMap.getIndex(irFunction.continuationParameter()!!.symbol)
         with(mv) {

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

@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.jvm.annotations.JVM_SYNTHETIC_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.annotations.STRICTFP_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.annotations.SYNCHRONIZED_ANNOTATION_FQ_NAME
@@ -43,10 +42,6 @@ open class FunctionCodegen(
 ) {
     private val context = classCodegen.context
 
-    private val continuationClassCodegen = lazy {
-        classCodegen.createLocalClassCodegen(irFunction.continuationClass(), irFunction).also { it.generate() }
-    }
-
     fun generate(smapOverride: DefaultSourceMapper? = null): JvmMethodGenericSignature =
         try {
             doGenerate(smapOverride)
@@ -56,8 +51,7 @@ open class FunctionCodegen(
 
     private fun doGenerate(smapOverride: DefaultSourceMapper?): JvmMethodGenericSignature {
         val signature = context.methodSignatureMapper.mapSignatureWithGeneric(irFunction)
-
+        val flags = irFunction.calculateMethodFlags()
-        val flags = calculateMethodFlags(irFunction.isStatic)
         var methodVisitor = createMethod(flags, signature)
 
         if (context.state.generateParametersMetadata && flags.and(Opcodes.ACC_SYNTHETIC) == 0) {
@@ -75,11 +69,7 @@ open class FunctionCodegen(
                         TypeReference.newTypeReference(TypeReference.METHOD_RETURN).value, path, descr, visible
                     )
                 }
-            }.genAnnotations(
+            }.genAnnotations(irFunction, signature.asmMethod.returnType, irFunction.returnType)
-                irFunction,
-                signature.asmMethod.returnType,
-                irFunction.returnType
-            )
             // Not generating parameter annotations for default stubs fixes KT-7892, though
             // this certainly looks like a workaround for a javac bug.
             if (irFunction !is IrConstructor || !irFunction.parentAsClass.shouldNotGenerateConstructorParameterAnnotations()) {
@@ -87,29 +77,18 @@ open class FunctionCodegen(
             }
         }
 
+        val continuationClassCodegen = lazy {
+            classCodegen.createLocalClassCodegen(irFunction.continuationClass()!!, irFunction).also { it.generate() }
+        }
         if (!context.state.classBuilderMode.generateBodies || flags.and(Opcodes.ACC_ABSTRACT) != 0 || irFunction.isExternal) {
             generateAnnotationDefaultValueIfNeeded(methodVisitor)
         } else {
-            val frameMap = createFrameMapWithReceivers()
+            val frameMap = irFunction.createFrameMapWithReceivers()
             if (irFunction.hasContinuation() || irFunction.isInvokeSuspendOfLambda()) {
-                if (irFunction is IrSimpleFunction && irFunction.parentAsClass.declarations.any {
+                methodVisitor = generateStateMachine(irFunction, classCodegen, methodVisitor, flags, signature) {
-                        it is IrSimpleFunction && it.attributeOwnerId == irFunction.attributeOwnerId &&
+                    // This has to be done lazily to avoid generating the class if tail call optimization makes it redundant.
-                                it.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+                    if (irFunction.isSuspend) continuationClassCodegen.value.visitor else classCodegen.visitor
-                    }
-                ) {
-                    // Force generation of fake continuation for inliner.
-                    continuationClassCodegen.value
                 }
-                // This has to be done lazily to avoid generating the class if tail call optimization makes it redundant.
-                val getContinuation = {
-                    if (irFunction.isSuspend)
-                        continuationClassCodegen.value.visitor
-                    else
-                        classCodegen.visitor
-                }
-                methodVisitor = generateStateMachine(
-                    irFunction, classCodegen, methodVisitor, flags, signature, getContinuation, psiElement()
-                )
             }
             context.state.globalInlineContext.enterDeclaration(irFunction.suspendFunctionOriginal().descriptor)
             try {
@@ -121,10 +100,9 @@ open class FunctionCodegen(
             methodVisitor.visitMaxs(-1, -1)
         }
         methodVisitor.visitEnd()
-        if (continuationClassCodegen.isInitialized()) {
+        if (continuationClassCodegen.isInitialized() || irFunction.alwaysNeedsContinuation()) {
             continuationClassCodegen.value.done()
         }
-
         return signature
     }
 
@@ -136,16 +114,6 @@ open class FunctionCodegen(
     private fun IrClass.shouldNotGenerateConstructorParameterAnnotations() =
         isAnonymousObject || origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS || origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
 
-    private fun psiElement(): KtElement =
-        (if (irFunction.isSuspend)
-            irFunction.symbol.descriptor.psiElement ?: irFunction.parentAsClass.descriptor.psiElement
-        else
-            context.suspendLambdaToOriginalFunctionMap[irFunction.parentAsClass.attributeOwnerId]!!.symbol.descriptor.psiElement)
-                as KtElement
-
-    private fun IrFunction.continuationClass(): IrClass =
-            (body as IrBlockBody).statements.first { it is IrClass && it.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS } as IrClass
-
     private fun IrFunction.getVisibilityForDefaultArgumentStub(): Int =
         when (visibility) {
             Visibilities.PUBLIC -> Opcodes.ACC_PUBLIC
@@ -153,59 +121,48 @@ open class FunctionCodegen(
             else -> throw IllegalStateException("Default argument stub should be either public or package private: ${ir2string(this)}")
         }
 
-    private fun calculateMethodFlags(isStatic: Boolean): Int {
+    private fun IrFunction.calculateMethodFlags(): Int {
-        if (irFunction.origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER) {
+        if (origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER) {
-            return irFunction.getVisibilityForDefaultArgumentStub() or Opcodes.ACC_SYNTHETIC.let {
+            return getVisibilityForDefaultArgumentStub() or Opcodes.ACC_SYNTHETIC.let {
-                if (irFunction is IrConstructor) it else it or Opcodes.ACC_STATIC
+                if (this is IrConstructor) it else it or Opcodes.ACC_STATIC
             }
         }
 
-        val visibility = irFunction.getVisibilityAccessFlag()
+        val isVararg = valueParameters.lastOrNull()?.varargElementType != null
-        val staticFlag = if (isStatic) Opcodes.ACC_STATIC else 0
+        val isBridge = origin == IrDeclarationOrigin.BRIDGE || origin == IrDeclarationOrigin.BRIDGE_SPECIAL
-        val varargFlag = if (irFunction.valueParameters.lastOrNull()?.varargElementType != null) Opcodes.ACC_VARARGS else 0
+        val modalityFlag = when ((this as? IrSimpleFunction)?.modality) {
-        val deprecation = irFunction.deprecationFlags
-        val bridgeFlag = if (
-            irFunction.origin == IrDeclarationOrigin.BRIDGE ||
-            irFunction.origin == IrDeclarationOrigin.BRIDGE_SPECIAL
-        ) Opcodes.ACC_BRIDGE else 0
-        val modalityFlag = when ((irFunction as? IrSimpleFunction)?.modality) {
             Modality.FINAL -> when {
-                irFunction.origin == JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER -> 0
+                origin == JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER -> 0
-                classCodegen.irClass.isInterface && irFunction.body != null -> 0
+                parentAsClass.isInterface && body != null -> 0
-                !classCodegen.irClass.isAnnotationClass || irFunction.isStatic -> Opcodes.ACC_FINAL
+                parentAsClass.isAnnotationClass && !isStatic -> Opcodes.ACC_ABSTRACT
-                else -> Opcodes.ACC_ABSTRACT
+                else -> Opcodes.ACC_FINAL
             }
             Modality.ABSTRACT -> Opcodes.ACC_ABSTRACT
-            else -> if (classCodegen.irClass.isJvmInterface && irFunction.body == null) Opcodes.ACC_ABSTRACT else 0 //TODO transform interface modality on lowering to DefaultImpls
+            // TODO transform interface modality on lowering to DefaultImpls
+            else -> if (parentAsClass.isJvmInterface && body == null) Opcodes.ACC_ABSTRACT else 0
         }
-        val nativeFlag = if (irFunction.isExternal) Opcodes.ACC_NATIVE else 0
+        val isSynthetic = origin.isSynthetic || hasAnnotation(JVM_SYNTHETIC_ANNOTATION_FQ_NAME) ||
-        val syntheticFlag =
+                (isSuspend && Visibilities.isPrivate(visibility) && !isInline) || isReifiable()
-            if (irFunction.origin.isSynthetic || irFunction.hasAnnotation(JVM_SYNTHETIC_ANNOTATION_FQ_NAME) ||
+        val isStrict = hasAnnotation(STRICTFP_ANNOTATION_FQ_NAME)
-                    (irFunction.isSuspend && Visibilities.isPrivate(irFunction.visibility) && !irFunction.isInline) ||
+        val isSynchronized = hasAnnotation(SYNCHRONIZED_ANNOTATION_FQ_NAME)
-                    irFunction.isReifiable()
-            ) Opcodes.ACC_SYNTHETIC
-            else 0
-        val strictFpFlag = if (irFunction.hasAnnotation(STRICTFP_ANNOTATION_FQ_NAME)) Opcodes.ACC_STRICT else 0
-        val synchronizedFlag = if (irFunction.hasAnnotation(SYNCHRONIZED_ANNOTATION_FQ_NAME)) Opcodes.ACC_SYNCHRONIZED else 0
 
-        return visibility or
+        return getVisibilityAccessFlag() or modalityFlag or deprecationFlags or
-                modalityFlag or
+                (if (isStatic) Opcodes.ACC_STATIC else 0) or
-                staticFlag or
+                (if (isVararg) Opcodes.ACC_VARARGS else 0) or
-                varargFlag or
+                (if (isExternal) Opcodes.ACC_NATIVE else 0) or
-                deprecation or
+                (if (isBridge) Opcodes.ACC_BRIDGE else 0) or
-                nativeFlag or
+                (if (isSynthetic) Opcodes.ACC_SYNTHETIC else 0) or
-                bridgeFlag or
+                (if (isStrict) Opcodes.ACC_STRICT else 0) or
-                syntheticFlag or
+                (if (isSynchronized) Opcodes.ACC_SYNCHRONIZED else 0)
-                strictFpFlag or
-                synchronizedFlag
     }
 
     protected open fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor =
         classCodegen.visitor.newMethod(
             irFunction.OtherOrigin,
             flags,
-            signature.asmMethod.name, signature.asmMethod.descriptor,
+            signature.asmMethod.name,
-            if (flags.and(Opcodes.ACC_SYNTHETIC) != 0) null else signature.genericsSignature,
+            signature.asmMethod.descriptor,
+            signature.genericsSignature.takeIf { flags.and(Opcodes.ACC_SYNTHETIC) == 0 },
             getThrownExceptions(irFunction)?.toTypedArray()
         )
 
@@ -223,10 +180,7 @@ open class FunctionCodegen(
 
     private fun generateAnnotationDefaultValueIfNeeded(methodVisitor: MethodVisitor) {
         getAnnotationDefaultValueExpression()?.let { defaultValueExpression ->
-            val annotationCodegen = object: AnnotationCodegen(
+            val annotationCodegen = object : AnnotationCodegen(classCodegen, context) {
-                classCodegen,
-                context
-            ) {
                 override fun visitAnnotation(descr: String?, visible: Boolean): AnnotationVisitor {
                     return methodVisitor.visitAnnotationDefault()
                 }
@@ -249,26 +203,18 @@ open class FunctionCodegen(
             ?.expression
     }
 
-    private fun IrFrameMap.enterDispatchReceiver(parameter: IrValueParameter) {
+    private fun IrFunction.createFrameMapWithReceivers(): IrFrameMap {
-        val type = context.typeMapper.mapTypeAsDeclaration(parameter.type)
-        enter(parameter, type)
-    }
-
-    private fun createFrameMapWithReceivers(): IrFrameMap {
         val frameMap = IrFrameMap()
-
+        val receiver = if (this is IrConstructor) parentAsClass.thisReceiver else dispatchReceiverParameter
-        if (irFunction is IrConstructor) {
+        receiver?.let {
-            frameMap.enterDispatchReceiver(irFunction.constructedClass.thisReceiver!!)
+            frameMap.enter(it, context.typeMapper.mapTypeAsDeclaration(it.type))
-        } else if (irFunction.dispatchReceiverParameter != null) {
-            frameMap.enterDispatchReceiver(irFunction.dispatchReceiverParameter!!)
         }
-        irFunction.extensionReceiverParameter?.let {
+        extensionReceiverParameter?.let {
             frameMap.enter(it, context.typeMapper.mapType(it))
         }
-        for (parameter in irFunction.valueParameters) {
+        for (parameter in valueParameters) {
             frameMap.enter(parameter, context.typeMapper.mapType(parameter.type))
         }
-
         return frameMap
     }
 }