KT-14294 for to stdlib to firstOrNull/lastOrNull: intention is absent if there is additional var before for loop

#KT-14294 Fixed

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

@@ -16,12 +16,18 @@
 
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
+import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
+import org.jetbrains.kotlin.idea.core.copied
 import org.jetbrains.kotlin.idea.core.replaced
-import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.idea.util.getResolutionScope
-import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
-import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
+import org.jetbrains.kotlin.resolve.scopes.utils.findClassifier
+import org.jetbrains.kotlin.resolve.scopes.utils.findFunction
+import org.jetbrains.kotlin.resolve.scopes.utils.findPackage
+import org.jetbrains.kotlin.resolve.scopes.utils.findVariable
 
 /**
  * Base class for [ResultTransformation]'s that replaces the loop-expression with the result call chain
@@ -30,8 +36,9 @@ abstract class ReplaceLoopResultTransformation(final override val loop: KtForExp
 
     override val commentSavingRange = PsiChildRange.singleElement(loop.unwrapIfLabeled())
 
-    override val expressionToBeReplacedByResultCallChain: KtExpression
+    override fun generateExpressionToReplaceLoopAndCheckErrors(resultCallChain: KtExpression): KtExpression {
-        get() = loop.unwrapIfLabeled()
+        return resultCallChain
+    }
 
     override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
         return loop.unwrapIfLabeled().replaced(resultCallChain)
@@ -48,8 +55,30 @@ abstract class AssignToVariableResultTransformation(
 
     override val commentSavingRange = PsiChildRange(initialization.initializationStatement, loop.unwrapIfLabeled())
 
-    override val expressionToBeReplacedByResultCallChain: KtExpression
+    override fun generateExpressionToReplaceLoopAndCheckErrors(resultCallChain: KtExpression): KtExpression {
-        get() = initialization.initializer
+        val psiFactory = KtPsiFactory(resultCallChain)
+        val initializationStatement = initialization.initializationStatement
+        if (initializationStatement is KtVariableDeclaration) {
+            val resolutionScope = loop.getResolutionScope()
+
+            fun isUniqueName(name: String): Boolean {
+                val identifier = Name.identifier(name)
+                return resolutionScope.findVariable(identifier, NoLookupLocation.FROM_IDE) == null
+                       && resolutionScope.findFunction(identifier, NoLookupLocation.FROM_IDE) == null
+                       && resolutionScope.findClassifier(identifier, NoLookupLocation.FROM_IDE) == null
+                       && resolutionScope.findPackage(identifier) == null
+            }
+            val uniqueName = KotlinNameSuggester.suggestNameByName("test", ::isUniqueName)
+
+            val copy = initializationStatement.copied()
+            copy.initializer!!.replace(resultCallChain)
+            copy.setName(uniqueName)
+            return copy
+        }
+        else {
+            return psiFactory.createExpressionByPattern("$0 = $1", initialization.variable.nameAsSafeName, resultCallChain)
+        }
+    }
 
     override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
         initialization.initializer.replace(resultCallChain)

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

@@ -79,7 +79,7 @@ interface ResultTransformation : Transformation {
 
     val commentSavingRange: PsiChildRange
 
-    val expressionToBeReplacedByResultCallChain: KtExpression
+    fun generateExpressionToReplaceLoopAndCheckErrors(resultCallChain: KtExpression): KtExpression
 
     /**
      * Implementations of this method are obliged to update [commentSavingRangeHolder] when deleting or adding any element into the tree

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

@@ -293,8 +293,8 @@ private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: MatchRe
         if (preservedSmartCastCount == smartCastCount) return true
 
         // not all smart cast expressions has been found in the result or have the same type after conversion, perform more expensive check
-        val expressionToBeReplaced = matchResult.transformationMatch.resultTransformation.expressionToBeReplacedByResultCallChain
+        val expression = matchResult.transformationMatch.resultTransformation.generateExpressionToReplaceLoopAndCheckErrors(callChain)
-        if (!tryChangeAndCheckErrors(expressionToBeReplaced, loop) { it.replace(callChain) }) return false
+        if (!tryChangeAndCheckErrors(loop) { it.replace(expression) }) return false
 
         return true
     }

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

@@ -135,8 +135,9 @@ object FindTransformationMatcher : TransformationMatcher {
             return generator.generate(chainedCallGenerator)
         }
 
-        override val expressionToBeReplacedByResultCallChain: KtExpression
+        override fun generateExpressionToReplaceLoopAndCheckErrors(resultCallChain: KtExpression): KtExpression {
-            get() = endReturn.returnedExpression!!
+            return KtPsiFactory(resultCallChain).createExpressionByPattern("return $0", resultCallChain)
+        }
 
         override fun convertLoop(resultCallChain: KtExpression, commentSavingRangeHolder: CommentSavingRangeHolder): KtExpression {
             endReturn.returnedExpression!!.replace(resultCallChain)

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

@@ -230,7 +230,7 @@ fun canChangeLocalVariableType(variable: KtProperty, newTypeText: String, loop:
 
 fun <TExpression : KtExpression> tryChangeAndCheckErrors(
         expressionToChange: TExpression,
-        scopeToExclude: KtElement,
+        scopeToExclude: KtElement? = null,
         performChange: (TExpression) -> Unit
 ): Boolean {
     val bindingContext = expressionToChange.analyze(BodyResolveMode.FULL)
@@ -246,19 +246,19 @@ fun <TExpression : KtExpression> tryChangeAndCheckErrors(
     val SCOPE_TO_EXCLUDE = Key<Unit>("SCOPE_TO_EXCLUDE")
 
     expressionToChange.putCopyableUserData(EXPRESSION, Unit)
-    scopeToExclude.putCopyableUserData(SCOPE_TO_EXCLUDE, Unit)
+    scopeToExclude?.putCopyableUserData(SCOPE_TO_EXCLUDE, Unit)
 
     val blockCopy = block.copied()
     val expressionCopy: TExpression
-    val scopeToExcludeCopy: KtElement
+    val scopeToExcludeCopy: KtElement?
     @Suppress("UNCHECKED_CAST")
     try {
         expressionCopy = blockCopy.findDescendantOfType<KtExpression> { it.getCopyableUserData(EXPRESSION) != null } as TExpression
-        scopeToExcludeCopy = blockCopy.findDescendantOfType<KtElement> { it.getCopyableUserData(SCOPE_TO_EXCLUDE) != null }!!
+        scopeToExcludeCopy = blockCopy.findDescendantOfType<KtElement> { it.getCopyableUserData(SCOPE_TO_EXCLUDE) != null }
     }
     finally {
         expressionToChange.putCopyableUserData(EXPRESSION, null)
-        scopeToExclude.putCopyableUserData(SCOPE_TO_EXCLUDE, null)
+        scopeToExclude?.putCopyableUserData(SCOPE_TO_EXCLUDE, null)
     }
 
     performChange(expressionCopy)

idea/testData/intentions/loopToCallChain/firstOrNull/KT14292.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterNotNull().firstOrNull{}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().firstOrNull{}'"
+
+fun getFirstValue() = "value"
+
+fun foo(list: List<String?>): String? {
+    var found: String? = null
+    val value = getFirstValue()
+    <caret>for (s in list)
+        if (s != null)
+            if(!s.startsWith("IMG:"))
+                if (s.contains(value)) {
+                    found = s
+                    break
+                }
+    return found
+}

idea/testData/intentions/loopToCallChain/firstOrNull/KT14292.kt.after

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterNotNull().firstOrNull{}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().firstOrNull{}'"
+
+fun getFirstValue() = "value"
+
+fun foo(list: List<String?>): String? {
+    val value = getFirstValue()
+    val <caret>found: String? = list
+            .filterNotNull()
+            .firstOrNull { !it.startsWith("IMG:") && it.contains(value) }
+    return found
+}

idea/testData/intentions/loopToCallChain/firstOrNull/KT14292.kt.after2

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterNotNull().firstOrNull{}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().firstOrNull{}'"
+
+fun getFirstValue() = "value"
+
+fun foo(list: List<String?>): String? {
+    val value = getFirstValue()
+    val <caret>found: String? = list
+            .asSequence()
+            .filterNotNull()
+            .firstOrNull { !it.startsWith("IMG:") && it.contains(value) }
+    return found
+}

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

@@ -0,0 +1,16 @@
+// WITH_RUNTIME
+// 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) {
+        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/smartCastNotBroken3.kt.after

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+// 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) {
+        val result: Any? = list
+                .map { it.length + (o as String).capitalize().hashCode() }
+                .map { it * o.length }
+                .firstOrNull { it > 1000 }
+    }
+}

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

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// 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) {
+        val result: Any? = list
+                .asSequence()
+                .map { it.length + (o as String).capitalize().hashCode() }
+                .map { it * o.length }
+                .firstOrNull { it > 1000 }
+    }
+}

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

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'map{}.map{}.firstOrNull{}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().map{}.map{}.firstOrNull{}'"
+fun foo(list: List<String>, o: Any) {
+    var result: Any? = ""
+    if (o is CharSequence) {
+        result = 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/smartCastNotBroken4.kt.after

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

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

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

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

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'map{}.map{}.filterTo(){}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().map{}.map{}.filterTo(){}'"
+fun foo(list: List<String>, o: Any, result: MutableCollection<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) {
+                result.add(x)
+            }
+        }
+    }
+}

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

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

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

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'map{}.map{}.filterTo(){}'"
+// INTENTION_TEXT_2: "Replace with 'asSequence().map{}.map{}.filterTo(){}'"
+fun foo(list: List<String>, o: Any, result: MutableCollection<Int>) {
+    if (o is CharSequence) {
+        list
+                .asSequence()
+                .map { it.length + (o as String).capitalize().hashCode() }
+                .map { it * o.length }
+                .filterTo(result) { it > 1000 }
+    }
+}

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

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

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

@@ -613,6 +613,12 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             doTest(fileName);
         }
 
+        @TestMetadata("KT14292.kt")
+        public void testKT14292() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/firstOrNull/KT14292.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("let.kt")
         public void testLet() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/firstOrNull/let.kt");
@@ -1174,6 +1180,24 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             doTest(fileName);
         }
 
+        @TestMetadata("smartCastNotBroken3.kt")
+        public void testSmartCastNotBroken3() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken3.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("smartCastNotBroken4.kt")
+        public void testSmartCastNotBroken4() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken4.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("smartCastNotBroken5.kt")
+        public void testSmartCastNotBroken5() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken5.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("smartCastNotNullRequired.kt")
         public void testSmartCastNotNullRequired() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotNullRequired.kt");
@@ -1216,6 +1240,12 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             doTest(fileName);
         }
 
+        @TestMetadata("smartCastRequired6.kt")
+        public void testSmartCastRequired6() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastRequired6.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("smartCastThisRequired.kt")
         public void testSmartCastThisRequired() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastThisRequired.kt");

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

@@ -8521,6 +8521,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("KT14292.kt")
+            public void testKT14292() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/firstOrNull/KT14292.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("let.kt")
             public void testLet() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/firstOrNull/let.kt");
@@ -9082,6 +9088,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("smartCastNotBroken3.kt")
+            public void testSmartCastNotBroken3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken3.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("smartCastNotBroken4.kt")
+            public void testSmartCastNotBroken4() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken4.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("smartCastNotBroken5.kt")
+            public void testSmartCastNotBroken5() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotBroken5.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("smartCastNotNullRequired.kt")
             public void testSmartCastNotNullRequired() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastNotNullRequired.kt");
@@ -9124,6 +9148,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("smartCastRequired6.kt")
+            public void testSmartCastRequired6() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastRequired6.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("smartCastThisRequired.kt")
             public void testSmartCastThisRequired() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/smartCasts/smartCastThisRequired.kt");