From b6630699aaf3f6521f16952ee7ae64c7f0b1c6d4 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Wed, 28 Nov 2018 16:33:56 +0300
Subject: [PATCH] SimplifyCallChainFix: implement 'modifyArguments' as lambda

---
 .../SimplifiableCallChainInspection.kt        | 17 +++++++++++-----
 .../collections/SimplifyCallChainFix.kt       | 20 ++++++++-----------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifiableCallChainInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifiableCallChainInspection.kt
index b08a59168f3..d89b64bbfcd 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifiableCallChainInspection.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifiableCallChainInspection.kt
@@ -8,13 +8,11 @@ package org.jetbrains.kotlin.idea.inspections.collections
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.js.resolve.JsPlatform
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.psi.qualifiedExpressionVisitor
 import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
-import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.builtIns
 import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 
@@ -40,14 +38,23 @@ class SimplifiableCallChainInspection : AbstractCallChainChecker() {
                 true
             } ?: return
 
-
+            val replacement = conversion.replacement
             val descriptor = holder.manager.createProblemDescriptor(
                 expression,
                 expression.firstCalleeExpression()!!.textRange.shiftRight(-expression.startOffset),
                 "Call chain on collection type may be simplified",
                 ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
                 isOnTheFly,
-                SimplifyCallChainFix(conversion.replacement)
+                SimplifyCallChainFix(replacement) { callExpression ->
+                    val lastArgumentName = if (replacement.startsWith("joinTo")) Name.identifier("transform") else null
+                    if (lastArgumentName != null) {
+                        val lastArgument = callExpression.valueArgumentList?.arguments?.singleOrNull()
+                        val argumentExpression = lastArgument?.getArgumentExpression()
+                        if (argumentExpression != null) {
+                            lastArgument.replace(createArgument(argumentExpression, lastArgumentName))
+                        }
+                    }
+                }
             )
             holder.registerProblem(descriptor)
         })
diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifyCallChainFix.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifyCallChainFix.kt
index bdf25b3c628..488fed96a91 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifyCallChainFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifyCallChainFix.kt
@@ -21,11 +21,14 @@ import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.idea.core.ShortenReferences
 import org.jetbrains.kotlin.idea.core.replaced
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
 
-class SimplifyCallChainFix(private val newCallText: String, private val removeReceiverOfFirstCall: Boolean = false) : LocalQuickFix {
+class SimplifyCallChainFix(
+    private val newCallText: String,
+    private val removeReceiverOfFirstCall: Boolean = false,
+    private val modifyArguments: KtPsiFactory.(KtCallExpression) -> Unit = {}
+) : LocalQuickFix {
     private val shortenedText = newCallText.split("(").joinToString(separator = "(") {
         it.substringAfterLast(".")
     }
@@ -48,17 +51,10 @@ class SimplifyCallChainFix(private val newCallText: String, private val removeRe
             if (!removeReceiverOfFirstCall && firstExpression is KtQualifiedExpression) firstExpression.receiverExpression else ""
 
         val firstCallExpression = AbstractCallChainChecker.getCallExpression(firstExpression) ?: return
-        val secondCallExpression = qualifiedExpression.selectorExpression as? KtCallExpression ?: return
-
-        val lastArgumentName = if (newCallText.startsWith("joinTo")) Name.identifier("transform") else null
-        if (lastArgumentName != null) {
-            val lastArgument = firstCallExpression.valueArgumentList?.arguments?.singleOrNull()
-            val argumentExpression = lastArgument?.getArgumentExpression()
-            if (argumentExpression != null) {
-                lastArgument.replace(factory.createArgument(argumentExpression, lastArgumentName))
-            }
-        }
+        factory.modifyArguments(firstCallExpression)
         val firstCallArgumentList = firstCallExpression.valueArgumentList
+
+        val secondCallExpression = qualifiedExpression.selectorExpression as? KtCallExpression ?: return
         val secondCallArgumentList = secondCallExpression.valueArgumentList
 
         fun KtValueArgumentList.getTextInsideParentheses(): String {