Autoimport completed class names

package.json

@@ -119,6 +119,15 @@
                     ],
                     "default": "off",
                     "description": "Traces the communication between VSCode and the language server."
+                },
+                "java.importOrder": {
+                    "type": "string",
+                    "enum": [
+                        "simple",
+                        "chromium"
+                    ],
+                    "default": "simple",
+                    "description": "Import order for organizing imports."
                 }
             }
         },

src/main/java/org/javacs/JavaLanguageServer.java

@@ -16,6 +16,9 @@
 import org.javacs.completion.SignatureProvider;
 import org.javacs.fold.FoldProvider;
 import org.javacs.hover.HoverProvider;
+import org.javacs.imports.AutoImportProvider;
+import org.javacs.imports.AutoImportProviderFactory;
+import org.javacs.imports.SimpleAutoImportProvider;
 import org.javacs.index.SymbolProvider;
 import org.javacs.lens.CodeLensProvider;
 import org.javacs.lsp.*;
@@ -33,6 +36,7 @@ class JavaLanguageServer extends LanguageServer {
     private JsonObject cacheSettings;
     private JsonObject settings = new JsonObject();
     private boolean modifiedBuild = true;
+    private AutoImportProvider autoImportProvider = SimpleAutoImportProvider.INSTANCE;
 
     JavaCompilerService compiler() {
         if (needsCompiler()) {
@@ -226,6 +230,7 @@ public void didChangeConfiguration(DidChangeConfigurationParams change) {
         var java = change.settings.getAsJsonObject().get("java");
         LOG.info("Received java settings " + java);
         settings = java.getAsJsonObject();
+        updateAutoImportProvider();
     }
 
     @Override
@@ -260,7 +265,7 @@ public void didChangeWatchedFiles(DidChangeWatchedFilesParams params) {
     public Optional<CompletionList> completion(TextDocumentPositionParams params) {
         if (!FileStore.isJavaFile(params.textDocument.uri)) return Optional.empty();
         var file = Paths.get(params.textDocument.uri);
-        var provider = new CompletionProvider(compiler());
+        var provider = new CompletionProvider(compiler(), autoImportProvider);
         var list = provider.complete(file, params.position.line + 1, params.position.character + 1);
         if (list == CompletionProvider.NOT_SUPPORTED) return Optional.empty();
         return Optional.of(list);
@@ -510,7 +515,7 @@ public void didCloseTextDocument(DidCloseTextDocumentParams params) {
 
     @Override
     public List<CodeAction> codeAction(CodeActionParams params) {
-        var provider = new CodeActionProvider(compiler());
+        var provider = new CodeActionProvider(compiler(), autoImportProvider);
         if (params.context.diagnostics.isEmpty()) {
             return provider.codeActionsForCursor(params);
         } else {
@@ -534,5 +539,19 @@ public void doAsyncWork() {
         }
     }
 
+    private void updateAutoImportProvider() {
+        if (!settings.has("importOrder")) {
+            return;
+        }
+        var name = settings.getAsJsonPrimitive("importOrder").getAsString();
+        try {
+            autoImportProvider = AutoImportProviderFactory.getByName(name);
+        } catch (IllegalArgumentException e) {
+            LOG.warning("Unknown import order: " + name);
+            LOG.warning("Falling back to the default import order");
+            autoImportProvider = SimpleAutoImportProvider.INSTANCE;
+        }
+    }
+
     private static final Logger LOG = Logger.getLogger("main");
 }

src/main/java/org/javacs/action/CodeActionProvider.java

@@ -13,14 +13,17 @@
 import javax.lang.model.element.*;
 import org.javacs.*;
 import org.javacs.FindTypeDeclarationAt;
+import org.javacs.imports.AutoImportProvider;
 import org.javacs.lsp.*;
 import org.javacs.rewrite.*;
 
 public class CodeActionProvider {
     private final CompilerProvider compiler;
+    private final AutoImportProvider autoImportProvider;
 
-    public CodeActionProvider(CompilerProvider compiler) {
+    public CodeActionProvider(CompilerProvider compiler, AutoImportProvider autoImportProvider) {
         this.compiler = compiler;
+        this.autoImportProvider = autoImportProvider;
     }
 
     public List<CodeAction> codeActionsForCursor(CodeActionParams params) {
@@ -171,7 +174,7 @@ private List<CodeAction> codeActionForDiagnostic(CompileTask task, Path file, Di
                 for (var qualifiedName : compiler.publicTopLevelTypes()) {
                     if (qualifiedName.endsWith("." + simpleName)) {
                         var title = "Import '" + qualifiedName + "'";
-                        var addImport = new AddImport(file, qualifiedName);
+                        var addImport = new AddImport(file, qualifiedName, autoImportProvider);
                         allImports.addAll(createQuickFix(title, addImport));
                     }
                 }

src/main/java/org/javacs/completion/CompletionProvider.java

@@ -8,6 +8,7 @@
 import com.sun.source.tree.Scope;
 import com.sun.source.tree.SwitchTree;
 import com.sun.source.tree.Tree;
+import com.sun.source.util.SourcePositions;
 import com.sun.source.util.TreePath;
 import com.sun.source.util.Trees;
 import java.nio.file.Path;
@@ -40,6 +41,7 @@
 import org.javacs.ParseTask;
 import org.javacs.SourceFileObject;
 import org.javacs.StringSearch;
+import org.javacs.imports.AutoImportProvider;
 import org.javacs.lsp.Command;
 import org.javacs.lsp.CompletionItem;
 import org.javacs.lsp.CompletionItemKind;
@@ -48,6 +50,7 @@
 
 public class CompletionProvider {
     private final CompilerProvider compiler;
+    private final AutoImportProvider autoImportProvider;
 
     public static final CompletionList NOT_SUPPORTED = new CompletionList(false, List.of());
     public static final int MAX_COMPLETION_ITEMS = 50;
@@ -116,8 +119,9 @@ public class CompletionProvider {
         "double",
     };
 
-    public CompletionProvider(CompilerProvider compiler) {
+    public CompletionProvider(CompilerProvider compiler, AutoImportProvider autoImportProvider) {
         this.compiler = compiler;
+        this.autoImportProvider = autoImportProvider;
     }
 
     public CompletionList complete(Path file, int line, int column) {
@@ -235,7 +239,7 @@ private CompletionList completeIdentifier(CompileTask task, TreePath path, Strin
         list.items = completeUsingScope(task, path, partial, endsWithParen);
         addStaticImports(task, path.getCompilationUnit(), partial, endsWithParen, list);
         if (!list.isIncomplete && partial.length() > 0 && Character.isUpperCase(partial.charAt(0))) {
-            addClassNames(path.getCompilationUnit(), partial, list);
+            addClassNames(path.getCompilationUnit(), Trees.instance(task.task).getSourcePositions(), partial, list);
         }
         addKeywords(path, partial, list);
         return list;
@@ -332,24 +336,23 @@ private boolean memberMatchesImport(Name staticImport, Element member) {
         return staticImport.contentEquals("*") || staticImport.contentEquals(member.getSimpleName());
     }
 
-    private void addClassNames(CompilationUnitTree root, String partial, CompletionList list) {
+    private void addClassNames(CompilationUnitTree root, SourcePositions sourcePositions, String partial, CompletionList list) {
         var packageName = Objects.toString(root.getPackageName(), "");
-        var uniques = new HashSet<String>();
         var previousSize = list.items.size();
         for (var className : compiler.packagePrivateTopLevelTypes(packageName)) {
-            if (!StringSearch.matchesPartialName(className, partial)) continue;
-            list.items.add(classItem(className));
-            uniques.add(className);
+            var item = classItem(className);
+            item.additionalTextEdits = autoImportProvider.addImport(className, root, sourcePositions);
+            list.items.add(item);
         }
         for (var className : compiler.publicTopLevelTypes()) {
             if (!StringSearch.matchesPartialName(simpleName(className), partial)) continue;
-            if (uniques.contains(className)) continue;
             if (list.items.size() > MAX_COMPLETION_ITEMS) {
                 list.isIncomplete = true;
                 break;
             }
-            list.items.add(classItem(className));
-            uniques.add(className);
+            var item = classItem(className);
+            item.additionalTextEdits = autoImportProvider.addImport(className, root, sourcePositions);
+            list.items.add(item);
         }
         LOG.info("...found " + (list.items.size() - previousSize) + " class names");
     }

src/main/java/org/javacs/imports/AutoImportProvider.java

@@ -0,0 +1,16 @@
+package org.javacs.imports;
+
+import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.util.SourcePositions;
+import java.util.List;
+import org.javacs.lsp.TextEdit;
+
+/**
+ * Provides the functionality to auto-import classes.
+ */
+public interface AutoImportProvider {
+    /**
+     * Computes edits to add an import statement of the given class name to the Java file.
+     */
+    List<TextEdit> addImport(String className, CompilationUnitTree root, SourcePositions sourcePositions);
+}

src/main/java/org/javacs/imports/AutoImportProviderFactory.java

@@ -0,0 +1,21 @@
+package org.javacs.imports;
+
+import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.util.SourcePositions;
+import org.javacs.lsp.TextEdit;
+
+/**
+ * The factory of AutoImportProvider.
+ */
+public class AutoImportProviderFactory {
+    public static AutoImportProvider getByName(String name) {
+        switch (name) {
+            case "", "default", "simple":
+                return SimpleAutoImportProvider.INSTANCE;
+            case "chromium":
+                return ChromiumAutoImportProvider.INSTANCE;
+            default:
+                throw new IllegalArgumentException("Unknown import order: " + name);
+        }
+    }
+}

src/main/java/org/javacs/imports/ChromiumAutoImportProvider.java

@@ -0,0 +1,65 @@
+package org.javacs.imports;
+
+import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.tree.ImportTree;
+import com.sun.source.util.SourcePositions;
+import java.util.List;
+import javax.tools.Diagnostic;
+import org.javacs.lsp.Position;
+import org.javacs.lsp.Range;
+import org.javacs.lsp.TextEdit;
+
+/**
+ * The import order for Chromium Java source files.
+ *
+ * https://chromium.googlesource.com/chromium/src/+/main/styleguide/java/java.md#import-order
+ */
+public class ChromiumAutoImportProvider implements AutoImportProvider {
+    public static final ChromiumAutoImportProvider INSTANCE = new ChromiumAutoImportProvider();
+
+    private final SectionedImportOrderHelper helper = new SectionedImportOrderHelper(ChromiumAutoImportProvider::sectionOf);
+
+    private ChromiumAutoImportProvider() {}
+
+    @Override
+    public List<TextEdit> addImport(String className, CompilationUnitTree root, SourcePositions sourcePositions) {
+        return helper.addImport(className, root, sourcePositions);
+    }
+
+    private static int sectionOf(String className, boolean isStatic) {
+        if (isStatic) {
+            return 0;
+        }
+        if (className.startsWith("com.google.android.apps.chrome.")) {
+            return 7;
+        }
+        if (className.startsWith("org.chromium.")) {
+            return 8;
+        }
+        if (className.startsWith("android.")) {
+            return 1;
+        }
+        if (className.startsWith("androidx.")) {
+            return 2;
+        }
+        if (className.startsWith("com.")) {
+            return 3;
+        }
+        if (className.startsWith("dalvik.")) {
+            return 4;
+        }
+        if (className.startsWith("junit.")) {
+            return 5;
+        }
+        if (className.startsWith("org.")) {
+            return 6;
+        }
+        if (className.startsWith("java.")) {
+            return 9;
+        }
+        if (className.startsWith("javax.")) {
+            return 10;
+        }
+        return 99;
+    }
+}