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 argumentsText = listOfNotNull(
-            secondCallArgumentList.takeIf { it?.arguments?.isNotEmpty() == true },
-            firstCallArgumentList.takeIf { it?.arguments?.isNotEmpty() == true }
-        ).let {
         val additionalArgument = conversion.additionalArgument
-            when {
+        val secondCallHasArguments = secondCallArgumentList?.arguments?.isNotEmpty() == true
-                it.isNotEmpty() -> it.joinToString(
+        val firstCallHasArguments = firstCallArgumentList?.arguments?.isNotEmpty() == true
-                    separator = ", ",
+        val argumentsText = listOfNotNull(
-                    prefix = "(",
+            secondCallArgumentList.takeIf { secondCallHasArguments }?.getTextInsideParentheses(),
-                    postfix = ")"
+            firstCallArgumentList.takeIf { firstCallHasArguments }?.getTextInsideParentheses(),
-                ) { callArgumentList ->
+            additionalArgument.takeIf { !firstCallHasArguments && !secondCallHasArguments },
-                    callArgumentList.getTextInsideParentheses()
+            lambdaExpression?.text
-                }
+        ).joinToString(separator = ",")
-                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()
         }