count() to merge with filterNotNull and filterIsInstance + changed test for KT-14191 to use sum()

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

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
-import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformationBase
 import org.jetbrains.kotlin.psi.*
 
 class CountTransformation(
@@ -31,10 +31,13 @@ class CountTransformation(
 ) : AssignToVariableResultTransformation(loop, initialization) {
 
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
-        if (previousTransformation !is FilterTransformation) return null
+        if (previousTransformation !is FilterTransformationBase) return null
         if (previousTransformation.indexVariable != null) return null
-        assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"}
+        val newFilter = if (filter == null)
-        return CountTransformation(loop, previousTransformation.inputVariable, initialization, previousTransformation.effectiveCondition())
+            previousTransformation.effectiveCondition
+        else
+            KtPsiFactory(filter).createExpressionByPattern("$0 && $1", previousTransformation.effectiveCondition, filter)
+        return CountTransformation(loop, previousTransformation.inputVariable, initialization, newFilter)
     }
 
     override val presentation: String

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

@@ -224,7 +224,7 @@ object FindTransformationMatcher : TransformationMatcher {
         assert(valueIfFound.isPhysical)
         assert(valueIfNotFound.isPhysical)
 
-        val filter = filterTransformation?.effectiveCondition()
+        val filter = filterTransformation?.effectiveCondition
 
         if (indexVariable != null) {
             if (filterTransformation == null) return null // makes no sense, indexVariable must be always null

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt

@@ -24,15 +24,26 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.blockExpressionsOrSingle
 
+abstract class FilterTransformationBase : SequenceTransformation {
+    abstract val effectiveCondition: KtExpression
+    abstract val inputVariable: KtCallableDeclaration
+    abstract val indexVariable: KtCallableDeclaration?
+
+    override val affectsIndex: Boolean
+        get() = true
+}
+
 class FilterTransformation(
         override val loop: KtForExpression,
-        val inputVariable: KtCallableDeclaration,
+        override val inputVariable: KtCallableDeclaration,
-        val indexVariable: KtCallableDeclaration?,
+        override val indexVariable: KtCallableDeclaration?,
         val condition: KtExpression,
         val isInverse: Boolean
-) : SequenceTransformation {
+) : FilterTransformationBase() {
 
-    fun effectiveCondition() = if (isInverse) condition.negate() else condition
+    override val effectiveCondition: KtExpression by lazy {
+        if (isInverse) condition.negate() else condition
+    }
 
     private val functionName = when {
         indexVariable != null -> "filterIndexed"
@@ -40,15 +51,12 @@ class FilterTransformation(
         else -> "filter"
     }
 
-    override val affectsIndex: Boolean
-        get() = true
-
     override val presentation: String
         get() = "$functionName{}"
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         val lambda = if (indexVariable != null)
-            generateLambda(inputVariable, indexVariable, effectiveCondition())
+            generateLambda(inputVariable, indexVariable, effectiveCondition)
         else
             generateLambda(inputVariable, condition)
         return chainedCallGenerator.generate("$0$1:'{}'", functionName, lambda)
@@ -97,7 +105,7 @@ class FilterTransformation(
 
                     val indexVariable = transformation.indexVariable ?: nextTransformation.indexVariable
                     val mergedCondition = KtPsiFactory(state.outerLoop).createExpressionByPattern(
-                            "$0 && $1", transformation.effectiveCondition(), nextTransformation.effectiveCondition())
+                            "$0 && $1", transformation.effectiveCondition, nextTransformation.effectiveCondition)
                     transformation = FilterTransformation(state.outerLoop, transformation.inputVariable, indexVariable, mergedCondition, isInverse = false) //TODO: build filterNot in some cases?
                     currentState = nextState
                 }
@@ -167,7 +175,7 @@ class FilterTransformation(
             ) {
                 val typeRef = effectiveCondition.typeReference
                 if (typeRef != null) {
-                    return FilterIsInstanceTransformation(loop, typeRef)
+                    return FilterIsInstanceTransformation(loop, inputVariable, typeRef)
                 }
             }
 
@@ -176,7 +184,7 @@ class FilterTransformation(
                 && effectiveCondition.right.isNullExpression()
                 && effectiveCondition.left.isSimpleName(inputVariable.nameAsSafeName)
             ) {
-                return FilterNotNullTransformation(loop)
+                return FilterNotNullTransformation(loop, inputVariable)
             }
 
             return FilterTransformation(loop, inputVariable, null, condition, isInverse)
@@ -186,11 +194,15 @@ class FilterTransformation(
 
 class FilterIsInstanceTransformation(
         override val loop: KtForExpression,
+        override val inputVariable: KtCallableDeclaration,
         private val type: KtTypeReference
-) : SequenceTransformation {
+) : FilterTransformationBase() {
 
-    override val affectsIndex: Boolean
+    override val effectiveCondition by lazy {
-        get() = true
+        KtPsiFactory(loop).createExpressionByPattern("$0 is $1", inputVariable.nameAsSafeName, type)
+    }
+
+    override val indexVariable: KtCallableDeclaration? get() = null
 
     override val presentation: String
         get() = "filterIsInstance<>()"
@@ -200,7 +212,17 @@ class FilterIsInstanceTransformation(
     }
 }
 
-class FilterNotNullTransformation(override val loop: KtForExpression) : SequenceTransformation {
+class FilterNotNullTransformation(
+        override val loop: KtForExpression,
+        override val inputVariable: KtCallableDeclaration
+) : FilterTransformationBase() {
+
+    override val effectiveCondition: KtExpression by lazy {
+        KtPsiFactory(loop).createExpressionByPattern("$0 != null", inputVariable.nameAsSafeName)
+    }
+
+    override val indexVariable: KtCallableDeclaration? get() = null
+
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? {
         if (previousTransformation is MapTransformation) {
             return MapTransformation(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, previousTransformation.mapping, mapNotNull = true)
@@ -208,9 +230,6 @@ class FilterNotNullTransformation(override val loop: KtForExpression) : Sequence
         return null
     }
 
-    override val affectsIndex: Boolean
-        get() = true
-
     override val presentation: String
         get() = "filterNotNull()"
 

idea/testData/intentions/loopToCallChain/KT14191.kt

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

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

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

idea/testData/intentions/loopToCallChain/KT14191.kt.after2

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

idea/testData/intentions/loopToCallChain/count/KT14191.kt

@@ -1,12 +0,0 @@
-// WITH_RUNTIME
-// INTENTION_TEXT: "Replace with 'filterNotNull().count()'"
-// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().count()'"
-fun f11(list: List<Any?>): Int{
-    var objs = 0
-    <caret>for (d in list) {
-        if (d != null) {
-            objs++
-        }
-    }
-    return objs
-}

idea/testData/intentions/loopToCallChain/count/KT14191.kt.after

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

idea/testData/intentions/loopToCallChain/count/KT14191.kt.after2

@@ -1,10 +0,0 @@
-// WITH_RUNTIME
-// INTENTION_TEXT: "Replace with 'filterNotNull().count()'"
-// INTENTION_TEXT_2: "Replace with 'asSequence().filterNotNull().count()'"
-fun f11(list: List<Any?>): Int{
-    val <caret>objs = list
-            .asSequence()
-            .filterNotNull()
-            .count()
-    return objs
-}

idea/testData/intentions/loopToCallChain/count/countIsInstance.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'count{}'"
+// IS_APPLICABLE_2: false
+fun f11(list: List<Any?>): Int{
+    var c = 0
+    <caret>for (d in list) {
+        if (d is String) {
+            c++
+        }
+    }
+    return c
+}

idea/testData/intentions/loopToCallChain/count/countIsInstance.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'count{}'"
+// IS_APPLICABLE_2: false
+fun f11(list: List<Any?>): Int{
+    val <caret>c = list.count { it is String }
+    return c
+}

idea/testData/intentions/loopToCallChain/count/countNotNull.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'count{}'"
+// IS_APPLICABLE_2: false
+fun f11(list: List<Any?>): Int{
+    var c = 0
+    <caret>for (d in list) {
+        if (d != null) {
+            c++
+        }
+    }
+    return c
+}

idea/testData/intentions/loopToCallChain/count/countNotNull.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'count{}'"
+// IS_APPLICABLE_2: false
+fun f11(list: List<Any?>): Int{
+    val <caret>c = list.count { it != null }
+    return c
+}

idea/testData/intentions/loopToCallChain/count/countSomethingAndNotNull.kt

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'count{}'"
+// IS_APPLICABLE_2: false
+fun f(list: List<Any?>): Int{
+    var c = 0
+    <caret>for (d in list) {
+        if (d == "") continue
+        if (d != null) {
+            if (d is Int) continue
+            c++
+        }
+    }
+    return c
+}

idea/testData/intentions/loopToCallChain/count/countSomethingAndNotNull.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'count{}'"
+// IS_APPLICABLE_2: false
+fun f(list: List<Any?>): Int{
+    val <caret>c = list.count { it != "" && it != null && it !is Int }
+    return c
+}

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

@@ -83,6 +83,12 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
         doTest(fileName);
     }
 
+    @TestMetadata("KT14191.kt")
+    public void testKT14191() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/KT14191.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("lastOrNull_ifAssign.kt")
     public void testLastOrNull_ifAssign() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/lastOrNull_ifAssign.kt");
@@ -235,9 +241,21 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/loopToCallChain/count"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }
 
-        @TestMetadata("KT14191.kt")
+        @TestMetadata("countIsInstance.kt")
-        public void testKT14191() throws Exception {
+        public void testCountIsInstance() throws Exception {
-            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/KT14191.kt");
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/countIsInstance.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("countNotNull.kt")
+        public void testCountNotNull() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/countNotNull.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("countSomethingAndNotNull.kt")
+        public void testCountSomethingAndNotNull() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/countSomethingAndNotNull.kt");
             doTest(fileName);
         }
 

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

@@ -7991,6 +7991,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("KT14191.kt")
+        public void testKT14191() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/KT14191.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("lastOrNull_ifAssign.kt")
         public void testLastOrNull_ifAssign() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/lastOrNull_ifAssign.kt");
@@ -8143,9 +8149,21 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/loopToCallChain/count"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
             }
 
-            @TestMetadata("KT14191.kt")
+            @TestMetadata("countIsInstance.kt")
-            public void testKT14191() throws Exception {
+            public void testCountIsInstance() throws Exception {
-                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/KT14191.kt");
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/countIsInstance.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("countNotNull.kt")
+            public void testCountNotNull() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/countNotNull.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("countSomethingAndNotNull.kt")
+            public void testCountSomethingAndNotNull() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/count/countSomethingAndNotNull.kt");
                 doTest(fileName);
             }