RemoveUnusedImportsTest

Author: Vincent Potucek

palantir-java-format/src/main/java/com/palantir/javaformat/java/RemoveUnusedImports.java

@@ -16,15 +16,21 @@
 
 package com.palantir.javaformat.java;
 
+import static java.lang.Math.max;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
 import com.google.common.base.CharMatcher;
 import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Range;
 import com.google.common.collect.RangeMap;
 import com.google.common.collect.RangeSet;
 import com.google.common.collect.TreeRangeMap;
 import com.google.common.collect.TreeRangeSet;
 import com.palantir.javaformat.Newlines;
+import com.palantir.javaformat.java.FormatterException;
 import com.sun.source.doctree.DocCommentTree;
 import com.sun.source.doctree.ReferenceTree;
 import com.sun.source.tree.CaseTree;
@@ -36,24 +42,36 @@
 import com.sun.source.util.TreePathScanner;
 import com.sun.source.util.TreeScanner;
 import com.sun.tools.javac.api.JavacTrees;
+import com.sun.tools.javac.file.JavacFileManager;
+import com.sun.tools.javac.parser.JavacParser;
+import com.sun.tools.javac.parser.ParserFactory;
 import com.sun.tools.javac.tree.DCTree;
 import com.sun.tools.javac.tree.DCTree.DCReference;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
 import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
-import com.sun.tools.javac.tree.JCTree.JCIdent;
 import com.sun.tools.javac.tree.JCTree.JCImport;
 import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.Log;
 import com.sun.tools.javac.util.Options;
+import java.io.IOError;
+import java.io.IOException;
 import java.lang.reflect.Method;
