MapReplaceableByEnumMap: add fix to replace Map with EnumMap IDEA-167983

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/InspectionGadgetsBundle.properties

@@ -1036,6 +1036,7 @@ field.may.be.static.problem.descriptor=Field <code>#ref</code> may be 'static' #
 method.may.be.static.problem.descriptor=Method <code>#ref()</code> may be 'static' #loc
 class.initializer.may.be.static.problem.descriptor=Class initializer may be 'static' #loc
 map.replaceable.by.enum.map.problem.descriptor=<code>#ref</code> replaceable with 'EnumMap' #loc
+map.replaceable.by.enum.map.fix.name=Replace with EnumMap
 set.replaceable.by.enum.set.problem.descriptor=<code>#ref</code> replaceable with 'EnumSet' #loc
 inner.class.may.be.static.problem.descriptor=Inner class <code>#ref</code> may be 'static' #loc
 string.buffer.must.have.initial.capacity.problem.descriptor=<code>new #ref()</code> without initial capacity #loc

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/performance/CollectionReplaceableByEnumCollectionVisitor.java

@@ -17,6 +17,7 @@ package com.siyeh.ig.performance;
 
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.PsiUtil;
 import com.siyeh.ig.BaseInspectionVisitor;
 import com.siyeh.ig.psiutils.ExpectedTypeUtils;
 import com.siyeh.ig.psiutils.TypeUtils;
@@ -24,6 +25,8 @@ import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
 
+import static com.intellij.util.ObjectUtils.tryCast;
+
 /**
  * @author Bas Leijdekkers
  */
@@ -78,11 +81,24 @@ abstract class CollectionReplaceableByEnumCollectionVisitor extends BaseInspecti
         return;
       }
     }
