Support suspend fun main in JVM

#KT-17679 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/FunctionCodegen.java

@@ -68,6 +68,7 @@ import java.util.*;
 
 import static org.jetbrains.kotlin.builtins.KotlinBuiltIns.isNullableAny;
 import static org.jetbrains.kotlin.codegen.AsmUtil.*;
+import static org.jetbrains.kotlin.codegen.CodegenUtilKt.generateBridgeForMainFunctionIfNecessary;
 import static org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings.METHOD_FOR_FUNCTION;
 import static org.jetbrains.kotlin.descriptors.CallableMemberDescriptor.Kind.DECLARATION;
 import static org.jetbrains.kotlin.descriptors.ModalityKt.isOverridable;
@@ -244,6 +245,8 @@ public class FunctionCodegen {
             parentBodyCodegen.addAdditionalTask(new JvmStaticInCompanionObjectGenerator(functionDescriptor, origin, state, parentBodyCodegen));
         }
 
+        generateBridgeForMainFunctionIfNecessary(state, v, functionDescriptor, jvmSignature, origin);
+
         boolean isOpenSuspendInClass =
                 functionDescriptor.isSuspend() &&
                 functionDescriptor.getModality() != Modality.ABSTRACT && isOverridable(functionDescriptor) &&

compiler/backend/src/org/jetbrains/kotlin/codegen/codegenUtil.kt

@@ -13,7 +13,9 @@ import org.jetbrains.kotlin.builtins.UnsignedTypes
 import org.jetbrains.kotlin.codegen.context.CodegenContext
 import org.jetbrains.kotlin.codegen.context.FieldOwnerContext
 import org.jetbrains.kotlin.codegen.context.PackageContext
+import org.jetbrains.kotlin.codegen.coroutines.continuationAsmType
 import org.jetbrains.kotlin.codegen.coroutines.unwrapInitialDescriptorForSuspendFunction
+import org.jetbrains.kotlin.codegen.inline.NUMBERED_FUNCTION_PREFIX
 import org.jetbrains.kotlin.codegen.inline.ReificationArgument
 import org.jetbrains.kotlin.codegen.intrinsics.TypeIntrinsics
 import org.jetbrains.kotlin.codegen.optimization.common.asSequence
@@ -23,6 +25,7 @@ import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.deserialization.PLATFORM_DEPENDENT_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
+import org.jetbrains.kotlin.idea.MainFunctionDetector
 import org.jetbrains.kotlin.load.java.BuiltinMethodsWithSpecialGenericSignature.SpecialSignatureInfo
 import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.load.java.descriptors.JavaCallableMemberDescriptor
@@ -41,8 +44,11 @@ import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.resolve.isInlineClassType
+import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
+import org.jetbrains.kotlin.resolve.jvm.diagnostics.Synthetic
+import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedMemberDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedMemberDescriptor.CoroutinesCompatibilityMode
@@ -52,6 +58,7 @@ import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.utils.DFS
 import org.jetbrains.org.objectweb.asm.Label
