Initial support for "+=", filterTo and mapTo for collections

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

@@ -25,6 +25,8 @@ import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
  * An abstraction for generating a chained call that knows about receiver expression and handles proper formatting
  */
 interface ChainedCallGenerator {
+    val receiver: KtExpression
+
     /**
      * @param pattern pattern string for generating the part of the call to the right from the dot
      */
@@ -56,7 +58,7 @@ interface ResultTransformation : Transformation {
     fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? = null
 
     val commentSavingRange: PsiChildRange
-    val commentRestoringRange: PsiChildRange
+    fun commentRestoringRange(convertLoopResult: KtExpression): PsiChildRange
 
     fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression
 

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

@@ -20,6 +20,7 @@ import com.intellij.openapi.util.Key
 import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.result.AddToCollectionTransformation
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.result.FindAndAssignTransformation
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.result.FindAndReturnTransformation
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
@@ -46,7 +47,8 @@ object MatcherRegistrar {
 
     val resultMatchers: Collection<ResultTransformationMatcher> = listOf(
             FindAndReturnTransformation.Matcher,
-            FindAndAssignTransformation.Matcher
+            FindAndAssignTransformation.Matcher,
+            AddToCollectionTransformation.Matcher
     )
 }
 
@@ -98,7 +100,7 @@ fun convertLoop(loop: KtForExpression, matchResult: ResultTransformationMatch):
 
     val result = matchResult.resultTransformation.convertLoop(callChain)
 
-    commentSaver.restore(matchResult.resultTransformation.commentRestoringRange)
+    commentSaver.restore(matchResult.resultTransformation.commentRestoringRange(result))
 
     return result
 }
@@ -169,6 +171,9 @@ private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression):
 
     val psiFactory = KtPsiFactory(loop)
     val chainedCallGenerator = object : ChainedCallGenerator {
+        override val receiver: KtExpression
+            get() = callChain
+
         override fun generate(pattern: String, vararg args: Any): KtExpression {
             val newPattern = "\$${args.size}$lineBreak.$pattern"
             return psiFactory.createExpressionByPattern(newPattern, *args, callChain)

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

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.intentions.loopToCallChain.result
+
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getCallNameExpression
+
+class AddToCollectionTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        private val targetCollection: KtExpression
+) : ReplaceLoopTransformation(loop, inputVariable) {
+
+    override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
+        if (previousTransformation !is FilterTransformation) return null
+        return FilterToTransformation(loop, inputVariable, targetCollection, previousTransformation.effectiveCondition()) //TODO: use filterNotTo?
+    }
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        return KtPsiFactory(loop).createExpressionByPattern("$0 += $1", targetCollection, chainedCallGenerator.receiver)
+    }
+
+    /**
+     * Matches:
+     *     for (...) {
+     *         ...
+     *         collection.add(...)
+     *     }
+     */
+    object Matcher : ResultTransformationMatcher {
+        override fun match(state: MatchingState): ResultTransformationMatch? {
+            //TODO: pass indexVariable as null if not used
+            if (state.indexVariable != null) return null
+
+            val statement = state.statements.singleOrNull() ?: return null
+            //TODO: it can be implicit 'this' too
+            val qualifiedExpression = statement as? KtDotQualifiedExpression ?: return null
+            val targetCollection = qualifiedExpression.receiverExpression
+            //TODO: check that receiver is stable!
+            val callExpression = qualifiedExpression.selectorExpression as? KtCallExpression ?: return null
+            if (callExpression.getCallNameExpression()?.getReferencedName() != "add") return null
+            //TODO: check that it's MutableCollection's add
+            val argument = callExpression.valueArguments.singleOrNull() ?: return null
+            val argumentValue = argument.getArgumentExpression() ?: return null
+
+            //TODO: if variable is initialized with new collection than generate toList(), toSet()
+
+            val transformation = if (argumentValue.isVariableReference(state.workingVariable)) {
+                AddToCollectionTransformation(state.outerLoop, state.workingVariable, targetCollection)
+            }
+            else {
+                MapToTransformation(state.outerLoop, state.workingVariable, targetCollection, argumentValue)
+            }
+            return ResultTransformationMatch(transformation)
+        }
+    }
+}
+
+class FilterToTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        private val targetCollection: KtExpression,
+        private val filter: KtExpression
+) : ReplaceLoopTransformation(loop, inputVariable) {
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(inputVariable, filter)
+        return chainedCallGenerator.generate("filterTo($0) $1:'{}'", targetCollection, lambda)
+    }
+}
+
+class MapToTransformation(
+        loop: KtForExpression,
+        inputVariable: KtCallableDeclaration,
+        private val targetCollection: KtExpression,
+        private val mapping: KtExpression
+) : ReplaceLoopTransformation(loop, inputVariable) {
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(inputVariable, mapping)
+        return chainedCallGenerator.generate("mapTo($0) $1:'{}'", targetCollection, lambda)
+    }
+}

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

@@ -38,11 +38,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, generator, initialization, previousTransformation.buildRealCondition())
+        return FindAndAssignTransformation(loop, previousTransformation.inputVariable, generator, initialization, previousTransformation.effectiveCondition())
     }
 
     override val commentSavingRange = PsiChildRange(initialization.initializationStatement, loop.unwrapIfLabeled())
