Honor errors in code before when checking smart casts

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

@@ -18,9 +18,11 @@ package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.UserDataHolderBase
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.diagnostics.DiagnosticFactory
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
@@ -244,10 +246,19 @@ fun <TExpression : KtExpression> tryChangeAndCheckErrors(
     // we declare these keys locally to avoid possible race-condition problems if this code is executed in 2 threads simultaneously
     val EXPRESSION = Key<Unit>("EXPRESSION")
     val SCOPE_TO_EXCLUDE = Key<Unit>("SCOPE_TO_EXCLUDE")
+    val ERRORS_BEFORE = Key<Collection<DiagnosticFactory<*>>>("ERRORS_BEFORE")
 
     expressionToChange.putCopyableUserData(EXPRESSION, Unit)
     scopeToExclude?.putCopyableUserData(SCOPE_TO_EXCLUDE, Unit)
 
+    block.forEachDescendantOfType<PsiElement> { element ->
+        val errors = bindingContext.diagnostics.forElement(element)
+                .filter { it.severity == Severity.ERROR }
+        if (errors.isNotEmpty()) {
+            element.putCopyableUserData(ERRORS_BEFORE, errors.map { it.factory })
+        }
+    }
+
     val blockCopy = block.copied()
     val expressionCopy: TExpression
     val scopeToExcludeCopy: KtElement?
@@ -268,8 +279,11 @@ fun <TExpression : KtExpression> tryChangeAndCheckErrors(
                                                        contextExpression = block,
                                                        dataFlowInfo = bindingContext.getDataFlowInfo(block),
                                                        trace = DelegatingBindingTrace(bindingContext, "Temporary trace"))
-    //TODO: what if there were errors before?
-    return newBindingContext.diagnostics.none { it.severity == Severity.ERROR && !scopeToExcludeCopy.isAncestor(it.psiElement) }
+    return newBindingContext.diagnostics.none {
+        it.severity == Severity.ERROR
+            && !scopeToExcludeCopy.isAncestor(it.psiElement)
+            && it.factory !in (it.psiElement.getCopyableUserData(ERRORS_BEFORE) ?: emptyList())
+    }
 }
 
 private val NO_SIDE_EFFECT_STANDARD_CLASSES = setOf(

idea/testData/intentions/loopToCallChain/smartCasts/errorOutsideLoop.kt

@@ -0,0 +1,19 @@
+// WITH_RUNTIME
+// ERROR: Unresolved reference: unresolvedFun
+// INTENTION_TEXT: "Replace with 'map{}.map{}.firstOrNull{}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().map{}.map{}.firstOrNull{}'"
+fun foo(list: List<String>, o: Any) {
+    if (o is CharSequence) {
+        unresolvedFun()
+
+        var result: Any? = null
+        <caret>for (s in list) {
+            val a = s.length + (o as String).capitalize().hashCode()
+            val x = a * o.length
+            if (x > 1000) {
+                result = x
+                break
+            }
+        }
+    }
+}

idea/testData/intentions/loopToCallChain/smartCasts/errorOutsideLoop.kt.after

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+// ERROR: Unresolved reference: unresolvedFun
+// INTENTION_TEXT: "Replace with 'map{}.map{}.firstOrNull{}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().map{}.map{}.firstOrNull{}'"
+fun foo(list: List<String>, o: Any) {
+    if (o is CharSequence) {
+        unresolvedFun()
+
+        val result: Any? = list
+                .map { it.length + (o as String).capitalize().hashCode() }
+                .map { it * o.length }
+                .firstOrNull { it > 1000 }
+    }
+}

idea/testData/intentions/loopToCallChain/smartCasts/errorOutsideLoop.kt.after2

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+// ERROR: Unresolved reference: unresolvedFun
+// INTENTION_TEXT: "Replace with 'map{}.map{}.firstOrNull{}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().map{}.map{}.firstOrNull{}'"
+fun foo(list: List<String>, o: Any) {
+    if (o is CharSequence) {
+        unresolvedFun()
+
+        val result: Any? = list
+                .asSequence()
+                .map { it.length + (o as String).capitalize().hashCode() }
+                .map { it * o.length }
+                .firstOrNull { it > 1000 }
+    }
+}

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

@@ -1168,6 +1168,12 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/loopToCallChain/smartCasts"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }
 
+        @TestMetadata("errorOutsideLoop.kt")
+        public void testErrorOutsideLoop() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/errorOutsideLoop.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("smartCastNotBroken.kt")
         public void testSmartCastNotBroken() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken.kt");

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

@@ -9076,6 +9076,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/loopToCallChain/smartCasts"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
             }
 
+            @TestMetadata("errorOutsideLoop.kt")
+            public void testErrorOutsideLoop() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/errorOutsideLoop.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("smartCastNotBroken.kt")
             public void testSmartCastNotBroken() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken.kt");