+import org.jetbrains.org.objectweb.asm.Opcodes.*
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import org.jetbrains.org.objectweb.asm.commons.Method
@@ -464,4 +471,156 @@ fun recordCallLabelForLambdaArgument(declaration: KtFunctionLiteral, bindingTrac
     val call = callExpression.getResolvedCall(bindingTrace.bindingContext) ?: return
 
     storeLabelName(call.resultingDescriptor.name.asString())
-}
+}
+
+private val ARRAY_OF_STRINGS_TYPE = Type.getType("[Ljava/lang/String;")
+private val METHOD_DESCRIPTOR_FOR_MAIN = Type.getMethodDescriptor(Type.VOID_TYPE, ARRAY_OF_STRINGS_TYPE)
+
+fun generateBridgeForMainFunctionIfNecessary(
+    state: GenerationState,
+    packagePartClassBuilder: ClassBuilder,
+    functionDescriptor: FunctionDescriptor,
+    signatureOfRealDeclaration: JvmMethodGenericSignature,
+    origin: JvmDeclarationOrigin
+) {
+    val originElement = origin.element ?: return
+    if (functionDescriptor.name.asString() != "main" || !DescriptorUtils.isTopLevelDeclaration(functionDescriptor)) return
+    if (!MainFunctionDetector.isMain(functionDescriptor.unwrapInitialDescriptorForSuspendFunction(), false, true)) return
+
+    if (!functionDescriptor.isSuspend) return
+
+    val lambdaInternalName = generateLambdaForRunSuspend(
+        state,
+        originElement,
+        packagePartClassBuilder.thisName,
+        signatureOfRealDeclaration
+    )
+
+    packagePartClassBuilder.newMethod(
+        Synthetic(originElement, functionDescriptor),
+        ACC_PUBLIC or ACC_STATIC or ACC_SYNTHETIC,
+        "main",
+        METHOD_DESCRIPTOR_FOR_MAIN, null, null
+    ).apply {
+        visitCode()
+        visitTypeInsn(NEW, lambdaInternalName)
+        visitInsn(DUP)
+        visitVarInsn(ALOAD, 0)
+        visitMethodInsn(
+            INVOKESPECIAL,
+            lambdaInternalName,
+            "<init>",
+            METHOD_DESCRIPTOR_FOR_MAIN,
+            false
+        )
+
+        visitMethodInsn(
+            INVOKESTATIC,
+            "kotlin/coroutines/jvm/internal/RunSuspendKt", "runSuspend",
+            Type.getMethodDescriptor(
+                Type.VOID_TYPE,
+                Type.getObjectType(NUMBERED_FUNCTION_PREFIX + "1")
+            ),
+            false
+        )
+        visitInsn(RETURN)
+        visitEnd()
+    }
+}
+
+private fun generateLambdaForRunSuspend(
+    state: GenerationState,
+    originElement: PsiElement,
+    packagePartClassInternalName: String,
+    signatureOfRealDeclaration: JvmMethodGenericSignature
+): String {
+    val internalName = "$packagePartClassInternalName$$\$main"
+    val lambdaBuilder = state.factory.newVisitor(
+        JvmDeclarationOrigin.NO_ORIGIN,
+        Type.getObjectType(internalName),
+        originElement.containingFile
+    )
+
+    lambdaBuilder.defineClass(
+        originElement, state.classFileVersion,
+        ACC_FINAL or ACC_SUPER or ACC_SYNTHETIC,
+        internalName, null,
+        AsmTypes.LAMBDA.internalName,
+        arrayOf(NUMBERED_FUNCTION_PREFIX + "1")
+    )
+
+    lambdaBuilder.newField(
+        JvmDeclarationOrigin.NO_ORIGIN,
+        ACC_PRIVATE or ACC_FINAL,
+        "args",
+        ARRAY_OF_STRINGS_TYPE.descriptor, null, null
+    )
+
+    lambdaBuilder.newMethod(
+        JvmDeclarationOrigin.NO_ORIGIN,
+        AsmUtil.NO_FLAG_PACKAGE_PRIVATE or ACC_SYNTHETIC,
+        "<init>",
+        METHOD_DESCRIPTOR_FOR_MAIN, null, null
+    ).apply {
+        visitCode()
+        visitVarInsn(ALOAD, 0)
+        visitVarInsn(ALOAD, 1)
+        visitFieldInsn(
+            PUTFIELD,
+            lambdaBuilder.thisName,
+            "args",
+            ARRAY_OF_STRINGS_TYPE.descriptor
+        )
+
+        visitVarInsn(ALOAD, 0)
+        visitInsn(ICONST_1)
+        visitMethodInsn(
+            INVOKESPECIAL,
+            AsmTypes.LAMBDA.internalName,
+            "<init>",
+            Type.getMethodDescriptor(Type.VOID_TYPE, Type.INT_TYPE),
+            false
+        )
+        visitInsn(RETURN)
+        visitEnd()
+    }
+
+    lambdaBuilder.newMethod(
+        JvmDeclarationOrigin.NO_ORIGIN,
+        ACC_PUBLIC or ACC_FINAL or ACC_SYNTHETIC,
+        "invoke",
+        Type.getMethodDescriptor(AsmTypes.OBJECT_TYPE, AsmTypes.OBJECT_TYPE), null, null
+    ).apply {
+        visitCode()
+
+        visitVarInsn(ALOAD, 0)
+        visitFieldInsn(
+            GETFIELD,
+            lambdaBuilder.thisName,
+            "args",
+            ARRAY_OF_STRINGS_TYPE.descriptor
+        )
+
+        visitVarInsn(ALOAD, 1)
+        val continuationInternalName = state.languageVersionSettings.continuationAsmType().internalName
+
+        visitTypeInsn(
+            CHECKCAST,
+            continuationInternalName
+        )
+        visitMethodInsn(
+            INVOKESTATIC,
+            packagePartClassInternalName,
+            signatureOfRealDeclaration.asmMethod.name,
+            signatureOfRealDeclaration.asmMethod.descriptor,
+            false
+        )
+        visitInsn(ARETURN)
+        visitEnd()
+    }
+
+    writeSyntheticClassMetadata(lambdaBuilder, state)
+
+    lambdaBuilder.done()
+    return lambdaBuilder.thisName
+}

