diff --git a/java/java-impl/src/com/intellij/codeInspection/java18api/Java8MapApiInspection.java b/java/java-impl/src/com/intellij/codeInspection/java18api/Java8MapApiInspection.java index b5b14a15cc0f..8bd6bda41d16 100644 --- a/java/java-impl/src/com/intellij/codeInspection/java18api/Java8MapApiInspection.java +++ b/java/java-impl/src/com/intellij/codeInspection/java18api/Java8MapApiInspection.java @@ -14,6 +14,9 @@ import com.intellij.psi.codeStyle.CodeStyleManager; import com.intellij.psi.codeStyle.JavaCodeStyleManager; import com.intellij.psi.search.searches.ReferencesSearch; import com.intellij.psi.util.PsiTreeUtil; +import com.intellij.util.ArrayUtil; +import com.intellij.util.ObjectUtils; +import com.siyeh.ig.PsiReplacementUtil; import com.siyeh.ig.psiutils.*; import one.util.streamex.StreamEx; import org.jetbrains.annotations.Nls; @@ -21,6 +24,7 @@ import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import javax.swing.*; +import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.Locale; @@ -93,6 +97,30 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool { } } + @Override + public void visitForeachStatement(PsiForeachStatement statement) { + MapLoopCondition condition = MapLoopCondition.create(statement); + if (condition == null) return; + PsiMethodCallExpression putCall = condition.extractPut(statement); + if (putCall == null) return; + PsiExpression[] args = putCall.getArgumentList().getExpressions(); + if (args.length != 2) return; + PsiExpression key = args[0]; + if (!condition.isKeyAccess(key)) return; + PsiExpression value = args[1]; + if (condition.isEntrySet() && isUsedAsReference(value, condition)) return; + + ReplaceWithSingleMapOperation fix = ReplaceWithSingleMapOperation.create("replaceAll", putCall, value); + holder.registerProblem(statement.getFirstChild(), QuickFixBundle.message("java.8.map.api.inspection.description", fix.myMethodName), + fix); + } + + private boolean isUsedAsReference(@NotNull PsiElement value, @NotNull MapLoopCondition condition) { + VariableUsedAsReferenceVisitor visitor = new VariableUsedAsReferenceVisitor(condition.getIterParam(), "getValue", "getKey"); + value.accept(visitor); + return visitor.isUsed; + } + private void processMerge(MapCheckCondition condition, PsiStatement existsBranch, PsiStatement noneBranch) { @@ -221,6 +249,44 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool { return nameCandidate.toLowerCase(Locale.ENGLISH); } + private static class VariableUsedAsReferenceVisitor extends JavaRecursiveElementWalkingVisitor { + private final PsiVariable myVariable; + private final String[] myExcludeMethodNames; + + private boolean isUsed = false; + + private VariableUsedAsReferenceVisitor(PsiVariable variable, String... excludeMethodNames) { + myVariable = variable; + myExcludeMethodNames = excludeMethodNames; + } + + @Override + public void visitReferenceExpression(PsiReferenceExpression expression) { + if (isUsed) { + return; + } + super.visitReferenceExpression(expression); + if (expression.resolve() == myVariable) { + isUsed = true; + } + } + + @Override + public void visitMethodCallExpression(PsiMethodCallExpression expression) { + if (isUsed) { + return; + } + if (ArrayUtil.contains(expression.getMethodExpression().getReferenceName(), myExcludeMethodNames)) { + PsiReferenceExpression qualifier = + ObjectUtils.tryCast(expression.getMethodExpression().getQualifierExpression(), PsiReferenceExpression.class); + if (qualifier != null && qualifier.resolve() == myVariable) { + return; + } + } + super.visitMethodCallExpression(expression); + } + } + private static class ReplaceWithSingleMapOperation implements LocalQuickFix { private final String myMethodName; private final SmartPsiElementPointer myCallPointer; @@ -237,16 +303,19 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool { @Override public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) { - PsiElement conditional = PsiTreeUtil.getParentOfType(descriptor.getStartElement(), PsiIfStatement.class, PsiConditionalExpression.class); - if(conditional == null) return; - MapCheckCondition condition = fromConditional(conditional, true); + PsiElement outerElement = PsiTreeUtil.getParentOfType(descriptor.getStartElement(), PsiIfStatement.class, + PsiConditionalExpression.class, PsiForeachStatement.class); + if (outerElement == null) return; + MapCondition condition = outerElement instanceof PsiForeachStatement ? + MapLoopCondition.create((PsiForeachStatement)outerElement) : + fromConditional(outerElement, true); if(condition == null) return; PsiMethodCallExpression call = myCallPointer.getElement(); if (call == null) return; PsiExpressionList argsList = call.getArgumentList(); PsiExpression[] args = argsList.getExpressions(); if(args.length == 0) return; - if(myMethodName.equals("merge") && args.length != 2) return; + if ((myMethodName.equals("merge") || myMethodName.equals("replaceAll")) && args.length != 2) return; PsiExpression value = myValuePointer.getElement(); if (value == null) return; PsiElement result = myResultPointer.getElement(); @@ -285,21 +354,34 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool { } replacement = factory.createExpressionFromText(varName + " -> " + ct.text(value), value); } else if (myMethodName.equals("merge")) { + MapCheckCondition checkCondition = ObjectUtils.tryCast(condition, MapCheckCondition.class); + if (checkCondition == null) return; PsiExpression absentValue = args[1]; String aVar = JavaCodeStyleManager.getInstance(project).suggestUniqueVariableName("a", value, true); String bVar = JavaCodeStyleManager.getInstance(project).suggestUniqueVariableName("b", value, true); for(PsiElement e : PsiTreeUtil.collectElements(value, e -> PsiEquivalenceUtil.areElementsEquivalent(e, absentValue))) { ct.replace(e, factory.createIdentifier(bVar)); } - for(PsiElement e : PsiTreeUtil.collectElements(value, e -> condition.extractGetCall(e) != null || condition.isValueReference(e))) { + for (PsiElement e : PsiTreeUtil + .collectElements(value, e -> checkCondition.extractGetCall(e) != null || checkCondition.isValueReference(e))) { ct.replace(e, factory.createIdentifier(aVar)); } replacement = factory.createExpressionFromText("("+aVar+","+bVar+") -> "+ct.text(value), value); - } else { + } + else if (myMethodName.equals("replaceAll")) { + MapLoopCondition loopCondition = ObjectUtils.tryCast(condition, MapLoopCondition.class); + if (loopCondition == null) return; + JavaCodeStyleManager styleManager = JavaCodeStyleManager.getInstance(project); + String kVar = suggestKeyName(styleManager, loopCondition, value); + String vVar = styleManager.suggestUniqueVariableName("v", value, true); + replacement = createLambdaForLoopReplacement(factory, kVar, vVar, loopCondition, value, ct); + ct.delete(args); + } + else { replacement = ct.markUnchanged(value); } PsiElement newArg; - if(args.length == 2 && !myMethodName.equals("merge")) { + if (args.length == 2 && !myMethodName.equals("merge") && !myMethodName.equals("replaceAll")) { newArg = ct.replace(args[1], replacement); } else { newArg = argsList.add(replacement); @@ -307,18 +389,65 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool { if(newArg instanceof PsiLambdaExpression) { LambdaCanBeMethodReferenceInspection.replaceLambdaWithMethodReference((PsiLambdaExpression)newArg); } - if(PsiTreeUtil.isAncestor(conditional, result, true)) { - result = ct.replaceAndRestoreComments(conditional, result); + if (PsiTreeUtil.isAncestor(outerElement, result, true)) { + result = ct.replaceAndRestoreComments(outerElement, result); } else { - ct.deleteAndRestoreComments(conditional); + ct.deleteAndRestoreComments(outerElement); } - PsiVariable variable = condition.extractDeclaration(); + PsiVariable variable = condition instanceof MapCheckCondition ? ((MapCheckCondition)condition).extractDeclaration() : null; if (variable != null && !PsiTreeUtil.isAncestor(result, variable, true) && ReferencesSearch.search(variable).findFirst() == null) { new CommentTracker().deleteAndRestoreComments(variable); } CodeStyleManager.getInstance(project).reformat(result); } + @NotNull + private static String suggestKeyName(@NotNull JavaCodeStyleManager codeStyleManager, + @NotNull MapLoopCondition loopCondition, + @Nullable PsiElement value) { + String nameCandidate = "k"; + if (!loopCondition.isEntrySet()) { + String origName = loopCondition.getIterParam().getName(); + if (origName != null) { + nameCandidate = getNameCandidate(origName); + if (origName.equals(nameCandidate)) return nameCandidate; + } + } + return codeStyleManager.suggestUniqueVariableName(nameCandidate, value, true); + } + + @NotNull + private static PsiExpression createLambdaForLoopReplacement(@NotNull PsiElementFactory factory, + @NotNull String kVar, + @NotNull String vVar, + @NotNull MapLoopCondition loopCondition, + @NotNull PsiExpression value, + @NotNull CommentTracker tracker) { + if (!loopCondition.isEntrySet()) { + PsiParameter param = loopCondition.getIterParam(); + StreamEx.ofTree((PsiElement)value, e -> StreamEx.of(e.getChildren())) + .select(PsiReferenceExpression.class) + .filter(ref -> param.equals(ref.resolve())) + .forEach(ref -> ExpressionUtils.bindReferenceTo(ref, kVar)); + } + else { + if (value instanceof PsiMethodCallExpression) { + if (loopCondition.isKeyAccess(value)) return factory.createExpressionFromText("(" + kVar + "," + vVar + ") ->" + kVar, value); + if (loopCondition.isValueAccess(value)) return factory.createExpressionFromText("(" + kVar + "," + vVar + ") ->" + vVar, value); + } + Collection calls = PsiTreeUtil.collectElementsOfType(value, PsiMethodCallExpression.class); + for (PsiMethodCallExpression call : calls) { + if (loopCondition.isKeyAccess(call)) { + PsiReplacementUtil.replaceExpression(call, kVar, new CommentTracker()); + } + else if (loopCondition.isValueAccess(call)) { + PsiReplacementUtil.replaceExpression(call, vVar, new CommentTracker()); + } + } + } + return factory.createExpressionFromText("(" + kVar + "," + vVar + ") ->" + tracker.text(value), value); + } + @Nls @NotNull @Override @@ -336,6 +465,11 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool { @NotNull static ReplaceWithSingleMapOperation fromIf(String methodName, MapCheckCondition condition, PsiExpression value) { PsiMethodCallExpression call = condition.getCheckCall(); + return create(methodName, call, value); + } + + @NotNull + static ReplaceWithSingleMapOperation create(String methodName, PsiMethodCallExpression call, PsiExpression value) { PsiStatement result = PsiTreeUtil.getParentOfType(call, PsiStatement.class); LOG.assertTrue(result != null); return new ReplaceWithSingleMapOperation(methodName, call, value, result); diff --git a/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetPoly.java b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetPoly.java new file mode 100644 index 000000000000..920bc2801c14 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetPoly.java @@ -0,0 +1,13 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +class Main { + public void test() { + Map map = new HashMap<>(); + String defaultValue = "42"; + map.put("foo", "bar"); + map.replaceAll((k, v) -> v + k + defaultValue + "baz"); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetShadowVar.java b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetShadowVar.java new file mode 100644 index 000000000000..c0c795524c62 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetShadowVar.java @@ -0,0 +1,13 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + String k = "another var"; + Map map = new HashMap<>(); + map.put("foo", "bar"); + map.replaceAll((k1, v) -> k); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetSimple.java b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetSimple.java new file mode 100644 index 000000000000..3a7495e4d6c5 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetSimple.java @@ -0,0 +1,12 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.Map; + +public class Main { + public static void main(String[] args) { + Map vals = new HashMap<>(); + vals.put("foo", "bar"); + String defaultValue = "42"; + vals.replaceAll((k, v) -> defaultValue); + } +} diff --git a/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetValueAsValue.java b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetValueAsValue.java new file mode 100644 index 000000000000..5bbfe7d1aa99 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllEntrySetValueAsValue.java @@ -0,0 +1,12 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + Map map = new HashMap<>(); + map.put("foo", "bar"); + map.replaceAll((k, v) -> v); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetKeyAsValue.java b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetKeyAsValue.java new file mode 100644 index 000000000000..ca81bf14d1d1 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetKeyAsValue.java @@ -0,0 +1,12 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + Map map = new HashMap<>(); + map.put("foo", "bar"); + map.replaceAll((k, v) -> k); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetPoly.java b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetPoly.java new file mode 100644 index 000000000000..b3c9dbe24a48 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetPoly.java @@ -0,0 +1,13 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +class Main { + public void test() { + Map map = new HashMap<>(); + String defaultValue = "42"; + map.put("foo", "bar"); + map.replaceAll((k, v) -> k + defaultValue + "baz"); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetSimple.java b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetSimple.java new file mode 100644 index 000000000000..2ee73f2e2e8f --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/afterReplaceAllKeySetSimple.java @@ -0,0 +1,13 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public static void main(String[] args) { + Map vals = new HashMap<>(); + vals.put("foo", "bar"); + String defaultValue = "42"; + vals.replaceAll((k, v) -> defaultValue); + } +} diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetEntryAsValue.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetEntryAsValue.java new file mode 100644 index 000000000000..a88ba34fc67c --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetEntryAsValue.java @@ -0,0 +1,14 @@ +// "Replace with 'replaceAll' method call" "false" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + Map map = new HashMap<>(); + map.put("foo", "bar"); + for (Map.Entry entry : map.entrySet()) { + map.put(entry.getKey(), String.valueOf(entry)); + } + } +} diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetEntryCallAsValue.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetEntryCallAsValue.java new file mode 100644 index 000000000000..608b461a2026 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetEntryCallAsValue.java @@ -0,0 +1,14 @@ +// "Replace with 'replaceAll' method call" "false" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + Map map = new HashMap<>(); + map.put("foo", "bar"); + for (Map.Entry entry : map.entrySet()) { + map.put(entry.getKey(), entry.toString()); + } + } +} diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetPoly.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetPoly.java new file mode 100644 index 000000000000..33edfc2abf0f --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetPoly.java @@ -0,0 +1,15 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +class Main { + public void test() { + Map map = new HashMap<>(); + String defaultValue = "42"; + map.put("foo", "bar"); + for (Map.Entry entry : map.entrySet()) { + map.put(entry.getKey(), entry.getValue() + entry.getKey() + defaultValue + "baz"); + } + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetShadowVar.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetShadowVar.java new file mode 100644 index 000000000000..440e26c16ac2 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetShadowVar.java @@ -0,0 +1,15 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + String k = "another var"; + Map map = new HashMap<>(); + map.put("foo", "bar"); + for (Map.Entry entry : map.entrySet()) { + map.put(entry.getKey(), k); + } + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetSimple.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetSimple.java new file mode 100644 index 000000000000..c9fe6a369329 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetSimple.java @@ -0,0 +1,14 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.Map; + +public class Main { + public static void main(String[] args) { + Map vals = new HashMap<>(); + vals.put("foo", "bar"); + String defaultValue = "42"; + for (Entry entry : vals.entrySet()) { + vals.put(entry.getKey(), defaultValue); + } + } +} diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetValueAsValue.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetValueAsValue.java new file mode 100644 index 000000000000..088e17300e75 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllEntrySetValueAsValue.java @@ -0,0 +1,14 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + Map map = new HashMap<>(); + map.put("foo", "bar"); + for (Map.Entry entry : map.entrySet()) { + map.put(entry.getKey(), entry.getValue()); + } + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetKeyAsValue.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetKeyAsValue.java new file mode 100644 index 000000000000..996fa506b5c1 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetKeyAsValue.java @@ -0,0 +1,14 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public void test() { + Map map = new HashMap<>(); + map.put("foo", "bar"); + for (String key : map.keySet()) { + map.put(key, key); + } + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetPoly.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetPoly.java new file mode 100644 index 000000000000..5cc73ff31ab9 --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetPoly.java @@ -0,0 +1,15 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +class Main { + public void test() { + Map map = new HashMap<>(); + String defaultValue = "42"; + map.put("foo", "bar"); + for (String key : map.keySet()) { + map.put(key, key + defaultValue + "baz"); + } + } +} \ No newline at end of file diff --git a/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetSimple.java b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetSimple.java new file mode 100644 index 000000000000..399b9bfa7e2a --- /dev/null +++ b/java/java-tests/testData/inspection/java8MapApi/beforeReplaceAllKeySetSimple.java @@ -0,0 +1,15 @@ +// "Replace with 'replaceAll' method call" "true" + +import java.util.HashMap; +import java.util.Map; + +public class Main { + public static void main(String[] args) { + Map vals = new HashMap<>(); + vals.put("foo", "bar"); + String defaultValue = "42"; + for (String key : vals.keySet()) { + vals.put(key, defaultValue); + } + } +} diff --git a/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/Java8MigrationUtils.java b/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/Java8MigrationUtils.java index 5abec58b7b99..9434e3000b66 100644 --- a/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/Java8MigrationUtils.java +++ b/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/Java8MigrationUtils.java @@ -21,11 +21,14 @@ import com.intellij.psi.*; import com.intellij.psi.search.searches.ReferencesSearch; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.psi.util.PsiUtil; +import com.intellij.util.ObjectUtils; import one.util.streamex.StreamEx; import org.jetbrains.annotations.Contract; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; +import java.util.Map; + public class Java8MigrationUtils { @Nullable private static Java8MigrationUtils.MapCheckCondition tryExtract(PsiExpression fullCondition, @@ -159,10 +162,13 @@ public class Java8MigrationUtils { } + public interface MapCondition { + } + /** * Class represents check when working with map: there is 2 ways - when there is a value that matches the key, and when value doesn't exists */ - public static class MapCheckCondition { + public static class MapCheckCondition implements MapCondition { private final @Nullable PsiReferenceExpression myValueReference; private final PsiExpression myMapExpression; private final PsiExpression myKeyExpression; @@ -281,4 +287,108 @@ public class Java8MigrationUtils { return null; } } + + /** + * This class is used to store enhanced for loop info when iterating {@link Map} keys or entries. + * It supports two iteration patterns: + * 1. for (String value : map.keySet()) + * 2. for (Entry entry : map.entrySet()) + */ + public static class MapLoopCondition implements MapCondition { + private final PsiParameter myIterParam; + private final boolean myIsEntrySet; + private final PsiReferenceExpression myMapExpression; + + private MapLoopCondition(@NotNull PsiParameter iterParam, boolean isEntrySet, @NotNull PsiReferenceExpression mapExpression) { + myIterParam = iterParam; + myIsEntrySet = isEntrySet; + myMapExpression = mapExpression; + } + + /** + * Try to extract {@link Map#put(Object, Object)} call from for each statement body. + * This call is expected to be applied to the same map instance as for each iteration, otherwise null will be returned. + * + * @see ControlFlowUtils#stripBraces(PsiStatement) + */ + @Nullable + public PsiMethodCallExpression extractPut(@NotNull PsiForeachStatement statement) { + PsiExpressionStatement putStatement = + ObjectUtils.tryCast(ControlFlowUtils.stripBraces(statement.getBody()), PsiExpressionStatement.class); + if (putStatement == null) return null; + PsiMethodCallExpression putCall = extractMapMethodCall(putStatement.getExpression(), "put"); + if (putCall == null || !isMap(putCall.getMethodExpression().getQualifierExpression())) return null; + return putCall; + } + + /** + * Check if given expression either entry.getKey() call (for entry set based loop) or + * loop parameter reference (for key based loop). + */ + public boolean isKeyAccess(@NotNull PsiExpression expression) { + if (!myIsEntrySet) { + PsiReferenceExpression keyRef = ObjectUtils.tryCast(expression, PsiReferenceExpression.class); + return keyRef != null && keyRef.resolve() == myIterParam; + } + return isParamCall(expression, "getKey"); + } + + /** + * Check if given expression is entry.getValue() call (for entry set based loop). + */ + public boolean isValueAccess(@NotNull PsiExpression expression) { + return myIsEntrySet && isParamCall(expression, "getValue"); + } + + /** + * Check if given call is invoked on loop iteration parameter. + */ + public boolean isParamCall(@NotNull PsiMethodCallExpression call) { + PsiReferenceExpression qualifier = + ObjectUtils.tryCast(call.getMethodExpression().getQualifierExpression(), PsiReferenceExpression.class); + return qualifier != null && qualifier.resolve() == myIterParam; + } + + public PsiParameter getIterParam() { + return myIterParam; + } + + public boolean isEntrySet() { + return myIsEntrySet; + } + + private boolean isParamCall(@NotNull PsiExpression expression, @NotNull String expectedName) { + PsiMethodCallExpression call = ObjectUtils.tryCast(expression, PsiMethodCallExpression.class); + if (call == null) return false; + String name = call.getMethodExpression().getReferenceName(); + return expectedName.equals(name) && isParamCall(call); + } + + private boolean isMap(@Nullable PsiElement element) { + return element != null && PsiEquivalenceUtil.areElementsEquivalent(myMapExpression, element); + } + + /** + * Create {@link MapLoopCondition} from enhanced for statement. + * Loop condition instance created only if iteration is done using {@link Map#keySet()} or {@link Map#entrySet()} methods. + */ + @Nullable + public static MapLoopCondition create(@NotNull PsiForeachStatement statement) { + PsiExpression iteratedValue = statement.getIteratedValue(); + PsiParameter iterParam = statement.getIterationParameter(); + PsiMethodCallExpression iterCall = extractMapMethodCall(iteratedValue, "keySet"); + if (iterCall != null) return create(iterParam, false, iterCall.getMethodExpression().getQualifierExpression()); + + iterCall = extractMapMethodCall(iteratedValue, "entrySet"); + if (iterCall != null) return create(iterParam, true, iterCall.getMethodExpression().getQualifierExpression()); + + return null; + } + + private static MapLoopCondition create(@NotNull PsiParameter iterParam, boolean isEntrySet, @Nullable PsiExpression qualifier) { + PsiReferenceExpression ref = ObjectUtils.tryCast(qualifier, PsiReferenceExpression.class); + if (ref == null) return null; + return new MapLoopCondition(iterParam, isEntrySet, ref); + } + } } \ No newline at end of file