Fixed some cases when smart cast problems were not detected

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/baseTransformations.kt

@@ -27,6 +27,9 @@ abstract class ReplaceLoopResultTransformation(override val loop: KtForExpressio
 
     override val commentSavingRange = PsiChildRange.singleElement(loop.unwrapIfLabeled())
 
+    override val expressionToBeReplacedByResultCallChain: KtExpression
+        get() = loop.unwrapIfLabeled()
+
     override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
         return loop.unwrapIfLabeled().replaced(resultCallChain)
     }
@@ -39,6 +42,9 @@ abstract class AssignToVariableResultTransformation(
 
     override val commentSavingRange = PsiChildRange(initialization.initializationStatement, loop.unwrapIfLabeled())
 
+    override val expressionToBeReplacedByResultCallChain: KtExpression
+        get() = initialization.initializer
+
     override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
         initialization.initializer.replace(resultCallChain)
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt

@@ -78,6 +78,8 @@ interface ResultTransformation : Transformation {
 
     val commentSavingRange: PsiChildRange
 
+    val expressionToBeReplacedByResultCallChain: KtExpression
+
     /**
      * Implementations of this method are obliged to update [commentSavingRangeHolder] when deleting or adding any element into the tree
      * except for the loop itself and the result element returned from this method

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt

@@ -190,21 +190,21 @@ private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: MatchRe
     val SMARTCAST_KEY = Key<KotlinType>("SMARTCAST_KEY")
     val IMPLICIT_RECEIVER_SMARTCAST_KEY = Key<KotlinType>("IMPLICIT_RECEIVER_SMARTCAST")
 
-    var smartCastsFound = false
+    var smartCastCount = 0
     try {
         loop.forEachDescendantOfType<KtExpression> { expression ->
             bindingContext[BindingContext.SMARTCAST, expression]?.let {
                 expression.putCopyableUserData(SMARTCAST_KEY, it)
-                smartCastsFound = true
+                smartCastCount++
             }
 
             bindingContext[BindingContext.IMPLICIT_RECEIVER_SMARTCAST, expression]?.let {
                 expression.putCopyableUserData(IMPLICIT_RECEIVER_SMARTCAST_KEY, it)
-                smartCastsFound = true
+                smartCastCount++
             }
         }
 
-        if (!smartCastsFound) return true // optimization
+        if (smartCastCount == 0) return true // optimization
 
         val callChain = matchResult.generateCallChain(loop)
 
@@ -214,23 +214,40 @@ private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: MatchRe
 
         val smartCastBroken = callChain.anyDescendantOfType<KtExpression> { expression ->
             val smartCastType = expression.getCopyableUserData(SMARTCAST_KEY)
-            if (smartCastType != null && newBindingContext[BindingContext.SMARTCAST, expression] != smartCastType && newBindingContext.getType(expression) != smartCastType) {
-                return@anyDescendantOfType true
+            if (smartCastType != null) {
+                if (newBindingContext[BindingContext.SMARTCAST, expression] != smartCastType && newBindingContext.getType(expression) != smartCastType) {
+                    return@anyDescendantOfType true
+                }
+                smartCastCount--
             }
 
             val implicitReceiverSmartCastType = expression.getCopyableUserData(IMPLICIT_RECEIVER_SMARTCAST_KEY)
-            if (implicitReceiverSmartCastType != null && newBindingContext[BindingContext.IMPLICIT_RECEIVER_SMARTCAST, expression] != implicitReceiverSmartCastType) {
-                return@anyDescendantOfType true
+            if (implicitReceiverSmartCastType != null) {
+                if (newBindingContext[BindingContext.IMPLICIT_RECEIVER_SMARTCAST, expression] != implicitReceiverSmartCastType) {
+                    return@anyDescendantOfType true
+                }
+                smartCastCount--
             }
 
             false
         }
 
-        return !smartCastBroken
+        if (smartCastBroken) return false
+
+        assert(smartCastCount >= 0)
+        if (smartCastCount > 0) { // not all smart cast expressions has been found in the result, perform more expensive check
+            val expressionToBeReplaced = matchResult.transformationMatch.resultTransformation.expressionToBeReplacedByResultCallChain
+            if (!tryChangeAndCheckErrors(expressionToBeReplaced, loop) { it.replace(callChain) }) return false
+        }
+
+        return true
     }
     finally {
-        if (smartCastsFound) {
-            loop.forEachDescendantOfType<KtExpression> { it.putCopyableUserData(SMARTCAST_KEY, null) }
+        if (smartCastCount > 0) {
+            loop.forEachDescendantOfType<KtExpression> {
+                it.putCopyableUserData(SMARTCAST_KEY, null)
+                it.putCopyableUserData(IMPLICIT_RECEIVER_SMARTCAST_KEY, null)
+            }
         }
     }
 }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindTransformationMatcher.kt

@@ -127,6 +127,9 @@ object FindTransformationMatcher : TransformationMatcher {
             return generator.generate(chainedCallGenerator)
         }
 
+        override val expressionToBeReplacedByResultCallChain: KtExpression
+            get() = endReturn.returnedExpression!!
+
         override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
             endReturn.returnedExpression!!.replace(resultCallChain)
             loop.deleteWithLabels()

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt

@@ -162,35 +162,43 @@ fun KtExpression.isSimpleCollectionInstantiation(): CollectionKind? {
 }
 
 fun canChangeLocalVariableType(variable: KtProperty, newTypeText: String, loop: KtForExpression): Boolean {
-    val bindingContext = variable.analyze(BodyResolveMode.FULL)
+    return tryChangeAndCheckErrors(variable, loop) {
+        it.typeReference = KtPsiFactory(it).createType(newTypeText)
+    }
+}
+
+fun <TExpression : KtExpression> tryChangeAndCheckErrors(
+        expressionToChange: TExpression,
+        scopeToExclude: KtElement,
+        performChange: (TExpression) -> Unit
+): Boolean {
+    val bindingContext = expressionToChange.analyze(BodyResolveMode.FULL)
 
     // analyze the closest block which is not used as expression
-    val block = variable.parents
+    val block = expressionToChange.parents
                         .filterIsInstance<KtBlockExpression>()
                         .firstOrNull { bindingContext[BindingContext.USED_AS_EXPRESSION, it] != true }
-                ?: return false
+                ?: return true
 
-    val KEY = Key<Unit>("KEY")
-    block.putCopyableUserData(KEY, Unit)
-    variable.putCopyableUserData(KEY, Unit)
-    loop.putCopyableUserData(KEY, Unit)
+    val EXPRESSION = Key<Unit>("EXPRESSION")
+    val SCOPE_TO_EXCLUDE = Key<Unit>("SCOPE_TO_EXCLUDE")
+    expressionToChange.putCopyableUserData(EXPRESSION, Unit)
+    scopeToExclude.putCopyableUserData(SCOPE_TO_EXCLUDE, Unit)
 
-    val fileCopy = block.containingFile.copied()
-    val blockCopy: KtBlockExpression
-    val variableCopy: KtProperty
-    val loopCopy: KtForExpression
+    val blockCopy = block.copied()
+    val expressionCopy: TExpression
+    val scopeToExcludeCopy: KtElement
+    @Suppress("UNCHECKED_CAST")
     try {
-        blockCopy = fileCopy.findDescendantOfType<KtBlockExpression> { it.getCopyableUserData(KEY) != null }!!
-        variableCopy = blockCopy.findDescendantOfType<KtProperty> { it.getCopyableUserData(KEY) != null }!!
-        loopCopy = blockCopy.findDescendantOfType<KtForExpression> { it.getCopyableUserData(KEY) != null }!!
+        expressionCopy = blockCopy.findDescendantOfType<KtExpression> { it.getCopyableUserData(EXPRESSION) != null } as TExpression
+        scopeToExcludeCopy = blockCopy.findDescendantOfType<KtElement> { it.getCopyableUserData(SCOPE_TO_EXCLUDE) != null }!!
     }
     finally {
-        block.putCopyableUserData(KEY, null)
-        variable.putCopyableUserData(KEY, null)
-        loop.putCopyableUserData(KEY, null)
+        expressionToChange.putCopyableUserData(EXPRESSION, null)
+        scopeToExclude.putCopyableUserData(SCOPE_TO_EXCLUDE, null)
     }
 
-    variableCopy.typeReference = KtPsiFactory(block).createType(newTypeText)
+    performChange(expressionCopy)
 
     val resolutionScope = block.getResolutionScope(bindingContext, block.getResolutionFacade())
     val newBindingContext = blockCopy.analyzeInContext(scope = resolutionScope,
@@ -198,7 +206,7 @@ fun canChangeLocalVariableType(variable: KtProperty, newTypeText: String, loop:
                                                        dataFlowInfo = bindingContext.getDataFlowInfo(block),
                                                        trace = DelegatingBindingTrace(bindingContext, "Temporary trace"))
     //TODO: what if there were errors before?
-    return newBindingContext.diagnostics.none { it.severity == Severity.ERROR && !loopCopy.isAncestor(it.psiElement) }
+    return newBindingContext.diagnostics.none { it.severity == Severity.ERROR && !scopeToExcludeCopy.isAncestor(it.psiElement) }
 }
 
 private val NO_SIDE_EFFECT_STANDARD_CLASSES = setOf(

idea/testData/intentions/loopToCallChain/smartCastNotBroken2.kt.after

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'map{}.map{}.firstOrNull{}'"
+fun foo(list: List<String>, o: Any): Int? {
+    if (o is CharSequence) {
+        <caret>return list
+                .map { it.length + (o as String).capitalize().hashCode() }
+                .map { it * o.length }
+                .firstOrNull { it > 1000 }
+    }
+    return 0
+}

idea/testData/intentions/loopToCallChain/smartCastNotBroken2.kt.todo

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'map{}.map{}.firstOrNull{}'"
+fun foo(list: List<String>, o: Any): Int? {
+    if (o is CharSequence) {
+        <caret>for (s in list) {
+            val a = s.length + (o as String).capitalize().hashCode()
+            val x = a * o.length
+            if (x > 1000) return x
+        }
+        return null
+    }
+    return 0
+}

idea/testData/intentions/loopToCallChain/smartCastRequired3.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<Any>): String? {
+    <caret>for (o in list) {
+        if (bar(o as String)) {
+            return o
+        }
+    }
+    return null
+}
+
+fun bar(s: String): Boolean = true

idea/testData/intentions/loopToCallChain/smartCastRequired4.kt

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<Any>): String? {
+    var v: String? = null
+    <caret>for (o in list) {
+        if (bar(o as String)) {
+            v = o
+            break
+        }
+    }
+    return v
+}
+
+fun bar(s: String): Boolean = true

idea/testData/intentions/loopToCallChain/smartCastRequired5.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<Any>, target: MutableCollection<String>) {
+    <caret>for (o in list) {
+        if (bar(o as String)) {
+            target.add(o)
+        }
+    }
+}
+
+fun bar(s: String): Boolean = true