More cases supported for "any"

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

@@ -27,7 +27,7 @@ import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
 class FindAndAssignTransformation(
         private val loop: KtForExpression,
         override val inputVariable: KtCallableDeclaration,
-        private val stdlibFunName: String,
+        private val generator: (chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression,
         private val initialDeclaration: KtProperty,
         private val filter: KtExpression? = null
 ) : ResultTransformation {
@@ -35,20 +35,14 @@ class FindAndAssignTransformation(
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
         if (previousTransformation !is FilterTransformation) return null
         assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"}
-        return FindAndAssignTransformation(loop, previousTransformation.inputVariable, stdlibFunName, initialDeclaration, previousTransformation.buildRealCondition())
+        return FindAndAssignTransformation(loop, previousTransformation.inputVariable, generator, initialDeclaration, previousTransformation.buildRealCondition())
     }
 
     override val commentSavingRange = PsiChildRange(initialDeclaration, loop.unwrapIfLabeled())
     override val commentRestoringRange = commentSavingRange.withoutLastStatement()
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        return if (filter == null) {
+        return generator(chainedCallGenerator, filter)
-            chainedCallGenerator.generate("$stdlibFunName()")
-        }
-        else {
-            val lambda = generateLambda(inputVariable, filter)
-            chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda)
-        }
     }
 
     override fun convertLoop(resultCallChain: KtExpression): KtExpression {
@@ -107,9 +101,9 @@ class FindAndAssignTransformation(
             val usageCountInLoop = ReferencesSearch.search(declarationBeforeLoop, LocalSearchScope(state.outerLoop)).count()
             if (usageCountInLoop != 1) return null // this should be the only usage of this variable inside the loop
 
-            val stdlibFunName = stdlibFunNameForFind(right, initializer, state.workingVariable, findFirst) ?: return null
+            val generator = buildFindOperationGenerator(right, initializer, state.workingVariable, findFirst) ?: return null
 
-            val transformation = FindAndAssignTransformation(state.outerLoop, state.workingVariable, stdlibFunName, declarationBeforeLoop)
+            val transformation = FindAndAssignTransformation(state.outerLoop, state.workingVariable, generator, declarationBeforeLoop)
             return ResultTransformationMatch(transformation)
         }
     }

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

@@ -27,7 +27,7 @@ import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
 class FindAndReturnTransformation(
         private val loop: KtForExpression,
         override val inputVariable: KtCallableDeclaration,
-        private val stdlibFunName: String,
+        private val generator: (chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression,
         private val endReturn: KtReturnExpression,
         private val filter: KtExpression? = null
 ) : ResultTransformation {
@@ -35,7 +35,7 @@ class FindAndReturnTransformation(
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
         if (previousTransformation !is FilterTransformation) return null
         assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"}
-        return FindAndReturnTransformation(loop, previousTransformation.inputVariable, stdlibFunName, endReturn, previousTransformation.buildRealCondition())
+        return FindAndReturnTransformation(loop, previousTransformation.inputVariable, generator, endReturn, previousTransformation.buildRealCondition())
     }
 
     override val commentSavingRange = PsiChildRange(loop.unwrapIfLabeled(), endReturn)
@@ -43,13 +43,7 @@ class FindAndReturnTransformation(
     override val commentRestoringRange = commentSavingRange.withoutFirstStatement()
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        return if (filter == null) {
+        return generator(chainedCallGenerator, filter)
-            chainedCallGenerator.generate("$stdlibFunName()")
-        }
-        else {
-            val lambda = generateLambda(inputVariable, filter)
-            chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda)
-        }
     }
 
     override fun convertLoop(resultCallChain: KtExpression): KtExpression {
@@ -78,9 +72,9 @@ class FindAndReturnTransformation(
             val returnValueInLoop = returnInLoop.returnedExpression ?: return null
             val returnValueAfterLoop = returnAfterLoop.returnedExpression ?: return null
 
-            val stdlibFunName = stdlibFunNameForFind(returnValueInLoop, returnValueAfterLoop, state.workingVariable, findFirst = true) ?: return null
+            val generator = buildFindOperationGenerator(returnValueInLoop, returnValueAfterLoop, state.workingVariable, findFirst = true) ?: return null
 
-            val transformation = FindAndReturnTransformation(state.outerLoop, state.workingVariable, stdlibFunName, returnAfterLoop)
+            val transformation = FindAndReturnTransformation(state.outerLoop, state.workingVariable, generator, returnAfterLoop)
             return ResultTransformationMatch(transformation)
         }
     }

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

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
@@ -95,16 +96,48 @@ fun KtProperty.hasWriteUsages(): Boolean {
     }
 }
 
-fun stdlibFunNameForFind(valueIfFound: KtExpression, valueIfNotFound: KtExpression, workingVariable: KtCallableDeclaration, findFirst: Boolean): String? {
+fun buildFindOperationGenerator(
-    return when {
+        valueIfFound: KtExpression,
-        valueIfNotFound.isNullExpression() && valueIfFound.isVariableReference(workingVariable) -> if (findFirst) "firstOrNull" else "lastOrNull"
+        valueIfNotFound: KtExpression,
+        workingVariable: KtCallableDeclaration,
+        findFirst: Boolean
+): ((chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression)?  {
+    assert(valueIfFound.isPhysical)
+    assert(valueIfNotFound.isPhysical)
+
+    fun generateChainedCall(stdlibFunName: String, chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?): KtExpression {
+        return if (filter == null) {
+            chainedCallGenerator.generate("$stdlibFunName()")
+        }
+        else {
+            val lambda = generateLambda(workingVariable, filter)
+            chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda)
+        }
+    }
+
+
+    val stdlibFunName = when {
+        valueIfNotFound.isNullExpression() && valueIfFound.isVariableReference(workingVariable) -> if (findFirst) "firstOrNull" else "lastOrNull" //TODO: ?: if not null
 
         valueIfFound.isTrueConstant() && valueIfNotFound.isFalseConstant() -> "any"
 
         valueIfFound.isFalseConstant() && valueIfNotFound.isTrueConstant() -> "none"
 
-        else -> /*TODO: allow other constants*/ null
+        workingVariable.hasUsages(listOf(valueIfFound)) -> /*TODO*/ return null
+
+        // initial value is compile-time constant
+        ConstantExpressionEvaluator.getConstant(valueIfNotFound, valueIfNotFound.analyze(BodyResolveMode.PARTIAL)) != null -> {
+            return { chainedCallGenerator, filter ->
+                val chainedCall = generateChainedCall("any", chainedCallGenerator, filter)
+                KtPsiFactory(chainedCall).createExpressionByPattern("if ($0) $1 else $2", chainedCall, valueIfFound, valueIfNotFound)
+            }
+
+        }
+
+        else -> return null
     }
+
+    return { chainedCallGenerator, filter -> generateChainedCall(stdlibFunName, chainedCallGenerator, filter) }
 }
 
 fun KtExpressionWithLabel.isBreakOrContinueOfLoop(loop: KtLoopExpression): Boolean {

idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+fun foo(list: List<String>) {
+    var result = 0
+    <caret>for (s in list) {
+        if (s.length > 0) {
+            result = 1
+            break
+        }
+    }
+}

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

@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+fun foo(list: List<String>) {
+    <caret>val result = if (list.any { it.length > 0 }) 1 else 0
+}

idea/testData/intentions/loopToCallChain/any_ifAssign_nonBooleanResult2.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<String>) {
+    var result = takeInt()
+    <caret>for (s in list) {
+        if (s.length > 0) {
+            result = 1
+            break
+        }
+    }
+}
+
+fun takeInt(): Int = 0