Support inline suspend functions

A lot of additional work was required to support them:
- Suspension points are being identified by two markers
  instead of one pointing to suspend function call
- Approach with replacing return type of suspend function does not work anymore.
  So we decode suspension return type as an argument for begin marker
- It became necessary to perform variables liveness analysis
  (see comment in org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor.spillVariables)

.idea/dictionaries/dzharkov.xml

@@ -4,6 +4,7 @@
       <w>checkcast</w>
       <w>coroutine</w>
       <w>insn</w>
+      <w>liveness</w>
     </words>
   </dictionary>
 </component>

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

@@ -2592,7 +2592,15 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
 
         Callable callable = resolveToCallable(fd, superCallTarget != null, resolvedCall);
 
-        return callable.invokeMethodWithArguments(resolvedCall, receiver, this);
+
+        StackValue result = callable.invokeMethodWithArguments(resolvedCall, receiver, this);
+
+        if (CoroutineCodegenUtilKt.isSuspensionPoint(resolvedCall)) {
+            // Suspension points should behave like they leave actual values on stack, while real methods return VOID
+            return new OperationStackValue(getSuspensionReturnTypeByResolvedCall(resolvedCall), ((OperationStackValue) result).getLambda());
+        }
+
+        return result;
     }
 
     @Nullable
@@ -2685,14 +2693,20 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         }
 
         if (isSuspensionPoint) {
+            v.aconst(getSuspensionReturnTypeByResolvedCall(resolvedCall).getDescriptor());
             v.invokestatic(
                     CoroutineCodegenUtilKt.COROUTINE_MARKER_OWNER,
-                    CoroutineCodegenUtilKt.SUSPENSION_POINT_MARKER_NAME, "()V", false);
+                    CoroutineCodegenUtilKt.BEFORE_SUSPENSION_POINT_MARKER_NAME,
+                    "(Ljava/lang/String;)V", false);
         }
 
         callGenerator.genCall(callableMethod, resolvedCall, defaultMaskWasGenerated, this);
 
         if (isSuspensionPoint) {
+            StackValue.coerce(Type.VOID_TYPE, getSuspensionReturnTypeByResolvedCall(resolvedCall), v);
+            v.invokestatic(
+                    CoroutineCodegenUtilKt.COROUTINE_MARKER_OWNER,
+                    CoroutineCodegenUtilKt.AFTER_SUSPENSION_POINT_MARKER_NAME, "()V", false);
             addInlineMarker(v, false);
         }
 
@@ -2703,6 +2717,18 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         }
     }
 
+    @NotNull
+    private Type getSuspensionReturnTypeByResolvedCall(@NotNull ResolvedCall<?> resolvedCall) {
+        assert resolvedCall.getResultingDescriptor() instanceof SimpleFunctionDescriptor
+                : "Suspension point resolved call should be built on SimpleFunctionDescriptor";
+
+        KotlinType returnType = org.jetbrains.kotlin.resolve.coroutine.CoroutineUtilKt
+                .getSuspensionPointReturnType((SimpleFunctionDescriptor) resolvedCall.getResultingDescriptor());
+
+        assert returnType != null : "Return type of suspension point should not be null";
+        return typeMapper.mapType(returnType);
+    }
+
     @NotNull
     private CallGenerator getOrCreateCallGenerator(
             @NotNull CallableDescriptor descriptor,
@@ -2719,7 +2745,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
 
         if (!isInline) return defaultCallGenerator;
 
-        FunctionDescriptor original = DescriptorUtils.unwrapFakeOverride((FunctionDescriptor) descriptor.getOriginal());
+        FunctionDescriptor original =
+                unwrapInitialSignatureDescriptor(DescriptorUtils.unwrapFakeOverride((FunctionDescriptor) descriptor.getOriginal()));
         if (isDefaultCompilation) {
             return new InlineCodegenForDefaultBody(original, this, state);
         }
@@ -2728,6 +2755,12 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         }
     }
 