-    override val commentRestoringRange = commentSavingRange.withoutLastStatement()
+
+    private val commentRestoringRange = commentSavingRange.withoutLastStatement()
+
+    override fun commentRestoringRange(convertLoopResult: KtExpression) = commentRestoringRange
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         return generator(chainedCallGenerator, filter)

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

@@ -35,12 +35,14 @@ 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, generator, endReturn, previousTransformation.buildRealCondition())
+        return FindAndReturnTransformation(loop, previousTransformation.inputVariable, generator, endReturn, previousTransformation.effectiveCondition())
     }
 
     override val commentSavingRange = PsiChildRange(loop.unwrapIfLabeled(), endReturn)
 
-    override val commentRestoringRange = commentSavingRange.withoutFirstStatement()
+    private val commentRestoringRange = commentSavingRange.withoutFirstStatement()
+
+    override fun commentRestoringRange(convertLoopResult: KtExpression) = commentRestoringRange
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         return generator(chainedCallGenerator, filter)

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

@@ -29,13 +29,13 @@ class FilterTransformation(
         val isInverse: Boolean
 ) : SequenceTransformation {
 
-    fun buildRealCondition() = if (isInverse) condition.negate() else condition
+    fun effectiveCondition() = if (isInverse) condition.negate() else condition
 
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? {
         if (previousTransformation !is FilterTransformation) return null
         assert(previousTransformation.inputVariable == inputVariable)
         val mergedCondition = KtPsiFactory(condition).createExpressionByPattern(
-                "$0 && $1", previousTransformation.buildRealCondition(), buildRealCondition())
+                "$0 && $1", previousTransformation.effectiveCondition(), effectiveCondition())
         return FilterTransformation(inputVariable, mergedCondition, isInverse = false) //TODO: build filterNot in some cases?
     }
 
@@ -92,22 +92,22 @@ class FilterTransformation(
                 condition: KtExpression,
                 isInverse: Boolean): SequenceTransformation {
 
-            val realCondition = if (isInverse) condition.negate() else condition
+            val effectiveCondition = if (isInverse) condition.negate() else condition
 
-            if (realCondition is KtIsExpression
-                && !realCondition.isNegated
-                && realCondition.leftHandSide.isSimpleName(inputVariable.nameAsSafeName) // we cannot use isVariableReference here because expression can be non-physical
+            if (effectiveCondition is KtIsExpression
+                && !effectiveCondition.isNegated
+                && effectiveCondition.leftHandSide.isSimpleName(inputVariable.nameAsSafeName) // we cannot use isVariableReference here because expression can be non-physical
             ) {
-                val typeRef = realCondition.typeReference
+                val typeRef = effectiveCondition.typeReference
                 if (typeRef != null) {
                     return FilterIsInstanceTransformation(inputVariable, typeRef)
                 }
             }
 
-            if (realCondition is KtBinaryExpression
-                && realCondition.operationToken == KtTokens.EXCLEQ
-                && realCondition.right.isNullExpression()
-                && realCondition.left.isSimpleName(inputVariable.nameAsSafeName)
+            if (effectiveCondition is KtBinaryExpression
+                && effectiveCondition.operationToken == KtTokens.EXCLEQ
+                && effectiveCondition.right.isNullExpression()
+                && effectiveCondition.left.isSimpleName(inputVariable.nameAsSafeName)
             ) {
                 return FilterNotNullTransformation(inputVariable)
             }

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

@@ -36,7 +36,6 @@ 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.types.typeUtil.TypeNullability
 import org.jetbrains.kotlin.types.typeUtil.nullability
@@ -257,3 +256,17 @@ fun KtExpression.detectInitializationBeforeLoop(loop: KtForExpression): Variable
     val initializer = assignment.right ?: return null
     return VariableInitialization(variable, assignment, initializer)
 }
+
+abstract class ReplaceLoopTransformation(
+        protected val loop: KtForExpression,
+        override val inputVariable: KtCallableDeclaration
+): ResultTransformation {
+
+    override val commentSavingRange = PsiChildRange.singleElement(loop.unwrapIfLabeled())
+
+    override fun commentRestoringRange(convertLoopResult: KtExpression) = PsiChildRange.singleElement(convertLoopResult)
+
+    override fun convertLoop(resultCallChain: KtExpression): KtExpression {
+        return loop.unwrapIfLabeled().replaced(resultCallChain)
+    }
+}

idea/testData/intentions/loopToCallChain/addToCollection.kt

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        target.add(s)
+    }
+}

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

@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>target += list
+}

idea/testData/intentions/loopToCallChain/filterTo.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        if (s.length > 0)
+            target.add(s)
+    }
+}

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

@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<String>) {
+    list.filterTo(target) { it.length > 0 }
+}

idea/testData/intentions/loopToCallChain/mapTo.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<Int>) {
+    <caret>for (s in list) {
+        if (s.length > 0)
+            target.add(s.hashCode())
+    }
+}

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

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun foo(list: List<String>, target: MutableList<Int>) {
+    <caret>list
+            .filter { it.length > 0 }
+            .mapTo(target) { it.hashCode() }
+}