-    if (!expectedType.isAssignableFrom(TypeUtils.getType(getReplacementCollectionName(), expression)) &&
-        !isReplaceableType((PsiClassType)expectedType)) {
+    PsiClassType replacementCollectionType = TypeUtils.getType(getReplacementCollectionName(), expression);
+    if (!expectedType.isAssignableFrom(replacementCollectionType) && !isReplaceableType((PsiClassType)expectedType)) {
       return;
     }
-    registerNewExpressionError(expression);
+    PsiElement parent = PsiUtil.skipParenthesizedExprUp(expression).getParent();
+    PsiLocalVariable localVariable = tryCast(parent, PsiLocalVariable.class);
+    if (localVariable != null) {
+      PsiClass localVariableClass = PsiUtil.resolveClassInClassTypeOnly(localVariable.getType());
+      if (localVariableClass != null && getBaseCollectionName().equals(localVariableClass.getQualifiedName())) {
+        registerNewExpressionError(expression, localVariable);
+      }
+      else {
+        registerNewExpressionError(expression);
+      }
+    }
+    else {
+      registerNewExpressionError(expression);
+    }
   }
 
   private boolean isReplaceableType(PsiClassType classType) {

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/performance/MapReplaceableByEnumMapInspection.java

@@ -15,15 +15,28 @@
  */
 package com.siyeh.ig.performance;
 
-import com.intellij.psi.CommonClassNames;
+import com.intellij.codeInspection.ProblemDescriptor;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.*;
+import com.intellij.psi.codeStyle.JavaCodeStyleManager;
+import com.intellij.psi.impl.PsiDiamondTypeUtil;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.PsiUtil;
 import com.siyeh.InspectionGadgetsBundle;
 import com.siyeh.ig.BaseInspection;
 import com.siyeh.ig.BaseInspectionVisitor;
+import com.siyeh.ig.InspectionGadgetsFix;
+import com.siyeh.ig.psiutils.CommentTracker;
+import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
+
+import static com.intellij.util.ObjectUtils.tryCast;
 
 public class MapReplaceableByEnumMapInspection extends BaseInspection {
 
@@ -44,6 +57,61 @@ public class MapReplaceableByEnumMapInspection extends BaseInspection {
     return new MapReplaceableByEnumMapVisitor();
   }
 
+  @Nullable
+  @Override
+  protected InspectionGadgetsFix buildFix(Object... infos) {
+    if (infos.length != 1) return null;
+    PsiLocalVariable localVariable = (PsiLocalVariable)infos[0];
+    PsiExpression initializer = PsiUtil.skipParenthesizedExprDown(localVariable.getInitializer());
+    PsiNewExpression newExpression = tryCast(initializer, PsiNewExpression.class);
+    if (newExpression == null) return null;
+    PsiExpressionList argumentList = newExpression.getArgumentList();
+    if (argumentList == null || !argumentList.isEmpty()) return null;
+    PsiJavaCodeReferenceElement classReference = newExpression.getClassReference();
+    if (classReference == null) return null;
+    PsiReferenceParameterList parameterList = classReference.getParameterList();
+    if (parameterList == null) return null;
+    PsiClassType classType = tryCast(newExpression.getType(), PsiClassType.class);
+    if (classType == null) return null;
+    if (classType.getParameterCount() != 2) return null;
+    PsiType[] parameters = classType.getParameters();
+    PsiType enumParameter = parameters[0];
+    String parameterListText = Arrays.stream(parameters).map(p -> p.getCanonicalText()).collect(Collectors.joining(",", "<", ">"));
+    PsiClass probablyEnum = PsiUtil.resolveClassInClassTypeOnly(enumParameter);
+    if (probablyEnum == null || !probablyEnum.isEnum()) return null;
+    String text = "new java.util.EnumMap" + parameterListText + "(" + enumParameter.getCanonicalText() + ".class)";
+    return new EnumMapReplacingFix(text);
+  }
+
+  private static class EnumMapReplacingFix extends InspectionGadgetsFix {
+    private final String newEnumMapText;
+
+    private EnumMapReplacingFix(String text) {newEnumMapText = text;}
+
+    @Override
+    protected void doFix(Project project, ProblemDescriptor descriptor) {
+      PsiNewExpression newExpression = PsiTreeUtil.getParentOfType(descriptor.getStartElement(), PsiNewExpression.class);
+      if (newExpression == null) return;
+      PsiElement result = new CommentTracker().replaceAndRestoreComments(newExpression, newEnumMapText);
+      PsiDiamondTypeUtil.removeRedundantTypeArguments(result);
+      JavaCodeStyleManager.getInstance(project).shortenClassReferences(result);
+    }
+
+    @Nls(capitalization = Nls.Capitalization.Sentence)
+    @NotNull
+    @Override
+    public String getName() {
+      return getFamilyName();
+    }
+
+    @Nls(capitalization = Nls.Capitalization.Sentence)
+    @NotNull
+    @Override
+    public String getFamilyName() {
+      return InspectionGadgetsBundle.message("map.replaceable.by.enum.map.fix.name");
+    }
+  }
+
   private static class MapReplaceableByEnumMapVisitor extends CollectionReplaceableByEnumCollectionVisitor {
 
     @Override

plugins/InspectionGadgets/test/com/siyeh/igfixes/performance/map_replaceable_with_enum_map/Simple.after.java

@@ -0,0 +1,11 @@
+import java.util.EnumMap;
+import java.util.HashMap;
+import java.util.Map;
+
+class EnumTest {
+  enum A {foo, bar;}
+
+  void foo() {
+    Map<A, String> m = new EnumMap<>(A.class);
+  }
+}

plugins/InspectionGadgets/test/com/siyeh/igfixes/performance/map_replaceable_with_enum_map/Simple.java

@@ -0,0 +1,10 @@
+import java.util.HashMap;
+import java.util.Map;
+
+class EnumTest {
+  enum A {foo, bar;}
+
+  void foo() {
+    Map<A, String> m = new <caret>HashMap<>();
+  }
+}

plugins/InspectionGadgets/testsrc/com/siyeh/ig/fixes/performance/MapReplaceableByEnumMapFixTest.java

@@ -0,0 +1,34 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.siyeh.ig.fixes.performance;
+
+import com.siyeh.InspectionGadgetsBundle;
+import com.siyeh.ig.IGQuickFixesTestCase;
+import com.siyeh.ig.performance.MapReplaceableByEnumMapInspection;
+
+public class MapReplaceableByEnumMapFixTest extends IGQuickFixesTestCase {
+
+  @Override
+  protected void setUp() throws Exception {
+    super.setUp();
+    myFixture.enableInspections(new MapReplaceableByEnumMapInspection());
+    myRelativePath = "performance/map_replaceable_with_enum_map";
+    myDefaultHint = InspectionGadgetsBundle.message("map.replaceable.by.enum.map.fix.name");
+    myFixture.addClass("package java.util;\n" +
+                       "\n" +
+                       "public class EnumMap<K extends Enum<K>, V> extends AbstractMap<K, V>\n" +
+                       "  implements java.io.Serializable, Cloneable\n" +
+                       "{\n" +
+                       "  public EnumMap(Class<K> keyType) {\n" +
+                       "    \n" +
+                       "  }\n" +
+                       "}");
+
+  }
+
+
+  public void testSimple() {
+    doTest();
+  }
+
+
+}