From b638e91c2bb6de8387cb8cbd78d3cd48f533c398 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Tue, 24 Jul 2018 14:30:03 +0300
Subject: [PATCH] Call chain --> Sequence: minor improvements

Related to KT-15476
---
 .../ConvertCallChainIntoSequenceInspection.kt | 50 ++++++++++++-------
 1 file changed, 32 insertions(+), 18 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/collections/ConvertCallChainIntoSequenceInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/collections/ConvertCallChainIntoSequenceInspection.kt
index 13b4a6a39e4..8e37f039a5f 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/collections/ConvertCallChainIntoSequenceInspection.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/collections/ConvertCallChainIntoSequenceInspection.kt
@@ -11,18 +11,24 @@ import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiWhiteSpace
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.DefaultBuiltIns
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.inspections.AbstractKotlinInspection
 import org.jetbrains.kotlin.idea.intentions.callExpression
-import org.jetbrains.kotlin.idea.intentions.getCallableDescriptor
 import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.KotlinType
 
 class ConvertCallChainIntoSequenceInspection : AbstractKotlinInspection() {
@@ -50,12 +56,13 @@ private class ConvertCallChainIntoSequenceFix : LocalQuickFix {
 
     override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
         val expression = descriptor.psiElement as? KtQualifiedExpression ?: return
-        val calls = expression.collectCallExpression().reversed()
+        val context = expression.analyze(BodyResolveMode.PARTIAL)
+        val calls = expression.collectCallExpression(context).reversed()
         val firstCall = calls.firstOrNull() ?: return
         val lastCall = calls.lastOrNull() ?: return
-        val first = firstCall.parent as? KtQualifiedExpression ?: return
-        val last = lastCall.parent as? KtQualifiedExpression ?: return
-        val endWithTermination = lastCall.isTermination()
+        val first = firstCall.getQualifiedExpressionForSelector() ?: return
+        val last = lastCall.getQualifiedExpressionForSelector() ?: return
+        val endWithTermination = lastCall.isTermination(context)
 
         val psiFactory = KtPsiFactory(expression)
         val dot = buildString {
@@ -93,17 +100,19 @@ private class ConvertCallChainIntoSequenceFix : LocalQuickFix {
 private fun KtQualifiedExpression.findTarget(): Pair<KtQualifiedExpression, KtCallExpression>? {
     if (parent is KtQualifiedExpression) return null
 
-    val calls = collectCallExpression()
+    val context = analyze(BodyResolveMode.PARTIAL)
+    val calls = collectCallExpression(context)
     if (calls.isEmpty()) return null
 
-    val receiverType = (calls.last().parent as? KtQualifiedExpression)?.receiverExpression?.getCallableDescriptor()?.returnType
+    val receiverType =
+        (calls.last().getQualifiedExpressionForSelector())?.receiverExpression?.getResolvedCall(context)?.resultingDescriptor?.returnType
     if (receiverType?.isCollection() != true) return null
 
-    val qualified = calls.first().parent as? KtQualifiedExpression ?: return null
+    val qualified = calls.first().getQualifiedExpressionForSelector() ?: return null
     return qualified to calls.last()
 }
 
-private fun KtQualifiedExpression.collectCallExpression(): List<KtCallExpression> {
+private fun KtQualifiedExpression.collectCallExpression(context: BindingContext): List<KtCallExpression> {
     val calls = mutableListOf<KtCallExpression>()
 
     fun collect(qualified: KtQualifiedExpression) {
@@ -118,29 +127,34 @@ private fun KtQualifiedExpression.collectCallExpression(): List<KtCallExpression
 
     val transformationCalls = calls
         .asSequence()
-        .dropWhile { !it.isTransformationOrTermination() }
-        .takeWhile { it.isTransformationOrTermination() && !it.hasReturn() }
+        .dropWhile { !it.isTransformationOrTermination(context) }
+        .takeWhile { it.isTransformationOrTermination(context) && !it.hasReturn() }
         .toList()
     if (transformationCalls.size < 2) return emptyList()
 
     return transformationCalls
 }
 
-private fun KotlinType.isCollection(): Boolean =
-    constructor.declarationDescriptor?.fqNameSafe == KotlinBuiltIns.FQ_NAMES.collection || constructor.supertypes.any { it.isCollection() }
+private fun KotlinType.isCollection(): Boolean {
+    val classDescriptor = constructor.declarationDescriptor as? ClassDescriptor ?: return false
+    val className = classDescriptor.name.asString()
+    val builtIns = DefaultBuiltIns.Instance
+    return className.endsWith("List") && classDescriptor.isSubclassOf(builtIns.list)
+            || className.endsWith("Set") && classDescriptor.isSubclassOf(builtIns.set)
+}
 
 private fun KtCallExpression.hasReturn(): Boolean = valueArguments.any { arg ->
     arg.anyDescendantOfType<KtReturnExpression> { it.labelQualifier == null }
 }
 
-private fun KtCallExpression.isTransformationOrTermination(): Boolean {
+private fun KtCallExpression.isTransformationOrTermination(context: BindingContext): Boolean {
     val fqName = transformationAndTerminations[calleeExpression?.text] ?: return false
-    return fqName == getCallableDescriptor()?.fqNameSafe
+    return fqName == getResolvedCall(context)?.resultingDescriptor?.fqNameSafe
 }
 
-private fun KtCallExpression.isTermination(): Boolean {
+private fun KtCallExpression.isTermination(context: BindingContext): Boolean {
     val fqName = terminations[calleeExpression?.text] ?: return false
-    return fqName == getCallableDescriptor()?.fqNameSafe
+    return fqName == getResolvedCall(context)?.resultingDescriptor?.fqNameSafe
 }
 
 private val transformations = listOf(