Basic support for "add" to collection used inside the loop

This commit is contained in:
Valentin Kipyatkov
2016-04-06 19:20:08 +03:00
parent d3721e9462
commit f51c5a19dd
20 changed files with 272 additions and 46 deletions
@@ -38,6 +38,9 @@ interface ChainedCallGenerator {
*/
interface Transformation {
val inputVariable: KtCallableDeclaration
val loop: KtForExpression
fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression
}
/**
@@ -48,7 +51,6 @@ interface SequenceTransformation : Transformation {
val affectsIndex: Boolean
fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression
}
/**
@@ -60,8 +62,6 @@ interface ResultTransformation : Transformation {
val commentSavingRange: PsiChildRange
fun commentRestoringRange(convertLoopResult: KtExpression): PsiChildRange
fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression
fun convertLoop(resultCallChain: KtExpression): KtExpression
}
@@ -104,5 +104,8 @@ interface ResultTransformationMatcher {
class ResultTransformationMatch(
val resultTransformation: ResultTransformation,
val sequenceTransformations: List<SequenceTransformation> = listOf()
)
val sequenceTransformations: List<SequenceTransformation>
) {
constructor(resultTransformation: ResultTransformation, vararg sequenceTransformations: SequenceTransformation)
: this(resultTransformation, sequenceTransformations.asList())
}
@@ -19,6 +19,7 @@ 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.idea.intentions.loopToCallChain.sequence.FlatMapTransformation
import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.MapTransformation
import org.jetbrains.kotlin.psi.*
import org.jetbrains.kotlin.psi.psiUtil.getCallNameExpression
@@ -26,7 +27,7 @@ class AddToCollectionTransformation(
loop: KtForExpression,
inputVariable: KtCallableDeclaration,
private val targetCollection: KtExpression
) : ReplaceLoopTransformation(loop, inputVariable) {
) : ReplaceLoopResultTransformation(loop, inputVariable) {
override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
return when (previousTransformation) {
@@ -68,10 +69,42 @@ class AddToCollectionTransformation(
//TODO: check that it's MutableCollection's add
val argument = callExpression.valueArguments.singleOrNull() ?: return null
val argumentValue = argument.getArgumentExpression() ?: return null
val argumentIsInputVariable = argumentValue.isVariableReference(state.inputVariable)
//TODO: if variable is initialized with new collection than generate toList(), toSet()
//TODO: collection can be used as mutable collection or even ArrayList!
val collectionInitialization = targetCollection.detectInitializationBeforeLoop(state.outerLoop)
if (collectionInitialization != null) {
val collectionKind = collectionInitialization.initializer.isSimpleCollectionInstantiation()
when (collectionKind) {
CollectionKind.LIST -> {
val transformation = if (argumentIsInputVariable) {
AssignToListTransformation(state.outerLoop, state.inputVariable, collectionInitialization)
}
else {
val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, argumentValue)
AssignSequenceTransformationResultTransformation(mapTransformation, collectionInitialization)
}
return ResultTransformationMatch(transformation)
}
val transformation = if (argumentValue.isVariableReference(state.inputVariable)) {
CollectionKind.SET -> {
if (argumentIsInputVariable) {
val transformation = AssignToSetTransformation(state.outerLoop, state.inputVariable, collectionInitialization)
return ResultTransformationMatch(transformation)
}
else {
val mapTransformation = MapTransformation(state.outerLoop, state.inputVariable, argumentValue)
val transformation = AssignToSetTransformation(
state.outerLoop,
state.inputVariable/*TODO: it's not correct and it looks like not all transformations should have inputVariable*/,
collectionInitialization)
return ResultTransformationMatch(transformation, mapTransformation)
}
}
}
}
val transformation = if (argumentIsInputVariable) {
AddToCollectionTransformation(state.outerLoop, state.inputVariable, targetCollection)
}
else {
@@ -87,7 +120,7 @@ class FilterToTransformation(
inputVariable: KtCallableDeclaration,
private val targetCollection: KtExpression,
private val filter: KtExpression
) : ReplaceLoopTransformation(loop, inputVariable) {
) : ReplaceLoopResultTransformation(loop, inputVariable) {
override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
val lambda = generateLambda(inputVariable, filter)
@@ -100,7 +133,7 @@ class MapToTransformation(
inputVariable: KtCallableDeclaration,
private val targetCollection: KtExpression,
private val mapping: KtExpression
) : ReplaceLoopTransformation(loop, inputVariable) {
) : ReplaceLoopResultTransformation(loop, inputVariable) {
override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
val lambda = generateLambda(inputVariable, mapping)
@@ -113,10 +146,37 @@ class FlatMapToTransformation(
inputVariable: KtCallableDeclaration,
private val targetCollection: KtExpression,
private val transform: KtExpression
) : ReplaceLoopTransformation(loop, inputVariable) {
) : ReplaceLoopResultTransformation(loop, inputVariable) {
override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
val lambda = generateLambda(inputVariable, transform)
return chainedCallGenerator.generate("flatMapTo($0) $1:'{}'", targetCollection, lambda)
}
}
}
class AssignToListTransformation(
loop: KtForExpression,
inputVariable: KtCallableDeclaration,
initialization: VariableInitialization
) : AssignToVariableResultTransformation(loop, inputVariable, initialization) {
override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
if (previousTransformation !is FilterTransformation) return null
return AssignSequenceTransformationResultTransformation(previousTransformation, initialization)
}
override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
return chainedCallGenerator.generate("toList()")
}
}
class AssignToSetTransformation(
loop: KtForExpression,
inputVariable: KtCallableDeclaration,
initialization: VariableInitialization
) : AssignToVariableResultTransformation(loop, inputVariable, initialization) {
override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
return chainedCallGenerator.generate("toSet()")
}
}
@@ -23,17 +23,16 @@ import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
import org.jetbrains.kotlin.lexer.KtTokens
import org.jetbrains.kotlin.psi.*
import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
class FindAndAssignTransformation(
private val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration,
loop: KtForExpression,
inputVariable: KtCallableDeclaration,
private val generator: (chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression,
private val initialization: VariableInitialization,
initialization: VariableInitialization,
private val filter: KtExpression? = null
) : ResultTransformation {
) : AssignToVariableResultTransformation(loop, inputVariable, initialization) {
override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
if (previousTransformation !is FilterTransformation) return null
@@ -41,27 +40,10 @@ class FindAndAssignTransformation(
return FindAndAssignTransformation(loop, previousTransformation.inputVariable, generator, initialization, previousTransformation.effectiveCondition())
}
override val commentSavingRange = PsiChildRange(initialization.initializationStatement, loop.unwrapIfLabeled())
private val commentRestoringRange = commentSavingRange.withoutLastStatement()
override fun commentRestoringRange(convertLoopResult: KtExpression) = commentRestoringRange
override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
return generator(chainedCallGenerator, filter)
}
override fun convertLoop(resultCallChain: KtExpression): KtExpression {
initialization.initializer.replace(resultCallChain)
loop.deleteWithLabels()
if (!initialization.variable.hasWriteUsages()) { // change variable to 'val' if possible
initialization.variable.valOrVarKeyword.replace(KtPsiFactory(initialization.variable).createValKeyword())
}
return initialization.initializationStatement
}
/**
* Matches:
* val variable = ...
@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.psi.KtReturnExpression
import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
class FindAndReturnTransformation(
private val loop: KtForExpression,
override val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration,
private val generator: (chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?) -> KtExpression,
private val endReturn: KtReturnExpression,
@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.psi.*
import org.jetbrains.kotlin.psi.psiUtil.blockExpressionsOrSingle
class FilterTransformation(
override val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration,
val condition: KtExpression,
val isInverse: Boolean
@@ -36,7 +37,7 @@ class FilterTransformation(
assert(previousTransformation.inputVariable == inputVariable)
val mergedCondition = KtPsiFactory(condition).createExpressionByPattern(
"$0 && $1", previousTransformation.effectiveCondition(), effectiveCondition())
return FilterTransformation(inputVariable, mergedCondition, isInverse = false) //TODO: build filterNot in some cases?
return FilterTransformation(loop, inputVariable, mergedCondition, isInverse = false) //TODO: build filterNot in some cases?
}
override val affectsIndex: Boolean
@@ -73,14 +74,14 @@ class FilterTransformation(
val then = ifStatement.then ?: return null
if (state.statements.size == 1) {
val transformation = createFilterTransformation(state.inputVariable, condition, isInverse = false)
val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, condition, isInverse = false)
val newState = state.copy(statements = listOf(then))
return SequenceTransformationMatch(transformation, newState)
}
else {
val continueExpression = then.blockExpressionsOrSingle().singleOrNull() as? KtContinueExpression ?: return null
if (!continueExpression.isBreakOrContinueOfLoop(state.innerLoop)) return null
val transformation = createFilterTransformation(state.inputVariable, condition, isInverse = true)
val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, condition, isInverse = true)
val newState = state.copy(statements = state.statements.drop(1))
return SequenceTransformationMatch(transformation, newState)
}
@@ -88,6 +89,7 @@ class FilterTransformation(
//TODO: choose filter or filterNot depending on condition
private fun createFilterTransformation(
loop: KtForExpression,
inputVariable: KtCallableDeclaration,
condition: KtExpression,
isInverse: Boolean): SequenceTransformation {
@@ -100,7 +102,7 @@ class FilterTransformation(
) {
val typeRef = effectiveCondition.typeReference
if (typeRef != null) {
return FilterIsInstanceTransformation(inputVariable, typeRef)
return FilterIsInstanceTransformation(loop, inputVariable, typeRef)
}
}
@@ -109,15 +111,16 @@ class FilterTransformation(
&& effectiveCondition.right.isNullExpression()
&& effectiveCondition.left.isSimpleName(inputVariable.nameAsSafeName)
) {
return FilterNotNullTransformation(inputVariable)
return FilterNotNullTransformation(loop, inputVariable)
}
return FilterTransformation(inputVariable, condition, isInverse)
return FilterTransformation(loop, inputVariable, condition, isInverse)
}
}
}
class FilterIsInstanceTransformation(
override val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration,
private val type: KtTypeReference
) : SequenceTransformation {
@@ -131,6 +134,7 @@ class FilterIsInstanceTransformation(
}
class FilterNotNullTransformation(
override val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration
) : SequenceTransformation {
@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.psi.KtForExpression
import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
class FlatMapTransformation(
override val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration,
val transform: KtExpression
) : SequenceTransformation {
@@ -63,7 +64,7 @@ class FlatMapTransformation(
val nestedLoopBody = nestedLoop.body ?: return null
val newWorkingVariable = nestedLoop.loopParameter ?: return null
val transformation = FlatMapTransformation(state.inputVariable, transform)
val transformation = FlatMapTransformation(state.outerLoop, state.inputVariable, transform)
val newState = state.copy(
innerLoop = nestedLoop,
statements = listOf(nestedLoopBody),
@@ -19,9 +19,11 @@ package org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence
import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
import org.jetbrains.kotlin.psi.KtCallableDeclaration
import org.jetbrains.kotlin.psi.KtExpression
import org.jetbrains.kotlin.psi.KtForExpression
import org.jetbrains.kotlin.psi.KtProperty
class MapTransformation(
override val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration,
val mapping: KtExpression
) : SequenceTransformation {
@@ -50,7 +52,7 @@ class MapTransformation(
if (declaration.hasWriteUsages()) return null
val restStatements = state.statements.drop(1)
val transformation = MapTransformation(state.inputVariable, initializer)
val transformation = MapTransformation(state.outerLoop, state.inputVariable, initializer)
val newState = state.copy(statements = restStatements, inputVariable = declaration)
return SequenceTransformationMatch(transformation, newState)
}
@@ -21,12 +21,14 @@ import com.intellij.psi.PsiWhiteSpace
import com.intellij.psi.search.LocalSearchScope
import com.intellij.psi.search.searches.ReferencesSearch
import org.jetbrains.kotlin.KtNodeTypes
import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
import org.jetbrains.kotlin.descriptors.VariableDescriptor
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.caches.resolve.resolveToDescriptor
import org.jetbrains.kotlin.idea.core.replaced
import org.jetbrains.kotlin.idea.imports.importableFqName
import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isNullExpression
import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
import org.jetbrains.kotlin.idea.references.mainReference
@@ -36,6 +38,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.calls.callUtil.getResolvedCall
import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
import org.jetbrains.kotlin.types.typeUtil.TypeNullability
import org.jetbrains.kotlin.types.typeUtil.nullability
@@ -246,6 +249,9 @@ fun KtExpression.detectInitializationBeforeLoop(loop: KtForExpression): Variable
val variable = this.mainReference.resolve() as? KtProperty ?: return null
val statementBeforeLoop = loop.previousStatement() //TODO: support initialization not right before the loop
// do not allow any other usages of this variable inside the loop
if (ReferencesSearch.search(variable, LocalSearchScope(loop)).count() > 1) return null
if (statementBeforeLoop == variable) {
val initializer = variable.initializer ?: return null
return VariableInitialization(variable, variable, initializer)
@@ -253,13 +259,14 @@ fun KtExpression.detectInitializationBeforeLoop(loop: KtForExpression): Variable
val assignment = statementBeforeLoop?.asAssignment() ?: return null
if (!assignment.left.isVariableReference(variable)) return null
val initializer = assignment.right ?: return null
return VariableInitialization(variable, assignment, initializer)
}
abstract class ReplaceLoopTransformation(
protected val loop: KtForExpression,
override val inputVariable: KtCallableDeclaration
abstract class ReplaceLoopResultTransformation(
override val loop: KtForExpression,
final override val inputVariable: KtCallableDeclaration
): ResultTransformation {
override val commentSavingRange = PsiChildRange.singleElement(loop.unwrapIfLabeled())
@@ -270,3 +277,56 @@ abstract class ReplaceLoopTransformation(
return loop.unwrapIfLabeled().replaced(resultCallChain)
}
}
abstract class AssignToVariableResultTransformation(
override val loop: KtForExpression,
final override val inputVariable: KtCallableDeclaration,
protected val initialization: VariableInitialization
) : ResultTransformation {
override val commentSavingRange = PsiChildRange(initialization.initializationStatement, loop.unwrapIfLabeled())
private val commentRestoringRange = commentSavingRange.withoutLastStatement()
override fun commentRestoringRange(convertLoopResult: KtExpression) = commentRestoringRange
override fun convertLoop(resultCallChain: KtExpression): KtExpression {
initialization.initializer.replace(resultCallChain)
loop.deleteWithLabels()
if (initialization.variable.isVar && !initialization.variable.hasWriteUsages()) { // change variable to 'val' if possible
initialization.variable.valOrVarKeyword.replace(KtPsiFactory(initialization.variable).createValKeyword())
}
return initialization.initializationStatement
}
}
class AssignSequenceTransformationResultTransformation(
private val sequenceTransformation: SequenceTransformation,
initialization: VariableInitialization
) : AssignToVariableResultTransformation(sequenceTransformation.loop, sequenceTransformation.inputVariable, initialization) {
override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
return sequenceTransformation.generateCode(chainedCallGenerator)
}
}
enum class CollectionKind {
LIST, SET/*, MAP*/
}
fun KtExpression.isSimpleCollectionInstantiation(): CollectionKind? {
val callExpression = this as? KtCallExpression ?: return null //TODO: it can be qualified too
if (callExpression.valueArguments.isNotEmpty()) return null
val bindingContext = callExpression.analyze(BodyResolveMode.PARTIAL)
val resolvedCall = callExpression.getResolvedCall(bindingContext) ?: return null
val constructorDescriptor = resolvedCall.resultingDescriptor as? ConstructorDescriptor ?: return null
val classDescriptor = constructorDescriptor.containingDeclaration
val classFqName = classDescriptor.importableFqName?.asString()
return when (classFqName) {
"java.util.ArrayList" -> CollectionKind.LIST
"java.util.HashSet", "java.util.LinkedHashSet" -> CollectionKind.SET
else -> null
}
}
@@ -0,0 +1,12 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(list: List<String>): List<String> {
val result = ArrayList<String>()
<caret>for (s in list) {
if (s.length > 0) {
result.add(s)
}
}
return result
}
@@ -0,0 +1,7 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(list: List<String>): List<String> {
<caret>val result = list.filter { it.length > 0 }
return result
}
+11
View File
@@ -0,0 +1,11 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(list: List<String>): List<Int> {
val result = ArrayList<Int>()
<caret>for (s in list) {
if (s.length > 0)
result.add(s.hashCode())
}
return result
}
@@ -0,0 +1,9 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(list: List<String>): List<Int> {
val result = list
.filter { it.length > 0 }
.map { it.hashCode() }
return result
}
@@ -0,0 +1,12 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(list: List<String>): List<Int> {
val result = ArrayList<Int>()
<caret>for (s in list) {
if (s.length > result.size) {
result.add(s.hashCode())
}
}
return result
}
@@ -0,0 +1,10 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(list: List<String>): List<Int> {
val result = ArrayList<Int>()
<caret>list
.filter { it.length > result.size }
.mapTo(result) { it.hashCode() }
return result
}
+10
View File
@@ -0,0 +1,10 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(map: Map<Int, String>): List<String> {
val result = ArrayList<String>()
<caret>for (s in map.values) {
result.add(s)
}
return result
}
@@ -0,0 +1,7 @@
// WITH_RUNTIME
import java.util.ArrayList
fun foo(map: Map<Int, String>): List<String> {
<caret>val result = map.values.toList()
return result
}
+10
View File
@@ -0,0 +1,10 @@
// WITH_RUNTIME
import java.util.HashSet
fun foo(map: Map<Int, String>): Collection<String> {
val result = HashSet<String>()
<caret>for (s in map.values) {
result.add(s)
}
return result
}
@@ -0,0 +1,7 @@
// WITH_RUNTIME
import java.util.HashSet
fun foo(map: Map<Int, String>): Collection<String> {
<caret>val result = map.values.toSet()
return result
}
@@ -0,0 +1,10 @@
// WITH_RUNTIME
import java.util.HashSet
fun foo(map: Map<Int, String>): Collection<Int> {
val result = HashSet<Int>()
<caret>for (s in map.values) {
result.add(s.length)
}
return result
}
@@ -0,0 +1,9 @@
// WITH_RUNTIME
import java.util.HashSet
fun foo(map: Map<Int, String>): Collection<Int> {
<caret>val result = map.values
.map { it.length }
.toSet()
return result
}