compiler/testData/diagnostics/tests/redeclarations/RedeclarationSuspendMainInMultiFile.kt

@@ -0,0 +1,5 @@
+// FILE: a.kt
+suspend fun main(args: Array<String>) {}
+
+// FILE: b.kt
+suspend fun main(args: Array<String>) {}

compiler/testData/diagnostics/tests/redeclarations/RedeclarationSuspendMainInMultiFile.txt

@@ -0,0 +1,4 @@
+package
+
+public suspend fun main(/*0*/ args: kotlin.Array<kotlin.String>): kotlin.Unit
+public suspend fun main(/*0*/ args: kotlin.Array<kotlin.String>): kotlin.Unit

compiler/testData/integration/smoke/helloAppSuspendMain/hello.compile.expected

@@ -0,0 +1 @@
+Return code: 0

compiler/testData/integration/smoke/helloAppSuspendMain/hello.kt

@@ -0,0 +1,35 @@
+package Hello
+
+import kotlin.concurrent.thread
+import kotlin.coroutines.suspendCoroutine
+import kotlin.coroutines.resume
+
+@kotlin.jvm.Volatile
+private var result = ""
+@kotlin.jvm.Volatile
+private var callback: Function0<Unit>? = null
+
+suspend fun appendAndSuspend(s: String) {
+    result += s
+
+    suspendCoroutine<Unit> { continuation ->
+        callback = {
+            continuation.resume(Unit)
+        }
+    }
+}
+
+suspend fun main(args: Array<String>) {
+    thread(isDaemon = true) {
+        while (true) {
+            val c = callback
+            c?.invoke()
+            Thread.sleep(500)
+        }
+    }
+
+    appendAndSuspend(args[0])
+    appendAndSuspend(args[1])
+    println(result)
+    callback = null
+}

compiler/testData/integration/smoke/helloAppSuspendMain/hello.run.expected

@@ -0,0 +1,4 @@
+OUT:
+OK
+
+Return code: 0

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -15926,6 +15926,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 runTest("compiler/testData/diagnostics/tests/redeclarations/RedeclarationMainInMultiFile.kt");
             }
 
+            @TestMetadata("RedeclarationSuspendMainInMultiFile.kt")
+            public void testRedeclarationSuspendMainInMultiFile() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/redeclarations/RedeclarationSuspendMainInMultiFile.kt");
+            }
+
             @TestMetadata("Redeclarations.kt")
             public void testRedeclarations() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/redeclarations/Redeclarations.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -15926,6 +15926,11 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                 runTest("compiler/testData/diagnostics/tests/redeclarations/RedeclarationMainInMultiFile.kt");
             }
 
+            @TestMetadata("RedeclarationSuspendMainInMultiFile.kt")
+            public void testRedeclarationSuspendMainInMultiFile() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/redeclarations/RedeclarationSuspendMainInMultiFile.kt");
+            }
+
             @TestMetadata("Redeclarations.kt")
             public void testRedeclarations() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/redeclarations/Redeclarations.kt");

compiler/tests/org/jetbrains/kotlin/integration/CompilerSmokeTest.java

@@ -41,6 +41,13 @@ public class CompilerSmokeTest extends CompilerSmokeTestBase {
         run("hello.run", "-cp", jar, "Hello.HelloKt");
     }
 
+    public void testHelloAppSuspendMain() throws Exception {
+        String jar = tmpdir.getAbsolutePath() + File.separator + "hello.jar";
+
+        assertEquals("compilation failed", 0, runCompiler("hello.compile", "-include-runtime", "hello.kt", "-d", jar));
+        run("hello.run", "-cp", jar, "Hello.HelloKt", "O", "K");
+    }
+
     public void testCompilationFailed() throws Exception {
         String jar = tmpdir.getAbsolutePath() + File.separator + "smoke.jar";
 

idea/testData/run/MainInTest/module/src/suspendMain.kt

@@ -0,0 +1,4 @@
+package suspendMain
+
+suspend fun main(args: Array<String>) { // yes
+}