+    @NotNull
+    private static FunctionDescriptor unwrapInitialSignatureDescriptor(@NotNull FunctionDescriptor function) {
+        if (function.getInitialSignatureDescriptor() != null) return function.getInitialSignatureDescriptor();
+        return function;
+    }
+
     @NotNull
     protected CallGenerator getOrCreateCallGeneratorForDefaultImplBody(@NotNull FunctionDescriptor descriptor, @Nullable KtNamedFunction function) {
         return getOrCreateCallGenerator(descriptor, function, null, true);

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformationClassBuilder.kt

@@ -16,11 +16,12 @@
 
 package org.jetbrains.kotlin.codegen.coroutines
 
+import com.intellij.util.containers.Stack
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.optimization.DeadCodeEliminationMethodTransformer
 import org.jetbrains.kotlin.codegen.optimization.FixStackWithLabelNormalizationMethodTransformer
 import org.jetbrains.kotlin.codegen.optimization.common.OptimizationBasicInterpreter
-import org.jetbrains.kotlin.codegen.optimization.common.asSequence
+import org.jetbrains.kotlin.codegen.optimization.common.analyzeLiveness
 import org.jetbrains.kotlin.codegen.optimization.common.insnListOf
 import org.jetbrains.kotlin.codegen.optimization.common.removeEmptyCatchBlocks
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
@@ -31,7 +32,6 @@ import org.jetbrains.org.objectweb.asm.MethodVisitor
 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
 import org.jetbrains.org.objectweb.asm.tree.*
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue
 
@@ -78,12 +78,14 @@ class CoroutineTransformerMethodVisitor(
         if (suspensionPoints.isEmpty()) return
 
         // Spill stack to variables before suspension points, try/catch blocks
-        FixStackWithLabelNormalizationMethodTransformer().transform("fake", methodNode)
+        FixStackWithLabelNormalizationMethodTransformer().transform(classBuilder.thisName, methodNode)
 
         // Remove unreachable suspension points
         // If we don't do this, then relevant frames will not be analyzed, that is unexpected from point of view of next steps (e.g. variable spilling)
-        DeadCodeEliminationMethodTransformer().transform("fake", methodNode)
+        DeadCodeEliminationMethodTransformer().transform(classBuilder.thisName, methodNode)
-        suspensionPoints.removeAll { it.suspensionCall.next == null && it.suspensionCall.previous == null }
+        suspensionPoints.removeAll {
+            it.suspensionCallBegin.next == null && it.suspensionCallBegin.previous == null
+        }
 
         processUninitializedStores(methodNode)
 
@@ -116,50 +118,98 @@ class CoroutineTransformerMethodVisitor(
             })
         }
 
