JVM_IR: Support is/as SuspendFunction operators

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

@@ -40,6 +40,7 @@ class JvmSymbols(
     private val storageManager = LockBasedStorageManager(this::class.java.simpleName)
     private val kotlinPackage: IrPackageFragment = createPackage(FqName("kotlin"))
     private val kotlinCoroutinesPackage: IrPackageFragment = createPackage(FqName("kotlin.coroutines"))
+    private val kotlinCoroutinesJvmInternalPackage: IrPackageFragment = createPackage(FqName("kotlin.coroutines.jvm.internal"))
     private val kotlinJvmPackage: IrPackageFragment = createPackage(FqName("kotlin.jvm"))
     private val kotlinJvmInternalPackage: IrPackageFragment = createPackage(FqName("kotlin.jvm.internal"))
     private val kotlinJvmFunctionsPackage: IrPackageFragment = createPackage(FqName("kotlin.jvm.functions"))
@@ -81,6 +82,7 @@ class JvmSymbols(
             parent = when (fqName.parent().asString()) {
                 "kotlin" -> kotlinPackage
                 "kotlin.coroutines" -> kotlinCoroutinesPackage
+                "kotlin.coroutines.jvm.internal" -> kotlinCoroutinesJvmInternalPackage
                 "kotlin.jvm.internal" -> kotlinJvmInternalPackage
                 "kotlin.jvm.functions" -> kotlinJvmFunctionsPackage
                 "java.lang" -> javaLangPackage
@@ -169,6 +171,8 @@ class JvmSymbols(
             klass.addTypeParameter("T", irBuiltIns.anyNType, Variance.IN_VARIANCE)
         }
 
+    val suspendFunctionInterface: IrClassSymbol = createClass(FqName("kotlin.coroutines.jvm.internal.SuspendFunction"), ClassKind.INTERFACE)
+
     val lambdaClass: IrClassSymbol = createClass(FqName("kotlin.jvm.internal.Lambda")) { klass ->
         klass.addConstructor().apply {
             addValueParameter("arity", irBuiltIns.intType)

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

@@ -103,14 +103,12 @@ internal fun IrFunction.isInvokeOfSuspendLambda(context: JvmBackendContext): Boo
 internal fun IrFunction.isInvokeSuspendOfContinuation(context: JvmBackendContext): Boolean =
     name.asString() == INVOKE_SUSPEND_METHOD_NAME && parentAsClass in context.suspendFunctionContinuations.values
 
-internal fun IrFunction.isInvokeOfCallableReference(): Boolean = isSuspend && name.asString() == "invoke" && (parent as? IrClass)?.let {
-    // TODO: Should we use different origin for lowered callable references?
-    it.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL && it.functions.any { it.name.asString() == "getSignature" }
-} == true
+internal fun IrFunction.isInvokeOfSuspendCallableReference(): Boolean = isSuspend && name.asString() == "invoke" &&
+        (parent as? IrClass)?.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
 
 internal fun IrFunction.isKnownToBeTailCall(): Boolean =
     origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER || origin == JvmLoweredDeclarationOrigin.SYNTHETIC_ACCESSOR ||
-            isInvokeOfCallableReference()
+            isInvokeOfSuspendCallableReference()
 
 internal fun IrFunction.shouldNotContainSuspendMarkers(context: JvmBackendContext): Boolean =
     isInvokeSuspendOfContinuation(context) || isKnownToBeTailCall()

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

@@ -1,6 +1,6 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- * that can be found in the license/LICENSE.txt file.
+ * 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
@@ -35,6 +35,8 @@ import org.jetbrains.org.objectweb.asm.Type
 
 class IrTypeMapper(private val context: JvmBackendContext) : KotlinTypeMapperBase() {
     internal val typeSystem = IrTypeCheckerContext(context.irBuiltIns)
+    private val IrTypeArgument.adjustedType
+        get() = (this as? IrTypeProjection)?.type ?: context.irBuiltIns.anyNType
 
     override fun mapClass(classifier: ClassifierDescriptor): Type =
         when (classifier) {
@@ -99,8 +101,8 @@ class IrTypeMapper(private val context: JvmBackendContext) : KotlinTypeMapperBas
 
         if (type.isSuspendFunction()) {
             val arguments =
-                type.arguments.dropLast(1).map { (it as IrTypeProjection).type } +
-                        context.ir.symbols.continuationClass.typeWith((type.arguments.last() as IrTypeProjection).type) +
+                type.arguments.dropLast(1).map { it.adjustedType } +
+                        context.ir.symbols.continuationClass.typeWith(type.arguments.last().adjustedType) +
                         context.irBuiltIns.anyNType
             val runtimeFunctionType = context.referenceClass(context.builtIns.getFunction(arguments.size - 1)).typeWith(arguments)
             return mapType(runtimeFunctionType, mode, sw)
@@ -251,4 +253,4 @@ class IrTypeMapper(private val context: JvmBackendContext) : KotlinTypeMapperBas
 
     private fun boxTypeIfNeeded(possiblyPrimitiveType: Type, needBoxedType: Boolean): Type =
         if (needBoxedType) AsmUtil.boxType(possiblyPrimitiveType) else possiblyPrimitiveType
-}
+}

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

@@ -18,7 +18,7 @@ import org.jetbrains.kotlin.backend.common.pop
 import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrBlock
-import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeOfCallableReference
+import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeOfSuspendCallableReference
 import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.config.coroutinesPackageFqName
 import org.jetbrains.kotlin.descriptors.Modality
@@ -433,7 +433,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             override fun visitFunction(declaration: IrFunction) {
                 super.visitFunction(declaration)
                 if (declaration.isSuspend && declaration !in suspendLambdas && !declaration.isInline &&
-                    !declaration.isInvokeOfCallableReference()
+                    !declaration.isInvokeOfSuspendCallableReference()
                 ) {
                     result.add(declaration)
                 }

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

@@ -188,6 +188,7 @@ internal class CallableReferenceLowering(private val context: JvmBackendContext)
             superTypes += superType
             if (samSuperType == null)
                 superTypes += functionSuperClass.typeWith(parameterTypes)
+            if (irFunctionReference.isSuspend) superTypes += context.ir.symbols.suspendFunctionInterface.typeWith()
             createImplicitParameterDeclarationWithWrappedDescriptor()
             copyAttributes(irFunctionReference)
         }

compiler/testData/codegen/box/coroutines/featureIntersection/callableReference/function/getArityViaFunctionImpl.kt

@@ -17,26 +17,26 @@ fun test(f: Function<*>, arity: Int) {
 }
 
 suspend fun foo(s: String, i: Int) {}
-//class A {
-//    suspend fun bar(s: String, i: Int) {}
-//}
-//
-//suspend fun Double.baz(s: String, i: Int) {}
+class A {
+    suspend fun bar(s: String, i: Int) {}
+}
+
+suspend fun Double.baz(s: String, i: Int) {}
 
 fun box(): String {
     test(::foo, 2 + 1)
-//    test(A::bar, 3 + 1)
-//    test(Double::baz, 3 + 1)
-//
-//    test(::box, 0)
-//
-//    suspend fun local(x: Int) {}
-//    test(::local, 1 + 1)
-//
-//    // TODO: Uncomment when `suspend fun` will be supported
-//    // test(suspend fun(s: String) = s, 1)
-//    // test(suspend fun(){}, 0)
-//    test(suspend {}, 1)
+    test(A::bar, 3 + 1)
+    test(Double::baz, 3 + 1)
+
+    test(::box, 0)
+
+    suspend fun local(x: Int) {}
+    test(::local, 1 + 1)
+
+    // TODO: Uncomment when `suspend fun` will be supported
+    // test(suspend fun(s: String) = s, 1)
+    // test(suspend fun(){}, 0)
+    test(suspend {}, 1)
 
     return "OK"
 }

compiler/testData/codegen/box/coroutines/featureIntersection/suspendFunctionIsAs.kt

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND: NATIVE
 // IGNORE_BACKEND: JS_IR, JS
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES