Quick-fix to add 'suspend' to the receiver of startCoroutine / createCoroutine #KT-15738 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/AddSuspendModifierFix.kt

@@ -17,16 +17,23 @@
 package org.jetbrains.kotlin.idea.quickfix
 
 import com.intellij.codeInsight.intention.IntentionAction
+import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.cfg.pseudocode.containingDeclarationForPseudocode
+import org.jetbrains.kotlin.descriptors.ValueDescriptor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
-class AddSuspendModifierFix(element: KtModifierListOwner): AddModifierFix(element, KtTokens.SUSPEND_KEYWORD) {
+class AddSuspendModifierFix(element: KtModifierListOwner, private val name: String?): AddModifierFix(element, KtTokens.SUSPEND_KEYWORD) {
 
     override fun getText() =
             when (element) {
-                is KtNamedFunction -> "Make ${element?.name ?: "containing function"} suspend"
+                is KtNamedFunction -> "Make ${name ?: "containing function"} suspend"
+                is KtTypeReference -> "Make ${name ?: "receiver"} type suspend"
                 else -> super.getText()
             }
 
@@ -35,7 +42,32 @@ class AddSuspendModifierFix(element: KtModifierListOwner): AddModifierFix(elemen
             val element = diagnostic.psiElement
             val function = (element as? KtElement)?.containingDeclarationForPseudocode as? KtNamedFunction ?: return null
 
-            return AddSuspendModifierFix(function)
+            return AddSuspendModifierFix(function, function.name)
+        }
+    }
+
+    object UnresolvedReferenceFactory : KotlinSingleIntentionActionFactory() {
+
+        private val suspendExtensionNames = setOf("startCoroutine", "createCoroutine")
+
+        override fun createAction(diagnostic: Diagnostic): IntentionAction? {
+            val refExpr = diagnostic.psiElement as? KtNameReferenceExpression ?: return null
+            if (refExpr.getReferencedName() !in suspendExtensionNames) return null
+
+            val callParent = refExpr.parent as? KtCallExpression ?: return null
+            val qualifiedGrandParent = callParent.parent as? KtQualifiedExpression ?: return null
+            if (callParent !== qualifiedGrandParent.selectorExpression || refExpr !== callParent.calleeExpression) return null
+            val receiver = qualifiedGrandParent.receiverExpression as? KtNameReferenceExpression ?: return null
+
+            val context = receiver.analyze(BodyResolveMode.PARTIAL)
+            val receiverDescriptor = context[BindingContext.REFERENCE_TARGET, receiver] as? ValueDescriptor ?: return null
+            if (!receiverDescriptor.type.isFunctionType) return null
+            val declaration = DescriptorToSourceUtils.descriptorToDeclaration(receiverDescriptor) as? KtCallableDeclaration
+                              ?: return null
+            if (declaration is KtFunction) return null
+            val variableTypeReference = declaration.typeReference ?: return null
+
+            return AddSuspendModifierFix(variableTypeReference, declaration.name)
         }
     }
 }

idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixRegistrar.kt

@@ -469,6 +469,8 @@ class QuickFixRegistrar : QuickFixContributor {
         ErrorsJs.WRONG_EXTERNAL_DECLARATION.registerFactory(MigrateExternalExtensionFix)
 
         ILLEGAL_SUSPEND_FUNCTION_CALL.registerFactory(AddSuspendModifierFix)
+        UNRESOLVED_REFERENCE.registerFactory(AddSuspendModifierFix.UnresolvedReferenceFactory)
+        UNRESOLVED_REFERENCE_WRONG_RECEIVER.registerFactory(AddSuspendModifierFix.UnresolvedReferenceFactory)
 
         UNSUPPORTED_FEATURE.registerFactory(EnableUnsupportedFeatureFix)
 

idea/testData/quickfix/modifiers/suspend/createCoroutine.kt

@@ -0,0 +1,10 @@
+// "Make block type suspend" "true"
+// WITH_RUNTIME
+
+import kotlin.coroutines.suspendCoroutine
+import kotlin.coroutines.createCoroutine
+
+suspend fun <T> suspending(): T {
+    val block: () -> T = { null!! }
+    return suspendCoroutine { block.<caret>createCoroutine(it) }
+}

idea/testData/quickfix/modifiers/suspend/createCoroutine.kt.after

@@ -0,0 +1,10 @@
+// "Make block type suspend" "true"
+// WITH_RUNTIME
+
+import kotlin.coroutines.suspendCoroutine
+import kotlin.coroutines.createCoroutine
+
+suspend fun <T> suspending(): T {
+    val block: suspend () -> T = { null!! }
+    return suspendCoroutine { block.createCoroutine(it) }
+}

idea/testData/quickfix/modifiers/suspend/startCoroutine.kt

@@ -0,0 +1,7 @@
+// "Make block type suspend" "true"
+// WITH_RUNTIME
+
+import kotlin.coroutines.suspendCoroutine
+import kotlin.coroutines.startCoroutine
+
+suspend fun <T> suspending(block: () -> T): T = suspendCoroutine { block.<caret>startCoroutine(it) }

idea/testData/quickfix/modifiers/suspend/startCoroutine.kt.after

@@ -0,0 +1,7 @@
+// "Make block type suspend" "true"
+// WITH_RUNTIME
+
+import kotlin.coroutines.suspendCoroutine
+import kotlin.coroutines.startCoroutine
+
+suspend fun <T> suspending(block: suspend () -> T): T = suspendCoroutine { block.startCoroutine(it) }

idea/testData/quickfix/modifiers/suspend/startCoroutineNoImport.kt

@@ -0,0 +1,7 @@
+// "Make block type suspend" "true"
+// WITH_RUNTIME
+// DISABLE-ERRORS
+
+import kotlin.coroutines.suspendCoroutine
+
+suspend fun <T> suspending(block: () -> T): T = suspendCoroutine { block.<caret>startCoroutine(it) }

idea/testData/quickfix/modifiers/suspend/startCoroutineNoImport.kt.after

@@ -0,0 +1,7 @@
+// "Make block type suspend" "true"
+// WITH_RUNTIME
+// DISABLE-ERRORS
+
+import kotlin.coroutines.suspendCoroutine
+
+suspend fun <T> suspending(block: suspend () -> T): T = suspendCoroutine { block.startCoroutine(it) }

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -6746,12 +6746,30 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("createCoroutine.kt")
+            public void testCreateCoroutine() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/modifiers/suspend/createCoroutine.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("init.kt")
             public void testInit() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/modifiers/suspend/init.kt");
                 doTest(fileName);
             }
 
+            @TestMetadata("startCoroutine.kt")
+            public void testStartCoroutine() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/modifiers/suspend/startCoroutine.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("startCoroutineNoImport.kt")
+            public void testStartCoroutineNoImport() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/modifiers/suspend/startCoroutineNoImport.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("topLevel.kt")
             public void testTopLevel() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/modifiers/suspend/topLevel.kt");