+        dropSuspensionMarkers(methodNode, suspensionPoints)
         methodNode.removeEmptyCatchBlocks()
 
     }
 
     private fun collectSuspensionPoints(methodNode: MethodNode): MutableList<SuspensionPoint> {
         val suspensionPoints = mutableListOf<SuspensionPoint>()
+        val beforeSuspensionPointMarkerStack = Stack<MethodInsnNode>()
 
-        for (methodInsn in methodNode.instructions.asSequence().filterIsInstance<MethodInsnNode>()) {
+        for (methodInsn in methodNode.instructions.toArray().filterIsInstance<MethodInsnNode>()) {
             if (methodInsn.owner != COROUTINE_MARKER_OWNER) continue
 
             when (methodInsn.name) {
-                SUSPENSION_POINT_MARKER_NAME -> {
+                BEFORE_SUSPENSION_POINT_MARKER_NAME -> {
-                    assert(methodInsn.next is MethodInsnNode) {
+                    beforeSuspensionPointMarkerStack.add(methodInsn)
-                        "Expected method call instruction after suspension point, but ${methodInsn.next} found"
+                }
+
+                AFTER_SUSPENSION_POINT_MARKER_NAME -> {
+                    assert(beforeSuspensionPointMarkerStack.isNotEmpty()) { "Expected non-empty stack beforeSuspensionPointMarkerStack" }
+                    val beforeSuspensionPointMarker = beforeSuspensionPointMarkerStack.pop()
+                    assert(beforeSuspensionPointMarker.previous is LdcInsnNode) {
+                        "Instruction before BEFORE_SUSPENSION_POINT_MARKER_NAME should be LDC with expected return type, " +
+                        "but ${beforeSuspensionPointMarker.previous} was found"
                     }
 
-                    suspensionPoints.add(SuspensionPoint(methodInsn.next as MethodInsnNode))
+                    val fakeReturnValueInsns = mutableListOf<AbstractInsnNode>()
+
+                    val suspensionPoint = SuspensionPoint(
+                            suspensionCallBegin = beforeSuspensionPointMarker,
+                            suspensionCallEnd = methodInsn,
+                            fakeReturnValueInsns = fakeReturnValueInsns,
+                            returnType = Type.getType((beforeSuspensionPointMarker.previous as LdcInsnNode).cst as String))
+                    suspensionPoints.add(suspensionPoint)
+
+                    // Drop type info from marker
+                    methodNode.instructions.remove(beforeSuspensionPointMarker.previous)
+                    beforeSuspensionPointMarker.desc = "()V"
+
+                    if (suspensionPoint.returnType != Type.VOID_TYPE) {
+                        val previous = suspensionPoint.suspensionCallEnd.previous
+                        assert(previous.opcode in DEFAULT_VALUE_OPCODES) {
+                            "Expected on of default value bytecodes, but ${previous.opcode} was found"
+                        }
+                        fakeReturnValueInsns.add(previous)
+
+                        if (previous.opcode == Opcodes.ACONST_NULL) {
+                            // Checkcast is needed to tell analyzer what type exactly should be returned from suspension point
+                            val checkCast = TypeInsnNode(Opcodes.CHECKCAST, suspensionPoint.returnType.internalName)
+                            methodNode.instructions.insert(previous, checkCast)
+                            fakeReturnValueInsns.add(checkCast)
+                        }
+                    }
                 }
             }
         }
 
-        // Drop markers
+        assert(beforeSuspensionPointMarkerStack.isEmpty()) { "Unbalanced suspension markers stack" }
-        suspensionPoints.forEach { methodNode.instructions.remove(it.suspensionCall.previous) }
 
         return suspensionPoints
     }
 
+    private fun dropSuspensionMarkers(methodNode: MethodNode, suspensionPoints: List<SuspensionPoint>) {
+        // Drop markers
+        suspensionPoints.forEach {
+            // before-marker
+            methodNode.instructions.remove(it.suspensionCallBegin)
+            // after-marker
+            methodNode.instructions.remove(it.suspensionCallEnd)
+        }
+    }
+
     private fun spillVariables(suspensionPoints: List<SuspensionPoint>, methodNode: MethodNode) {
         val instructions = methodNode.instructions
-        val frames = MethodTransformer.analyze("fake", methodNode, OptimizationBasicInterpreter())
+        val frames = MethodTransformer.analyze(classBuilder.thisName, methodNode, OptimizationBasicInterpreter())
         fun AbstractInsnNode.index() = instructions.indexOf(this)
 
         // We postpone these actions because they change instruction indices that we use when obtaining frames
         val postponedActions = mutableListOf<() -> Unit>()
         val maxVarsCountByType = mutableMapOf<Type, Int>()
+        val livenessFrames = analyzeLiveness(methodNode)
 
         for (suspension in suspensionPoints) {
-            val call = suspension.suspensionCall
+            val suspensionCallBegin = suspension.suspensionCallBegin
-            assert(frames[call.next.index()].stackSize == (if (Type.getReturnType(call.desc).sort == Type.VOID) 0 else 1)) {
+            val suspensionCallEnd = suspension.suspensionCallEnd
+            assert(frames[suspensionCallEnd.next.index()].stackSize == (if (suspension.returnType.sort == Type.VOID) 0 else 1)) {
                 "Stack should be spilled before suspension call"
             }
 
             val variableTypeInTable = arrayOfNulls<Type>(methodNode.maxLocals)
-            methodNode.localVariables.filter { it.start.index() < call.index() && call.index() < it.end.index() }.forEach {
+            methodNode.localVariables.filter {
+                it.start.index() < suspensionCallBegin.index() && suspensionCallEnd.index() < it.end.index()
+            }.forEach {
                 val type = Type.getType(it.desc)
                 variableTypeInTable[it.index] = type
                 if (type.size == 2) {
@@ -167,9 +217,24 @@ class CoroutineTransformerMethodVisitor(
                 }
             }
 
-            val frame = frames[call.index()]
+            val frame = frames[suspensionCallBegin.index()]
             val localsCount = frame.locals
             val varsCountByType = mutableMapOf<Type, Int>()
+
+            // We consider variable liveness to avoid problems with inline suspension functions:
+            // <spill variables>
+            // <inline suspension call with new variables initialized> *
+            // RETURN (appears only on further transformation phase)
+            // ...
+            // <spill variables before next suspension point>
+            //
+            // The problem is that during current phase (before inserting RETURN opcode) we suppose variables generated
+            // within inline suspension point as correctly initialized, thus trying to spill them.
+            // While after RETURN introduction these variables become uninitialized (at the same time they can't be used further).
+            // So we only spill variables that are alive at the begin of suspension point.
+            // NB: it's also rather useful for sake of optimization
+            val livenessFrame = livenessFrames[suspensionCallBegin.index()]
+
             // 0 - this
             // 1 - continuation argument
             // 2 - continuation exception
@@ -179,6 +244,7 @@ class CoroutineTransformerMethodVisitor(
                             .filter {
                                 val (index, value) = it
                                 value != BasicValue.UNINITIALIZED_VALUE
+                                    && livenessFrame.isAlive(index)
                                     && variableTypeInTable[index] !== UNINITIALIZED_TYPE
                                     && (variableTypeInTable[index] == null || variableTypeInTable[index]?.sort == value.type.sort)
                             }
@@ -195,7 +261,7 @@ class CoroutineTransformerMethodVisitor(
                 postponedActions.add {
                     with(instructions) {
                         // store variable before suspension call
-                        insertBefore(call, withInstructionAdapter {
+                        insertBefore(suspension.suspensionCallBegin, withInstructionAdapter {
                             load(0, AsmTypes.OBJECT_TYPE)
                             load(index, type)
                             StackValue.coerce(type, normalizedType, this)
@@ -203,7 +269,7 @@ class CoroutineTransformerMethodVisitor(
                         })
 
                         // restore variable after suspension call
-                        insert(call.tryCatchBlockEndLabelAfterSuspensionCall, withInstructionAdapter {
+                        insert(suspension.tryCatchBlockEndLabelAfterSuspensionCall, withInstructionAdapter {
                             load(0, AsmTypes.OBJECT_TYPE)
                             getfield(classBuilder.thisName, fieldName, normalizedType.descriptor)
                             StackValue.coerce(normalizedType, type, this)
@@ -230,33 +296,34 @@ class CoroutineTransformerMethodVisitor(
         }
     }
 
-    private val MethodInsnNode.tryCatchBlockEndLabelAfterSuspensionCall: LabelNode
+    /**
+     * See 'splitTryCatchBlocksContainingSuspensionPoint'
+     */
+    private val SuspensionPoint.tryCatchBlockEndLabelAfterSuspensionCall: LabelNode
         get() {
-            assert(next is LabelNode) {
+            assert(suspensionCallEnd.next is LabelNode) {
-                "Next instruction after ${this} should be a label, but ${next.javaClass}/${next.opcode} was found"
+                "Next instruction after ${this} should be a label, but " +
+                "${suspensionCallEnd.next.javaClass}/${suspensionCallEnd.next.opcode} was found"
             }
 
-            return next as LabelNode
+            return suspensionCallEnd.next as LabelNode
         }
 
     private fun transformCallAndReturnContinuationLabel(id: Int, suspension: SuspensionPoint, methodNode: MethodNode): LabelNode {
-        val call = suspension.suspensionCall
-        val method = Method(call.name, call.desc)
-        call.desc = Method(method.name, Type.VOID_TYPE, method.argumentTypes).descriptor
-
         val continuationLabel = LabelNode()
-
         with(methodNode.instructions) {
             // Save state
-            insertBefore(call,
+            insertBefore(suspension.suspensionCallBegin,
                          insnListOf(
                                  VarInsnNode(Opcodes.ALOAD, 0),
                                  *withInstructionAdapter { iconst(id) }.toArray(),
                                  FieldInsnNode(
                                          Opcodes.PUTFIELD, classBuilder.thisName, COROUTINE_LABEL_FIELD_NAME, Type.INT_TYPE.descriptor)))
 
+            // Drop default value that purpose was to simulate returned value by suspension method
+            suspension.fakeReturnValueInsns.forEach { methodNode.instructions.remove(it) }
 
-            insert(call.tryCatchBlockEndLabelAfterSuspensionCall, withInstructionAdapter {
+            insert(suspension.tryCatchBlockEndLabelAfterSuspensionCall, withInstructionAdapter {
                 // Exit
                 areturn(Type.VOID_TYPE)
                 // Mark place for continuation
@@ -274,7 +341,6 @@ class CoroutineTransformerMethodVisitor(
             }
             remove(possibleTryCatchBlockStart.previous)
 
-
             insert(possibleTryCatchBlockStart, withInstructionAdapter {
                 // Check if resumeWithException has been called
                 load(2, AsmTypes.OBJECT_TYPE)
@@ -288,7 +354,7 @@ class CoroutineTransformerMethodVisitor(
 
                 // Load continuation argument just like suspending function returns it
                 load(1, AsmTypes.OBJECT_TYPE)
-                StackValue.coerce(AsmTypes.OBJECT_TYPE, method.returnType, this)
+                StackValue.coerce(AsmTypes.OBJECT_TYPE, suspension.returnType, this)
             })
         }
 
@@ -300,7 +366,11 @@ class CoroutineTransformerMethodVisitor(
     // otherwise they would be treated as uninitialized within catch-block while they can be used there
     // How suspension point area will look like after all transformations:
     // <spill variables>
-    // INVOKEVIRTUAL suspensionMethod()
+    // INVOKESTATIC beforeSuspensionMarker
+    // INVOKEVIRTUAL suspensionMethod()V
+    // ACONST_NULL -- default value
+    // CHECKCAST SomeType
+    // INVOKESTATIC afterSuspensionMarker
     // L1: -- end of all TCB's that are containing the suspension point (inserted by this method)
     // RETURN
     // L2: -- continuation label (used for the TABLESWITCH)
@@ -309,11 +379,12 @@ class CoroutineTransformerMethodVisitor(
     // ...
     private fun splitTryCatchBlocksContainingSuspensionPoint(methodNode: MethodNode, suspensionPoint: SuspensionPoint) {
         val instructions = methodNode.instructions
-        val indexOfSuspension = instructions.indexOf(suspensionPoint.suspensionCall)
+        val beginIndex = instructions.indexOf(suspensionPoint.suspensionCallBegin)
+        val endIndex = instructions.indexOf(suspensionPoint.suspensionCallEnd)
 
         val firstLabel = LabelNode()
         val secondLabel = LabelNode()
-        instructions.insert(suspensionPoint.suspensionCall, firstLabel)
+        instructions.insert(suspensionPoint.suspensionCallEnd, firstLabel)
         // NOP is needed to preventing these label merge
         // Here between these labels additional instructions are supposed to be inserted (variables spilling, etc.)
         instructions.insert(firstLabel, InsnNode(Opcodes.NOP))
@@ -322,11 +393,15 @@ class CoroutineTransformerMethodVisitor(
         methodNode.tryCatchBlocks =
                 methodNode.tryCatchBlocks.flatMap {
                     val isContainingSuspensionPoint =
-                            instructions.indexOf(it.start) < indexOfSuspension && indexOfSuspension < instructions.indexOf(it.end)
+                            instructions.indexOf(it.start) < beginIndex && beginIndex < instructions.indexOf(it.end)
 
-                    if (isContainingSuspensionPoint)
+                    if (isContainingSuspensionPoint) {
+                        assert(instructions.indexOf(it.start) < endIndex && endIndex < instructions.indexOf(it.end)) {
+                            "Try catch block containing marker before suspension point should also contain the marker after suspension point"
+                        }
                         listOf(TryCatchBlockNode(it.start, firstLabel, it.handler, it.type),
                                TryCatchBlockNode(secondLabel, it.end, it.handler, it.type))
+                    }
                     else
                         listOf(it)
                 }
@@ -384,8 +459,28 @@ private fun Type.normalize() =
         else -> this
     }
 
-private class SuspensionPoint(val suspensionCall: MethodInsnNode) {
+/**
+ * Suspension call may consists of several instructions:
+ * INVOKESTATIC beforeSuspensionMarker
+ * INVOKEVIRTUAL suspensionMethod()V // actually it could be some inline method instead of plain call
+ * ACONST_NULL // It's only needed when the suspension point returns something beside VOID
+ * CHECKCAST Type
+ * INVOKESTATIC afterSuspensionMarker
+ */
+private class SuspensionPoint(
+        // INVOKESTATIC beforeSuspensionMarker
+        val suspensionCallBegin: AbstractInsnNode,
+        // INVOKESTATIC afterSuspensionMarker
+        val suspensionCallEnd: AbstractInsnNode,
+        val fakeReturnValueInsns: List<AbstractInsnNode>,
+        val returnType: Type
+) {
     lateinit var tryCatchBlocksContinuationLabel: LabelNode
 }
 
 private val UNINITIALIZED_TYPE = Type.getObjectType("uninitialized")
+private val DEFAULT_VALUE_OPCODES =
+        setOf(Opcodes.ICONST_0, Opcodes.LCONST_0, Opcodes.FCONST_0, Opcodes.DCONST_0, Opcodes.ACONST_NULL,
+              // GETSTATIC Unit.Instance
+              Opcodes.GETSTATIC)
+

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/coroutineCodegenUtil.kt

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.calls.util.CallMaker
+import org.jetbrains.kotlin.resolve.coroutine.REPLACED_SUSPENSION_POINT_KEY
 import org.jetbrains.kotlin.resolve.coroutine.SUSPENSION_POINT_KEY
 import org.jetbrains.kotlin.types.TypeConstructorSubstitution
 import org.jetbrains.kotlin.types.TypeSubstitutor
@@ -41,7 +42,8 @@ import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 val CONTINUATION_METHOD_ANNOTATION_DESC = "Lkotlin/ContinuationMethod;"
 
 const val COROUTINE_MARKER_OWNER = "kotlin/coroutines/Markers"
-const val SUSPENSION_POINT_MARKER_NAME = "suspensionPoint"
+const val BEFORE_SUSPENSION_POINT_MARKER_NAME = "beforeSuspensionPoint"
+const val AFTER_SUSPENSION_POINT_MARKER_NAME = "afterSuspensionPoint"
 const val HANDLE_EXCEPTION_MARKER_NAME = "handleException"
 const val HANDLE_EXCEPTION_ARGUMENT_MARKER_NAME = "handleExceptionArgument"
 
@@ -54,7 +56,7 @@ data class ResolvedCallWithRealDescriptor(val resolvedCall: ResolvedCall<*>, val
 // E.g. `fun <V> await(f: CompletableFuture<V>): V` instead of `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`
 // See `createCoroutineSuspensionFunctionView` and it's usages for clarification
 // But for call generation it's convenient to have `machine` (continuation) parameter/argument within resolvedCall.
-// So this function returns resolved call with descriptor looking like `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): V`
+// So this function returns resolved call with descriptor looking like `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`
 // and fake `this` expression that used as argument for second parameter
 fun ResolvedCall<*>.replaceSuspensionFunctionViewWithRealDescriptor(
         project: Project
@@ -63,11 +65,14 @@ fun ResolvedCall<*>.replaceSuspensionFunctionViewWithRealDescriptor(
     if (!isSuspensionPoint()) return null
 
     val initialSignatureDescriptor = function.initialSignatureDescriptor ?: return null
+    if (function.getUserData(REPLACED_SUSPENSION_POINT_KEY) == true) return null
+
     val newCandidateDescriptor =
             initialSignatureDescriptor.createCustomCopy {
-                // Here we know that last parameter should be Continuation<T> where T is return type
+                setPreserveSourceElement()
-                setReturnType(it.valueParameters.last().type.arguments.single().type)
+                setSignatureChange()
                 putUserData(SUSPENSION_POINT_KEY, true)
+                putUserData(REPLACED_SUSPENSION_POINT_KEY, true)
             }
 
     val newCall = ResolvedCallImpl(
@@ -88,8 +93,12 @@ fun ResolvedCall<*>.replaceSuspensionFunctionViewWithRealDescriptor(
             newCandidateDescriptor.valueParameters.last(),
             ExpressionValueArgument(arguments))
 
+    val newTypeArguments = newCandidateDescriptor.typeParameters.map {
+        Pair(it, typeArguments[candidateDescriptor.typeParameters[it.index]]!!.asTypeProjection())
+    }.toMap()
+
     newCall.setResultingSubstitutor(
-            TypeConstructorSubstitution.createByParametersMap(typeArguments.mapValues { it.value.asTypeProjection() }).buildSubstitutor())
+            TypeConstructorSubstitution.createByParametersMap(newTypeArguments).buildSubstitutor())
 
     return ResolvedCallWithRealDescriptor(newCall, thisExpression)
 }

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/ControlFlowGraph.kt

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.codegen.optimization.common
+
+import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
+import org.jetbrains.org.objectweb.asm.tree.InsnList
+import org.jetbrains.org.objectweb.asm.tree.MethodNode
+import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue
+
+
+class ControlFlowGraph private constructor(private val insns: InsnList) {
+    private val edges: Array<MutableList<Int>> = Array(insns.size()) { arrayListOf<Int>() }
+
+    fun getSuccessorsIndices(insn: AbstractInsnNode): List<Int> = edges[insns.indexOf(insn)]
+
+    companion object {
+        @JvmStatic
+        fun build(node: MethodNode): ControlFlowGraph {
+            val graph = ControlFlowGraph(node.instructions)
+
+            fun addEdge(from: Int, to: Int) {
+                graph.edges[from].add(to)
+            }
+
+            object : MethodAnalyzer<BasicValue>("fake", node, OptimizationBasicInterpreter()) {
+                override fun visitControlFlowEdge(insn: Int, successor: Int): Boolean {
+                    addEdge(insn, successor)
+                    return true
+                }
+
+                override fun visitControlFlowExceptionEdge(insn: Int, successor: Int): Boolean {
+                    addEdge(insn, successor)
+                    return true
+                }
+            }.analyze()
+
+            return graph
+        }
+    }
+}

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/OptimizationBasicInterpreter.java

@@ -21,7 +21,6 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.org.objectweb.asm.Opcodes;
 import org.jetbrains.org.objectweb.asm.Type;
 import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode;
-import org.jetbrains.org.objectweb.asm.tree.LdcInsnNode;
 import org.jetbrains.org.objectweb.asm.tree.analysis.AnalyzerException;
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicInterpreter;
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue;
@@ -57,6 +56,14 @@ public class OptimizationBasicInterpreter extends BasicInterpreter {
         }
     }
 
+    @Override
+    public BasicValue newOperation(@NotNull AbstractInsnNode insn) throws AnalyzerException {
+        if (insn.getOpcode() == Opcodes.ACONST_NULL) {
+            return newValue(Type.getObjectType("java/lang/Object"));
+        }
+        return super.newOperation(insn);
+    }
+
     @NotNull
     @Override
     public BasicValue merge(

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/Util.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.codegen.optimization.common
 
+import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Opcodes.*
 import org.jetbrains.org.objectweb.asm.tree.*
 import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue
@@ -125,3 +126,6 @@ val AbstractInsnNode.intConstant: Int? get() =
     }
 
 fun insnListOf(vararg insns: AbstractInsnNode) = InsnList().apply { insns.forEach { add(it) } }
+
+fun AbstractInsnNode.isStoreOperation(): Boolean = getOpcode() in Opcodes.ISTORE..Opcodes.ASTORE
+fun AbstractInsnNode.isLoadOperation(): Boolean = getOpcode() in Opcodes.ILOAD..Opcodes.ALOAD

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/variableLiveness.kt

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.codegen.optimization.common
+
+import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
+import org.jetbrains.org.objectweb.asm.tree.IincInsnNode
+import org.jetbrains.org.objectweb.asm.tree.MethodNode
+import org.jetbrains.org.objectweb.asm.tree.VarInsnNode
+import java.util.*
+
+
+class VariableLivenessFrame(val maxLocals: Int) {
+    private val bitSet = BitSet(maxLocals)
+
+    fun mergeFrom(other: VariableLivenessFrame) {
+        bitSet.or(other.bitSet)
+    }
+
+    fun markAlive(varIndex: Int) {
+        bitSet.set(varIndex, true)
+    }
+
+    fun markDead(varIndex: Int) {
+        bitSet.set(varIndex, false)
+    }
+
+    fun isAlive(varIndex: Int): Boolean = bitSet.get(varIndex)
+
+    override fun equals(other: Any?): Boolean {
+        if (other !is VariableLivenessFrame) return false
+        return bitSet == other.bitSet
+    }
+
+    override fun hashCode() = bitSet.hashCode()
+}
+
+fun analyzeLiveness(node: MethodNode): Array<VariableLivenessFrame> {
+    val graph = ControlFlowGraph.build(node)
+    val insnList = node.instructions
+    val frames = Array(insnList.size()) { VariableLivenessFrame(node.maxLocals) }
+    val insnArray = insnList.toArray()
+
+    val varVisibility = Array(node.maxLocals) { BitSet(insnArray.size) }
+    for (localVar in node.localVariables) {
+        varVisibility[localVar.index].set(insnList.indexOf(localVar.start), insnList.indexOf(localVar.end), true)
+    }
+
+    // see Figure 9.16 from Dragon book
+    var wereChanges: Boolean
+
+    do {
+        wereChanges = false
+        for (insn in insnArray) {
+            val index = insnList.indexOf(insn)
+            val newFrame = VariableLivenessFrame(node.maxLocals)
+            for (successorIndex in graph.getSuccessorsIndices(insn)) {
+                newFrame.mergeFrom(frames[successorIndex])
+            }
+
+            def(newFrame, insn)
+            use(newFrame, insn, index, varVisibility)
+
+            if (frames[index] != newFrame) {
+                frames[index] = newFrame
+                wereChanges = true
+            }
+        }
+
+    } while (wereChanges)
+
+    return frames
+}
+
+private fun def(frame: VariableLivenessFrame, insn: AbstractInsnNode) {
+    if (insn is VarInsnNode && insn.isStoreOperation()) {
+        frame.markDead(insn.`var`)
+    }
+}
+
+private fun use(frame: VariableLivenessFrame, insn: AbstractInsnNode, index: Int, varVisibility: Array<BitSet>) {
+    for (i in 0..frame.maxLocals - 1) {
+        if (varVisibility[i].get(index)) {
+            frame.markAlive(i)
+        }
+    }
+
+    if (insn is VarInsnNode && insn.isLoadOperation()) {
+        frame.markAlive(insn.`var`)
+    }
+    else if (insn is IincInsnNode) {
+        frame.markAlive(insn.`var`)
+    }
+}

compiler/frontend/src/org/jetbrains/kotlin/resolve/SuspendModifierChecker.kt

@@ -28,7 +28,6 @@ import org.jetbrains.kotlin.incremental.KotlinLookupLocation
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
-import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.sure
@@ -48,11 +47,6 @@ object SuspendModifierChecker : SimpleDeclarationChecker {
             diagnosticHolder.report(Errors.INAPPLICABLE_MODIFIER.on(suspendModifierElement, KtTokens.SUSPEND_KEYWORD, message))
         }
 
-        if (functionDescriptor.isInline) {
-            report("inline suspend functions are not supported")
-            return
-        }
-
         val isValidContinuation = functionDescriptor.valueParameters.lastOrNull()?.type?.isValidContinuation() ?: false
         if (!isValidContinuation) {
             report("last parameter of suspend function should have a type of Continuation<T>")

compiler/resolution/src/org/jetbrains/kotlin/resolve/coroutine/coroutineUtil.kt

@@ -23,12 +23,13 @@ import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver
 import org.jetbrains.kotlin.types.KotlinType
 
 val SUSPENSION_POINT_KEY: FunctionDescriptor.UserDataKey<Boolean> = object : FunctionDescriptor.UserDataKey<Boolean> {}
+val REPLACED_SUSPENSION_POINT_KEY: FunctionDescriptor.UserDataKey<Boolean> = object : FunctionDescriptor.UserDataKey<Boolean> {}
 
 // Returns suspension function as it's visible within coroutines:
 // E.g. `fun <V> await(f: CompletableFuture<V>): V` instead of `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`
 fun SimpleFunctionDescriptor.createCoroutineSuspensionFunctionView(): SimpleFunctionDescriptor? {
     if (!isSuspend) return null
-    val returnType = valueParameters.lastOrNull()?.returnType?.arguments?.getOrNull(0)?.type ?: return null
+    val returnType = getSuspensionPointReturnType() ?: return null
 
     val newOriginal =
             if (original !== this)
@@ -45,4 +46,6 @@ fun SimpleFunctionDescriptor.createCoroutineSuspensionFunctionView(): SimpleFunc
     }.build()!!
 }
 
+fun SimpleFunctionDescriptor.getSuspensionPointReturnType(): KotlinType? = valueParameters.lastOrNull()?.returnType?.arguments?.getOrNull(0)?.type
+
 class CoroutineReceiverValue(callableDescriptor: CallableDescriptor, receiverType: KotlinType) : ExtensionReceiver(callableDescriptor, receiverType)

compiler/testData/codegen/box/coroutines/inlineSuspendFunction.kt

@@ -0,0 +1,40 @@
+// WITH_RUNTIME
+class Controller {
+    fun withValue(v: String, x: Continuation<String>) {
+        x.resume(v)
+    }
+
+    suspend inline fun suspendInline(v: String, x: Continuation<String>) {
+        withValue(v, x)
+    }
+
+    suspend inline fun suspendInline(crossinline b: () -> String, x: Continuation<String>) {
+        suspendInline(b(), x)
+    }
+
+    suspend inline fun <reified T : Any> suspendInline(x: Continuation<String>) {
+        suspendInline({ T::class.java.simpleName }, x)
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    c(Controller()).resume(Unit)
+}
+
+class OK
+
+fun box(): String {
+    var result = ""
+
+    builder {
+        result = suspendInline("56")
+        if (result != "56") throw RuntimeException("fail 1")
+
+        result = suspendInline { "57" }
+        if (result != "57") throw RuntimeException("fail 2")
+
+        result = suspendInline<OK>()
+    }
+
+    return result
+}

compiler/testData/diagnostics/tests/coroutines/suspendApplicability.kt

@@ -8,7 +8,7 @@ class Controller {
 
     }
 
-    inline <!INAPPLICABLE_MODIFIER!>suspend<!> fun inlineFun(x: Continuation<Int>) {
+    inline suspend fun inlineFun(x: Continuation<Int>) {
 
     }
 

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -4159,6 +4159,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("inlineSuspendFunction.kt")
+        public void testInlineSuspendFunction() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/inlineSuspendFunction.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("inlinedTryCatchFinally.kt")
         public void testInlinedTryCatchFinally() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/inlinedTryCatchFinally.kt");

idea/testData/checker/infos/suspendApplicability.kt

@@ -8,7 +8,7 @@ class Controller {
 
     }
 
-    <warning descr="[NOTHING_TO_INLINE] Expected performance impact of inlining 'public final inline suspend fun inlineFun(x: Continuation<Int>): Unit defined in Controller' can be insignificant. Inlining works best for functions with lambda parameters"><info descr="null">inline</info></warning> <error descr="[INAPPLICABLE_MODIFIER] 'suspend' modifier is inapplicable. The reason is that inline suspend functions are not supported"><info descr="null">suspend</info></error> fun inlineFun(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
+    <warning descr="[NOTHING_TO_INLINE] Expected performance impact of inlining 'public final inline suspend fun inlineFun(x: Continuation<Int>): Unit defined in Controller' can be insignificant. Inlining works best for functions with lambda parameters"><info descr="null">inline</info></warning> <info descr="null">suspend</info> fun inlineFun(<warning descr="[UNUSED_PARAMETER] Parameter 'x' is never used">x</warning>: Continuation<Int>) {
 
     }