+import java.net.URI;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import javax.tools.Diagnostic;
+import javax.tools.DiagnosticCollector;
+import javax.tools.DiagnosticListener;
+import javax.tools.JavaFileObject;
+import javax.tools.SimpleJavaFileObject;
+import javax.tools.StandardLocation;
 
 /**
- * Removes unused imports from a source file. Imports that are only used in javadoc are also removed, and the references
- * in javadoc are replaced with fully qualified names.
+ * Removes unused imports from a source file. Imports that are only used in javadoc are also
+ * removed, and the references in javadoc are replaced with fully qualified names.
  */
 public class RemoveUnusedImports {
 
@@ -73,12 +91,10 @@ public class RemoveUnusedImports {
     // This is still reasonably effective in practice because type names differ
     // from other kinds of names in casing convention, and simple name
     // clashes between imported and declared types are rare.
-    private static final class UnusedImportScanner extends TreePathScanner<Void, Void> {
+    private static class UnusedImportScanner extends TreePathScanner<Void, Void> {
 
         private final Set<String> usedNames = new LinkedHashSet<>();
-
         private final Multimap<String, Range<Integer>> usedInJavadoc = HashMultimap.create();
-
         final JavacTrees trees;
         final DocTreeScanner docTreeSymbolScanner;
 
@@ -102,7 +118,7 @@ public Void visitIdentifier(IdentifierTree tree, Void unused) {
             return null;
         }
 
-        // TODO(fwindheuser): remove this override when pattern matching in switch is no longer a preview
+        // TODO(cushon): remove this override when pattern matching in switch is no longer a preview
         // feature, and TreePathScanner visits CaseTree#getLabels instead of CaseTree#getExpressions
         @SuppressWarnings("unchecked") // reflection
         @Override
@@ -111,15 +127,14 @@ public Void visitCase(CaseTree tree, Void unused) {
                 try {
                     scan((List<? extends Tree>) CASE_TREE_GET_LABELS.invoke(tree), null);
                 } catch (ReflectiveOperationException e) {
-                    throw new RuntimeException(e.getMessage(), e);
+                    throw new LinkageError(e.getMessage(), e);
                 }
             }
             return super.visitCase(tree, null);
         }
 
         private static final Method CASE_TREE_GET_LABELS = caseTreeGetLabels();
 
-        @SuppressWarnings("for-rollout:NullAway")
         private static Method caseTreeGetLabels() {
             try {
                 return CaseTree.class.getMethod("getLabels");
@@ -128,7 +143,6 @@ private static Method caseTreeGetLabels() {
             }
         }
 
-        @SuppressWarnings("for-rollout:VoidUsed")
         @Override
         public Void scan(Tree tree, Void unused) {
             if (tree == null) {
@@ -159,9 +173,10 @@ public Void visitIdentifier(com.sun.source.doctree.IdentifierTree node, Void aVo
             @Override
             public Void visitReference(ReferenceTree referenceTree, Void unused) {
                 DCReference reference = (DCReference) referenceTree;
-                long basePos = reference
-                        .pos((DCTree.DCDocComment) getCurrentPath().getDocComment())
-                        .getStartPosition();
+                long basePos =
+                        reference
+                                .pos((DCTree.DCDocComment) getCurrentPath().getDocComment())
+                                .getStartPosition();
                 // the position of trees inside the reference node aren't stored, but the qualifier's
                 // start position is the beginning of the reference node
                 if (reference.qualifierExpression != null) {
@@ -186,37 +201,78 @@ public ReferenceScanner(long basePos) {
                     this.basePos = basePos;
                 }
 
-                @SuppressWarnings("for-rollout:VoidUsed")
                 @Override
                 public Void visitIdentifier(IdentifierTree node, Void aVoid) {
                     usedInJavadoc.put(
                             node.getName().toString(),
                             basePos != -1
-                                    ? Range.closedOpen(
-                                            (int) basePos,
-                                            (int) basePos + node.getName().length())
+                                    ? Range.closedOpen((int) basePos, (int) basePos + node.getName().length())
                                     : null);
                     return super.visitIdentifier(node, aVoid);
                 }
             }
         }
     }
-
     public static String removeUnusedImports(final String contents) throws FormatterException {
         Context context = new Context();
         JCCompilationUnit unit = parse(context, contents);
-        if (unit == null) {
-            // error handling is done during formatting
-            return contents;
-        }
         UnusedImportScanner scanner = new UnusedImportScanner(JavacTrees.instance(context));
         scanner.scan(unit, null);
-        return applyReplacements(contents, buildReplacements(contents, unit, scanner.usedNames, scanner.usedInJavadoc));
+        String s = applyReplacements(
+                contents, buildReplacements(contents, unit, scanner.usedNames, scanner.usedInJavadoc));
+
+        // Normalize newlines while preserving important blank lines
+        String sep = Newlines.guessLineSeparator(contents);
+
+        // Ensure exactly one blank line after package declaration
+        s = s.replaceAll("(?m)^package .+" + sep + "\\s+" + sep, "package $1" + sep + sep);
+
+        // Ensure exactly one blank line between last import and class declaration
+        s = s.replaceAll("(?m)import .+" + sep + "\\s+" + sep + "(?=class|interface|enum|record)",
+                "import $1" + sep + sep);
+
+        // Remove multiple blank lines elsewhere in imports section
+        s = s.replaceAll("(?m)^import .+" + sep + "\\s+" + sep + "(?=import)", "import $1" + sep);
+
+        return s;
     }
 
-    private static JCCompilationUnit parse(Context context, String javaInput) throws FormatterException {
+    private static JCCompilationUnit parse(Context context, String javaInput)
+            throws FormatterException {
+        DiagnosticCollector<JavaFileObject> diagnostics = new DiagnosticCollector<>();
+        context.put(DiagnosticListener.class, diagnostics);
+        Options.instance(context).put("--enable-preview", "true");
         Options.instance(context).put("allowStringFolding", "false");
-        return Formatter.parseJcCompilationUnit(context, javaInput);
+        JCCompilationUnit unit;
+        try (JavacFileManager fileManager = new JavacFileManager(context, true, UTF_8)){
+            fileManager.setLocation(StandardLocation.PLATFORM_CLASS_PATH, ImmutableList.of());
+        } catch (IOException e) {
+            // impossible
+            throw new IOError(e);
+        }
+        SimpleJavaFileObject source =
+                new SimpleJavaFileObject(URI.create("source"), JavaFileObject.Kind.SOURCE) {
+                    @Override
+                    public CharSequence getCharContent(boolean ignoreEncodingErrors) throws IOException {
+                        return javaInput;
+                    }
+                };
+        Log.instance(context).useSource(source);
+        ParserFactory parserFactory = ParserFactory.instance(context);
+        JavacParser parser =
+                parserFactory.newParser(
+                        javaInput,
+                        /* keepDocComments= */ true,
+                        /* keepEndPos= */ true,
+                        /* keepLineMap= */ true);
+        unit = parser.parseCompilationUnit();
+        unit.sourcefile = source;
+        Iterable<Diagnostic<? extends JavaFileObject>> errorDiagnostics =
+                Iterables.filter(diagnostics.getDiagnostics(), Formatter::errorDiagnostic);
+        if (!Iterables.isEmpty(errorDiagnostics)) {
+            // error handling is done during formatting
+        }
+        return unit;
     }
 
     /** Construct replacements to fix unused imports. */
@@ -226,53 +282,63 @@ private static RangeMap<Integer, String> buildReplacements(
             Set<String> usedNames,
             Multimap<String, Range<Integer>> usedInJavadoc) {
         RangeMap<Integer, String> replacements = TreeRangeMap.create();
-        for (JCImport importTree : unit.getImports()) {
+        int size = unit.getImports().size();
+        JCTree lastImport = size > 0 ? unit.getImports().get(size - 1) : null;
+        for (JCTree importTree : unit.getImports()) {
             String simpleName = getSimpleName(importTree);
             if (!isUnused(unit, usedNames, usedInJavadoc, importTree, simpleName)) {
                 continue;
             }
             // delete the import
             int endPosition = importTree.getEndPosition(unit.endPositions);
-            endPosition = Math.max(CharMatcher.isNot(' ').indexIn(contents, endPosition), endPosition);
+            endPosition = max(CharMatcher.isNot(' ').indexIn(contents, endPosition), endPosition);
             String sep = Newlines.guessLineSeparator(contents);
+
+            // Check if there's an empty line after this import
+            boolean hasEmptyLineAfter = false;
+            if (endPosition + sep.length() * 2 <= contents.length()) {
+                String nextTwoLines = contents.substring(endPosition, endPosition + sep.length() * 2);
+                hasEmptyLineAfter = nextTwoLines.equals(sep + sep);
+            }
+
             if (endPosition + sep.length() < contents.length()
-                    && contents.subSequence(endPosition, endPosition + sep.length())
-                            .toString()
-                            .equals(sep)) {
+                    && contents.subSequence(endPosition, endPosition + sep.length()).toString().equals(sep)) {
                 endPosition += sep.length();
             }
+
+            // If this isn't the last import and there's an empty line after, preserve it
+            if ((size == 1 || importTree != lastImport) && !hasEmptyLineAfter) {
+                while (endPosition + sep.length() <= contents.length()
+                        && contents.regionMatches(endPosition, sep, 0, sep.length())) {
+                    endPosition += sep.length();
+                }
+            }
             replacements.put(Range.closedOpen(importTree.getStartPosition(), endPosition), "");
         }
         return replacements;
     }
-
-    private static String getSimpleName(ImportTree importTree) {
-        return importTree.getQualifiedIdentifier() instanceof JCIdent
-                ? ((JCIdent) importTree.getQualifiedIdentifier()).getName().toString()
-                : ((JCFieldAccess) importTree.getQualifiedIdentifier())
-                        .getIdentifier()
-                        .toString();
+    private static String getSimpleName(JCTree importTree) {
+        return getQualifiedIdentifier(importTree).getIdentifier().toString();
     }
 
     private static boolean isUnused(
             JCCompilationUnit unit,
             Set<String> usedNames,
             Multimap<String, Range<Integer>> usedInJavadoc,
-            ImportTree importTree,
+            JCTree importTree,
             String simpleName) {
-        String qualifier = ((JCFieldAccess) importTree.getQualifiedIdentifier())
-                .getExpression()
-                .toString();
+        JCFieldAccess qualifiedIdentifier = getQualifiedIdentifier(importTree);
+        String qualifier = qualifiedIdentifier.getExpression().toString();
         if (qualifier.equals("java.lang")) {
             return true;
         }
+        if(usedNames.contains(simpleName)){
+            return false;
+        }
         if (unit.getPackageName() != null && unit.getPackageName().toString().equals(qualifier)) {
             return true;
         }
-        if (importTree.getQualifiedIdentifier() instanceof JCFieldAccess
-                && ((JCFieldAccess) importTree.getQualifiedIdentifier())
-                        .getIdentifier()
-                        .contentEquals("*")) {
+        if (qualifiedIdentifier.getIdentifier().contentEquals("*") && !((JCImport) importTree).isStatic()) {
             return false;
         }
 
@@ -285,6 +351,15 @@ private static boolean isUnused(
         return true;
     }
 
+    private static JCFieldAccess getQualifiedIdentifier(JCTree importTree) {
+        // Use reflection because the return type is JCTree in some versions and JCFieldAccess in others
+        try {
+            return (JCFieldAccess) JCImport.class.getMethod("getQualifiedIdentifier").invoke(importTree);
+        } catch (ReflectiveOperationException e) {
+            throw new LinkageError(e.getMessage(), e);
+        }
+    }
+
     /** Applies the replacements to the given source, and re-format any edited javadoc. */
     private static String applyReplacements(String source, RangeMap<Integer, String> replacements) {
         // save non-empty fixed ranges for reformatting after fixes are applied
@@ -296,8 +371,7 @@ private static String applyReplacements(String source, RangeMap<Integer, String>
         // be applied in descending order without adjusting offsets.
         StringBuilder sb = new StringBuilder(source);
         int offset = 0;
-        for (Map.Entry<Range<Integer>, String> replacement :
-                replacements.asMapOfRanges().entrySet()) {
+        for (Map.Entry<Range<Integer>, String> replacement : replacements.asMapOfRanges().entrySet()) {
             Range<Integer> range = replacement.getKey();
             String replaceWith = replacement.getValue();
             int start = offset + range.lowerEndpoint();
@@ -310,4 +384,4 @@ private static String applyReplacements(String source, RangeMap<Integer, String>
         }
         return sb.toString();
     }
-}
+}