From 1bc72f0b323b7b57efb0da5af772620f4af02b2b Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Sun, 29 Dec 2019 11:02:36 +0900
Subject: [PATCH] "Surround with lambda" quickfix: suggest for suspend function

#KT-34371 Fixed
---
 .../jetbrains/kotlin/idea/quickfix/SurroundWithLambdaFix.kt | 4 ++--
 .../quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt    | 6 ++++++
 .../typeMismatch/paramTypeLambdaMatchSuspend.kt.after       | 6 ++++++
 .../kotlin/idea/quickfix/QuickFixTestGenerated.java         | 5 +++++
 4 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt
 create mode 100644 idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/SurroundWithLambdaFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/SurroundWithLambdaFix.kt
index ef12d922bad..061c7c32c64 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/SurroundWithLambdaFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/SurroundWithLambdaFix.kt
@@ -4,7 +4,7 @@ import com.intellij.codeInsight.intention.HighPriorityAction
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.builtins.isFunctionType
+import org.jetbrains.kotlin.builtins.isFunctionOrSuspendFunctionType
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
@@ -68,7 +68,7 @@ class SurroundWithLambdaFix(
                 }
             }
 
-            if (!expectedType.isFunctionType) return null
+            if (!expectedType.isFunctionOrSuspendFunctionType) return null
             if (expectedType.arguments.size != 1) return null
             val lambdaReturnType = expectedType.arguments[0].type
 
diff --git a/idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt b/idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt
new file mode 100644
index 00000000000..649333ef211
--- /dev/null
+++ b/idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt
@@ -0,0 +1,6 @@
+// "Surround with lambda" "true"
+fun foo(action: suspend () -> String) {}
+
+fun usage() {
+    foo("oraora"<caret>)
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt.after b/idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt.after
new file mode 100644
index 00000000000..3528080b619
--- /dev/null
+++ b/idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt.after
@@ -0,0 +1,6 @@
+// "Surround with lambda" "true"
+fun foo(action: suspend () -> String) {}
+
+fun usage() {
+    foo({ "oraora" })
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
index 1c0c7bc810b..5a11b33bf8d 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
@@ -13584,6 +13584,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSubclass.kt");
         }
 
+        @TestMetadata("paramTypeLambdaMatchSuspend.kt")
+        public void testParamTypeLambdaMatchSuspend() throws Exception {
+            runTest("idea/testData/quickfix/typeMismatch/paramTypeLambdaMatchSuspend.kt");
+        }
+
         @TestMetadata("paramTypeLambdaMismatch.kt")
         public void testParamTypeLambdaMismatch() throws Exception {
             runTest("idea/testData/quickfix/typeMismatch/paramTypeLambdaMismatch.kt");