Refactor code in "simplify call chain" (relates to KT-28576)

idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifyCallChainFix.kt

@@ -20,8 +20,11 @@ import com.intellij.codeInsight.actions.OptimizeImportsProcessor
 import com.intellij.codeInspection.LocalQuickFix
 import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.codeStyle.CodeStyleManager
 import org.jetbrains.kotlin.idea.core.ShortenReferences
+import org.jetbrains.kotlin.idea.core.moveFunctionLiteralOutsideParentheses
 import org.jetbrains.kotlin.idea.core.replaced
+import org.jetbrains.kotlin.idea.intentions.callExpression
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
 
@@ -47,8 +50,8 @@ class SimplifyCallChainFix(
             else -> ""
         }
 
-        val receiverExpressionOrEmptyString: Any =
+        val receiverExpressionOrEmptyString =
-            if (!removeReceiverOfFirstCall && firstExpression is KtQualifiedExpression) firstExpression.receiverExpression else ""
+            if (!removeReceiverOfFirstCall && firstExpression is KtQualifiedExpression) firstExpression.receiverExpression.text else ""
 
         val firstCallExpression = AbstractCallChainChecker.getCallExpression(firstExpression) ?: return
         factory.modifyArguments(firstCallExpression)
@@ -63,45 +66,35 @@ class SimplifyCallChainFix(
         }
 
         val lambdaExpression = firstCallExpression.lambdaArguments.singleOrNull()?.getLambdaExpression()
+        val additionalArgument = conversion.additionalArgument
+        val secondCallHasArguments = secondCallArgumentList?.arguments?.isNotEmpty() == true
+        val firstCallHasArguments = firstCallArgumentList?.arguments?.isNotEmpty() == true
         val argumentsText = listOfNotNull(
-            secondCallArgumentList.takeIf { it?.arguments?.isNotEmpty() == true },
+            secondCallArgumentList.takeIf { secondCallHasArguments }?.getTextInsideParentheses(),
-            firstCallArgumentList.takeIf { it?.arguments?.isNotEmpty() == true }
+            firstCallArgumentList.takeIf { firstCallHasArguments }?.getTextInsideParentheses(),
-        ).let {
+            additionalArgument.takeIf { !firstCallHasArguments && !secondCallHasArguments },
-            val additionalArgument = conversion.additionalArgument
+            lambdaExpression?.text
-            when {
+        ).joinToString(separator = ",")
-                it.isNotEmpty() -> it.joinToString(
-                    separator = ", ",
-                    prefix = "(",
-                    postfix = ")"
-                ) { callArgumentList ->
-                    callArgumentList.getTextInsideParentheses()
-                }
-                additionalArgument != null -> "($additionalArgument)"
-                else -> ""
-            }
-        }
 
         val newCallText = conversion.replacement
-        val newQualifiedExpression = if (lambdaExpression != null) factory.createExpressionByPattern(
+        val newQualifiedOrCallExpression = factory.createExpression(
-            "$0$1$2 $3 $4",
+            "$receiverExpressionOrEmptyString$operationSign$newCallText($argumentsText)"
-            receiverExpressionOrEmptyString,
-            operationSign,
-            newCallText,
-            argumentsText,
-            lambdaExpression.text
-        )
-        else factory.createExpressionByPattern(
-            "$0$1$2 $3",
-            receiverExpressionOrEmptyString,
-            operationSign,
-            newCallText,
-            argumentsText
         )
 
+        if (lambdaExpression != null) {
+            val callExpression = when (newQualifiedOrCallExpression) {
+                is KtQualifiedExpression -> newQualifiedOrCallExpression.callExpression
+                is KtCallExpression -> newQualifiedOrCallExpression
+                else -> null
+            }
+            callExpression?.moveFunctionLiteralOutsideParentheses()
+        }
+
         val project = qualifiedExpression.project
         val file = qualifiedExpression.containingKtFile
-        val result = qualifiedExpression.replaced(newQualifiedExpression)
+        val result = qualifiedExpression.replaced(newQualifiedOrCallExpression)
-        ShortenReferences.DEFAULT.process(result)
+        val reformatted = CodeStyleManager.getInstance(project).reformat(result)
+        ShortenReferences.DEFAULT.process(reformatted as KtElement)
         if (runOptimizeImports) {
             OptimizeImportsProcessor(project, file).run()
         }

idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToStringViaBuilder.kt.after

@@ -1,7 +1,7 @@
 // WITH_RUNTIME
 
 val x = listOf(1, 2, 3).joinToString(prefix = "= ", separator = " + ") {
-    val sb = StringBuilder()
+val sb = StringBuilder()
-    sb.append(it).append(" + ").append(it)
+sb.append(it).append(" + ").append(it)
-    sb
+sb
 }

idea/testData/inspectionsLocal/coroutines/deferredIsResult/complex.kt.after

@@ -8,9 +8,9 @@ suspend fun myFunction(context: CoroutineContext, switch: Int): Int {
     with (GlobalScope) {
         when (switch) {
             0 -> return withContext(Dispatchers.Default) {
-            val x = 123
+                val x = 123
-            x * x
+                x * x
-        }
+            }
             1 -> return withContext(context) { -1 }
             else -> return withContext(Dispatchers.Default) { 9 }
         }