fix RemoveUnusedImportsStep leftovers

Author: Vincent Potucek

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

@@ -16,8 +16,13 @@
 
 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;
@@ -36,20 +41,32 @@
 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
@@ -76,15 +93,12 @@ public class RemoveUnusedImports {
     private static final 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;
+        private final DocTreeScanner docTreeSymbolScanner = new DocTreeScanner();
+        private final JavacTrees trees;
 
         private UnusedImportScanner(JavacTrees trees) {
             this.trees = trees;
-            docTreeSymbolScanner = new DocTreeScanner();
         }
 
         /** Skip the imports themselves when checking for usage. */
@@ -202,21 +216,58 @@ public Void visitIdentifier(IdentifierTree node, Void aVoid) {
         }
     }
 
-    public static String removeUnusedImports(final String contents) throws FormatterException {
+    public static String removeUnusedImports(final String contents) {
         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, "$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)", "$1" + sep + sep);
+
+        // Remove multiple blank lines elsewhere in imports section
+        s = s.replaceAll("(?m)^(import .+)" + sep + "\\s+" + sep + "(?=import)", "$1" + sep);
+
+        return s;
     }
 
-    private static JCCompilationUnit parse(Context context, String javaInput) throws FormatterException {
+    private static JCCompilationUnit parse(Context context, String javaInput) {
+        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) {
+            throw new IOError(e);
+        }
+        SimpleJavaFileObject source = new SimpleJavaFileObject(URI.create("source"), JavaFileObject.Kind.SOURCE) {
+            @Override
+            public CharSequence getCharContent(boolean ignoreEncodingErrors) {
+                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 +277,66 @@ 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)) {
                 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 +349,16 @@ 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