Extract Function: Fix detection of suspend calls containing extracted parameters

#KT-16251 Fixed

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractableAnalysisUtil.kt

@@ -655,9 +655,6 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
         return AnalysisResult(null, Status.CRITICAL_ERROR, listOf(paramsInfo.errorMessage!!))
     }
 
-    val virtualContext = virtualBlock.analyze(BodyResolveMode.PARTIAL)
-    val isSuspendExpected = virtualContext.diagnostics.all().any { it.factory == Errors.ILLEGAL_SUSPEND_FUNCTION_CALL }
-
     val messages = ArrayList<ErrorMessage>()
 
     val modifiedVarDescriptorsForControlFlow = HashMap(modifiedVarDescriptorsWithExpressions)
@@ -703,20 +700,32 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
     val receiverParameter = if (receiverCandidates.size == 1 && !options.canWrapInWith) receiverCandidates.first() else null
     receiverParameter?.let { adjustedParameters.remove(it) }
 
+    var descriptor = ExtractableCodeDescriptor(
+            this,
+            bindingContext,
+            suggestFunctionNames(returnType),
+            getDefaultVisibility(),
+            adjustedParameters.sortedBy { it.name },
+            receiverParameter,
+            paramsInfo.typeParameters.sortedBy { it.originalDeclaration.name!! },
+            paramsInfo.replacementMap,
+            if (messages.isEmpty()) controlFlow else controlFlow.toDefault(),
+            returnType,
+            emptyList()
+    )
+
+    val body = ExtractionGeneratorConfiguration(
+            descriptor,
+            ExtractionGeneratorOptions(inTempFile = true, allowExpressionBody = false)
+    ).generateDeclaration().declaration.getGeneratedBody()
+    val virtualContext = body.analyzeFully()
+    if (virtualContext.diagnostics.all().any { it.factory == Errors.ILLEGAL_SUSPEND_FUNCTION_CALL }) {
+        descriptor = descriptor.copy(modifiers = listOf(KtTokens.SUSPEND_KEYWORD))
+    }
+
+
     return AnalysisResult(
-            ExtractableCodeDescriptor(
-                    this,
-                    bindingContext,
-                    suggestFunctionNames(returnType),
-                    getDefaultVisibility(),
-                    adjustedParameters.sortedBy { it.name },
-                    receiverParameter,
-                    paramsInfo.typeParameters.sortedBy { it.originalDeclaration.name!! },
-                    paramsInfo.replacementMap,
-                    if (messages.isEmpty()) controlFlow else controlFlow.toDefault(),
-                    returnType,
-                    if (isSuspendExpected) listOf(KtTokens.SUSPEND_KEYWORD) else emptyList()
-            ),
+            descriptor,
             if (messages.isEmpty()) Status.SUCCESS else Status.NON_CRITICAL_ERROR,
             messages
     )

idea/testData/refactoring/extractFunction/basic/suspendCallWithExtractedParameter.kt

@@ -0,0 +1,10 @@
+// PARAM_TYPES: D
+// PARAM_DESCRIPTOR: value-parameter d: D defined in test1
+class D {
+    suspend fun await() {}
+}
+
+// SIBLING:
+suspend fun test1(d: D) {
+    <selection>d.await()</selection>
+}

idea/testData/refactoring/extractFunction/basic/suspendCallWithExtractedParameter.kt.after

@@ -0,0 +1,14 @@
+// PARAM_TYPES: D
+// PARAM_DESCRIPTOR: value-parameter d: D defined in test1
+class D {
+    suspend fun await() {}
+}
+
+// SIBLING:
+suspend fun test1(d: D) {
+    __dummyTestFun__(d)
+}
+
+private suspend fun __dummyTestFun__(d: D) {
+    d.await()
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/ExtractionTestGenerated.java

@@ -1102,6 +1102,12 @@ public class ExtractionTestGenerated extends AbstractExtractionTest {
                 doExtractFunctionTest(fileName);
             }
 
+            @TestMetadata("suspendCallWithExtractedParameter.kt")
+            public void testSuspendCallWithExtractedParameter() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/suspendCallWithExtractedParameter.kt");
+                doExtractFunctionTest(fileName);
+            }
+
             @TestMetadata("topLevelValUnderSmartCast.kt")
             public void testTopLevelValUnderSmartCast() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/extractFunction/basic/topLevelValUnderSmartCast.kt");