Redundant async inspection: support case with explicit scope

#KT-28504 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/coroutines/RedundantAsyncInspection.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.inspections.coroutines
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
 import org.jetbrains.kotlin.idea.inspections.collections.AbstractCallChainChecker
 import org.jetbrains.kotlin.idea.inspections.collections.SimplifyCallChainFix
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
@@ -15,6 +16,7 @@ import org.jetbrains.kotlin.psi.KtQualifiedExpression
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.psi.qualifiedExpressionVisitor
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.model.ArgumentMatch
 import org.jetbrains.kotlin.resolve.calls.model.DefaultValueArgument
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
@@ -29,8 +31,8 @@ class RedundantAsyncInspection : AbstractCallChainChecker() {
             for ((parameterDescriptor, valueArgument) in firstResolvedCall.valueArguments) {
                 val default = valueArgument is DefaultValueArgument
                 when (parameterDescriptor.name.asString()) {
-                    "context" -> defaultContext = default
-                    "start" -> defaultStart = default
+                    CONTEXT_ARGUMENT_NAME -> defaultContext = default
+                    START_ARGUMENT_NAME -> defaultStart = default
                 }
             }
             true
@@ -45,12 +47,11 @@ class RedundantAsyncInspection : AbstractCallChainChecker() {
             val context = scopeExpression.analyze(BodyResolveMode.PARTIAL)
             val scopeDescriptor = (scopeExpression as? KtNameReferenceExpression)?.let { context[BindingContext.REFERENCE_TARGET, it] }
             if (scopeDescriptor?.fqNameSafe?.toString() != GLOBAL_SCOPE) {
-                // Temporary forbid cases with explicit non-global scope
-                return null
+                conversion = conversion.withArgument("${scopeExpression.text}.coroutineContext")
             }
         }
 
-        if (defaultContext!! && defaultStart!!) {
+        if (conversion.additionalArgument == null && defaultContext!! && defaultStart!!) {
             conversion = conversion.withArgument(
                 if (conversion === conversions[0]) {
                     DEFAULT_ASYNC_ARGUMENT
@@ -65,13 +66,35 @@ class RedundantAsyncInspection : AbstractCallChainChecker() {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean) =
         qualifiedExpressionVisitor(fun(expression) {
             val conversion = generateConversion(expression) ?: return
+            val contextArgument = conversion.additionalArgument
             val descriptor = holder.manager.createProblemDescriptor(
                 expression,
                 expression.firstCalleeExpression()!!.textRange.shiftRight(-expression.startOffset),
                 "Redundant 'async' call may be reduced to '${conversion.replacement}'",
                 ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
                 isOnTheFly,
-                SimplifyCallChainFix(conversion, removeReceiverOfFirstCall = true)
+                SimplifyCallChainFix(conversion, removeReceiverOfFirstCall = true) { callExpression ->
+                    if (contextArgument != null) {
+                        val call = callExpression.resolveToCall()
+                        if (call != null) {
+                            for (argument in callExpression.valueArguments) {
+                                val mapping = call.getArgumentMapping(argument) as? ArgumentMatch ?: continue
+                                if (mapping.valueParameter.name.asString() == CONTEXT_ARGUMENT_NAME) {
+                                    val name = argument.getArgumentName()?.asName
+                                    val expressionText = contextArgument + " + " + argument.getArgumentExpression()!!.text
+                                    argument.replace(
+                                        if (name == null) {
+                                            createArgument(expressionText)
+                                        } else {
+                                            createArgument("$name = $expressionText")
+                                        }
+                                    )
+                                    break
+                                }
+                            }
+                        }
+                    }
+                }
             )
             holder.registerProblem(descriptor)
         })
@@ -94,6 +117,10 @@ class RedundantAsyncInspection : AbstractCallChainChecker() {
 
         private const val GLOBAL_SCOPE = "kotlinx.coroutines.GlobalScope"
 
+        private const val CONTEXT_ARGUMENT_NAME = "context"
+
+        private const val START_ARGUMENT_NAME = "start"
+
         private const val DEFAULT_ASYNC_ARGUMENT = "$COROUTINE_PACKAGE.Dispatchers.Default"
 
         private const val DEFAULT_ASYNC_ARGUMENT_EXPERIMENTAL = "$COROUTINE_EXPERIMENTAL_PACKAGE.DefaultDispatcher"

idea/testData/inspectionsLocal/coroutines/redundantAsync/coroutines.lib.kt

@@ -41,4 +41,7 @@ suspend fun <T> withContext(
 
 suspend fun <R> coroutineScope(block: suspend CoroutineScope.() -> R): R = GlobalScope.block()
 
+operator fun CoroutineContext.plus(other: CoroutineContext): CoroutineContext {
+    TODO()
+}
 

idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScope.kt

@@ -1,9 +1,7 @@
 // WITH_RUNTIME
-// PROBLEM: none
 
 package kotlinx.coroutines
 
 suspend fun test(ctx: CoroutineContext, scope: CoroutineScope) {
-    // Does not work yet (1.3.11)
     scope.<caret>async(ctx) { 42 }.await()
 }

idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScope.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+package kotlinx.coroutines
+
+suspend fun test(ctx: CoroutineContext, scope: CoroutineScope) {
+    withContext(scope.coroutineContext + ctx) { 42 }
+}

idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScopeNamed.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+package kotlinx.coroutines
+
+suspend fun test(ctx: CoroutineContext, scope: CoroutineScope) {
+    scope.<caret>async(context = ctx) { 42 }.await()
+}

idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScopeNamed.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+package kotlinx.coroutines
+
+suspend fun test(ctx: CoroutineContext, scope: CoroutineScope) {
+    withContext(context = scope.coroutineContext + ctx) { 42 }
+}

idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScopeNoContext.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+package kotlinx.coroutines
+
+suspend fun test(scope: CoroutineScope) {
+    scope.<caret>async() { 42 }.await()
+}

idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScopeNoContext.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+package kotlinx.coroutines
+
+suspend fun test(scope: CoroutineScope) {
+    withContext(scope.coroutineContext) { 42 }
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -2075,6 +2075,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScope.kt");
             }
 
+            @TestMetadata("explicitScopeNamed.kt")
+            public void testExplicitScopeNamed() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScopeNamed.kt");
+            }
+
+            @TestMetadata("explicitScopeNoContext.kt")
+            public void testExplicitScopeNoContext() throws Exception {
+                runTest("idea/testData/inspectionsLocal/coroutines/redundantAsync/explicitScopeNoContext.kt");
+            }
+
             @TestMetadata("globalScope.kt")
             public void testGlobalScope() throws Exception {
                 runTest("idea/testData/inspectionsLocal/coroutines/redundantAsync/globalScope.kt");