KT-14209 for loop to stdlib: sum is not recognized if before was check for type

#KT-14209 Fixed

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.MapTransformation
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.KotlinType
 
@@ -66,7 +67,7 @@ abstract class SumTransformationBase(
 
             val value = statement.right ?: return null
 
-            val valueType = value.analyze(BodyResolveMode.PARTIAL).getType(value)?.toSupportedType() ?: return null
+            val valueType = value.typeWithSmartCast()?.toSupportedType() ?: return null
             val sumType = (variableInitialization.variable.resolveToDescriptorIfAny() as? VariableDescriptor)?.type?.toSupportedType() ?: return null
 
             val conversionFunctionName = when (sumType) {
@@ -140,6 +141,12 @@ abstract class SumTransformationBase(
                 else -> null
             }
         }
+
+        private fun KtExpression.typeWithSmartCast(): KotlinType? {
+            val bindingContext = analyze(BodyResolveMode.PARTIAL)
+            return bindingContext[BindingContext.SMARTCAST, this]?.defaultType
+                   ?: bindingContext.getType(this)
+        }
     }
 }
 

idea/testData/intentions/loopToCallChain/sum/smartCast.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIsInstance<>().sum()'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterIsInstance<>().sum()'"
+fun foo(list: List<Any>){
+    var result = 0
+    <caret>for (l in list)
+        if (l is Int)
+            result += l
+}

idea/testData/intentions/loopToCallChain/sum/smartCast.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIsInstance<>().sum()'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterIsInstance<>().sum()'"
+fun foo(list: List<Any>){
+    val <caret>result = list
+            .filterIsInstance<Int>()
+            .sum()
+}

idea/testData/intentions/loopToCallChain/sum/smartCast.kt.after2

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIsInstance<>().sum()'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterIsInstance<>().sum()'"
+fun foo(list: List<Any>){
+    val <caret>result = list
+            .asSequence()
+            .filterIsInstance<Int>()
+            .sum()
+}

idea/testData/intentions/loopToCallChain/sum/smartCastNotNull.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterNotNull().sum()'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().sum()'"
+fun foo(list: List<Int?>){
+    var result = 0
+    <caret>for (l in list)
+        if (l != null)
+            result += l
+}

idea/testData/intentions/loopToCallChain/sum/smartCastNotNull.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterNotNull().sum()'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().sum()'"
+fun foo(list: List<Int?>){
+    val <caret>result = list
+            .filterNotNull()
+            .sum()
+}

idea/testData/intentions/loopToCallChain/sum/smartCastNotNull.kt.after2

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterNotNull().sum()'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().sum()'"
+fun foo(list: List<Int?>){
+    val <caret>result = list
+            .asSequence()
+            .filterNotNull()
+            .sum()
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTest2Generated.java

@@ -1267,6 +1267,18 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             doTest(fileName);
         }
 
+        @TestMetadata("smartCast.kt")
+        public void testSmartCast() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/sum/smartCast.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("smartCastNotNull.kt")
+        public void testSmartCastNotNull() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/sum/smartCastNotNull.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("sumByDouble.kt")
         public void testSumByDouble() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/sum/sumByDouble.kt");

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -9175,6 +9175,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("smartCast.kt")
+            public void testSmartCast() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/sum/smartCast.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("smartCastNotNull.kt")
+            public void testSmartCastNotNull() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/sum/smartCastNotNull.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("sumByDouble.kt")
             public void testSumByDouble() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/sum/sumByDouble.kt");