diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/ForEachTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/ForEachTransformation.kt index 08799423f56..f13828a8408 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/ForEachTransformation.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/ForEachTransformation.kt @@ -17,9 +17,15 @@ package org.jetbrains.kotlin.idea.intentions.loopToCallChain.result import org.jetbrains.kotlin.idea.intentions.loopToCallChain.* -import org.jetbrains.kotlin.psi.* -import org.jetbrains.kotlin.psi.KtPsiUtil.isAssignment +import org.jetbrains.kotlin.idea.references.ReferenceAccess +import org.jetbrains.kotlin.idea.references.mainReference +import org.jetbrains.kotlin.idea.references.readWriteAccess +import org.jetbrains.kotlin.psi.KtCallableDeclaration +import org.jetbrains.kotlin.psi.KtExpression +import org.jetbrains.kotlin.psi.KtForExpression +import org.jetbrains.kotlin.psi.KtNameReferenceExpression import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType +import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector class ForEachTransformation( loop: KtForExpression, @@ -57,9 +63,21 @@ class ForEachTransformation( val statement = state.statements.singleOrNull() ?: return null - // check if contains assignment to non-qualified variable - in this case only use of lazy sequence is correct - if (!state.lazySequence - && statement.anyDescendantOfType { isAssignment(it) && it.left is KtNameReferenceExpression }) return null + if (!state.lazySequence) { + // check if it changes any variable that has other usages in the loop - then only lazy sequence is correct + val onlySequence = statement.anyDescendantOfType { nameExpr -> + if (nameExpr.getQualifiedExpressionForSelector() == null) { + // we don't use resolve for write access detection because '+=' which modifies a collection is kind of write access too + val isWrite = nameExpr.readWriteAccess(useResolveForReadWrite = false) != ReferenceAccess.READ + if (isWrite) { + val variable = nameExpr.mainReference.resolve() as? KtCallableDeclaration + if (variable != null && variable.countUsages(state.outerLoop) > 1) return@anyDescendantOfType true + } + } + false + } + if (onlySequence) return null + } //TODO: should we disallow it for complicated statements like loops, if, when? val transformation = ForEachTransformation(state.outerLoop, state.inputVariable, state.indexVariable, statement) diff --git a/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt b/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt index 518d0e76879..380d2993bfe 100644 --- a/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt +++ b/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt @@ -1,5 +1,5 @@ // WITH_RUNTIME -// INTENTION_TEXT: "Replace with 'filter{}.forEach{}'" +// IS_APPLICABLE: false // INTENTION_TEXT_2: "Replace with 'asSequence().filter{}.forEach{}'" fun foo(list: List): Int { var count = 0 diff --git a/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt.after b/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt.after deleted file mode 100644 index 2da9e9c874f..00000000000 --- a/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt.after +++ /dev/null @@ -1,10 +0,0 @@ -// WITH_RUNTIME -// INTENTION_TEXT: "Replace with 'filter{}.forEach{}'" -// INTENTION_TEXT_2: "Replace with 'asSequence().filter{}.forEach{}'" -fun foo(list: List): Int { - var count = 0 - list - .filter { it.length > count } - .forEach { count++ } - return count -} \ No newline at end of file diff --git a/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt.after2 b/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt.after2 index 172d29a0859..1a8467bc02c 100644 --- a/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt.after2 +++ b/idea/testData/intentions/loopToCallChain/count/variableUsedBefore.kt.after2 @@ -1,5 +1,5 @@ // WITH_RUNTIME -// INTENTION_TEXT: "Replace with 'filter{}.forEach{}'" +// IS_APPLICABLE: false // INTENTION_TEXT_2: "Replace with 'asSequence().filter{}.forEach{}'" fun foo(list: List): Int { var count = 0