Auto-import completed classes

src/main/java/org/javacs/Extractors.java

@@ -0,0 +1,26 @@
+package org.javacs;
+
+import java.util.regex.Pattern;
+
+public class Extractors {
+
+    private static final Pattern PACKAGE_EXTRACTOR = Pattern.compile("^([a-z][_a-zA-Z0-9]*\\.)*[a-z][_a-zA-Z0-9]*");
+
+    public static String packageName(String className) {
+        var m = PACKAGE_EXTRACTOR.matcher(className);
+        if (m.find()) {
+            return m.group();
+        }
+        return "";
+    }
+
+    private static final Pattern SIMPLE_EXTRACTOR = Pattern.compile("[A-Z][_a-zA-Z0-9]*$");
+
+    public static String simpleName(String className) {
+        var m = SIMPLE_EXTRACTOR.matcher(className);
+        if (m.find()) {
+            return m.group();
+        }
+        return "";
+    }
+}

src/main/java/org/javacs/JavaCompilerService.java

@@ -97,26 +97,6 @@ private CompileBatch compileBatch(Collection<? extends JavaFileObject> sources)
         return cachedCompile;
     }
 
-    private static final Pattern PACKAGE_EXTRACTOR = Pattern.compile("^([a-z][_a-zA-Z0-9]*\\.)*[a-z][_a-zA-Z0-9]*");
-
-    private String packageName(String className) {
-        var m = PACKAGE_EXTRACTOR.matcher(className);
-        if (m.find()) {
-            return m.group();
-        }
-        return "";
-    }
-
-    private static final Pattern SIMPLE_EXTRACTOR = Pattern.compile("[A-Z][_a-zA-Z0-9]*$");
-
-    private String simpleName(String className) {
-        var m = SIMPLE_EXTRACTOR.matcher(className);
-        if (m.find()) {
-            return m.group();
-        }
-        return "";
-    }
-
     private static final Cache<String, Boolean> cacheContainsWord = new Cache<>();
 
     private boolean containsWord(Path file, String word) {
@@ -206,7 +186,7 @@ public List<String> packagePrivateTopLevelTypes(String packageName) {
     }
 
     private boolean containsImport(Path file, String className) {
-        var packageName = packageName(className);
+        var packageName = Extractors.packageName(className);
         if (FileStore.packageName(file).equals(packageName)) return true;
         var star = packageName + ".*";
         for (var i : readImports(file)) {
@@ -273,8 +253,8 @@ public Path findTypeDeclaration(String className) {
         if (fastFind != NOT_FOUND) return fastFind;
         // In principle, the slow path can be skipped in many cases.
         // If we're spending a lot of time in findTypeDeclaration, this would be a good optimization.
-        var packageName = packageName(className);
-        var simpleName = simpleName(className);
+        var packageName = Extractors.packageName(className);
+        var simpleName = Extractors.simpleName(className);
         for (var f : FileStore.list(packageName)) {
             if (containsWord(f, simpleName) && containsType(f, className)) {
                 return f;
@@ -301,8 +281,8 @@ private Path findPublicTypeDeclaration(String className) {
 
     @Override
     public Path[] findTypeReferences(String className) {
-        var packageName = packageName(className);
-        var simpleName = simpleName(className);
+        var packageName = Extractors.packageName(className);
+        var simpleName = Extractors.simpleName(className);
         var candidates = new ArrayList<Path>();
         for (var f : FileStore.all()) {
             if (containsWord(f, packageName) && containsImport(f, className) && containsWord(f, simpleName)) {

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

@@ -2,6 +2,7 @@
 
 import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.tree.ImportTree;
 import com.sun.source.tree.MemberReferenceTree;
 import com.sun.source.tree.MemberSelectTree;
 import com.sun.source.tree.MethodTree;
@@ -15,13 +16,16 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.ExecutableElement;
@@ -35,6 +39,7 @@
 import org.javacs.CompileTask;
 import org.javacs.CompilerProvider;
 import org.javacs.CompletionData;
+import org.javacs.Extractors;
 import org.javacs.FileStore;
 import org.javacs.JsonHelper;
 import org.javacs.ParseTask;
@@ -45,6 +50,8 @@
 import org.javacs.lsp.CompletionItemKind;
 import org.javacs.lsp.CompletionList;
 import org.javacs.lsp.InsertTextFormat;
+import org.javacs.lsp.TextEdit;
+import org.javacs.rewrite.AddImport;
 
 public class CompletionProvider {
     private final CompilerProvider compiler;
@@ -336,9 +343,19 @@ private void addClassNames(CompilationUnitTree root, String partial, CompletionL
         var packageName = Objects.toString(root.getPackageName(), "");
         var uniques = new HashSet<String>();
         var previousSize = list.items.size();
+
+        var fileImports =
+                root.getImports()
+                        .stream()
+                        .map(ImportTree::getQualifiedIdentifier)
+                        .map(Tree::toString)
+                        .collect(Collectors.toSet());
+
+        var filePath = Path.of(root.getSourceFile().toUri());
+
         for (var className : compiler.packagePrivateTopLevelTypes(packageName)) {
             if (!StringSearch.matchesPartialName(className, partial)) continue;
-            list.items.add(classItem(className));
+            list.items.add(classItem(fileImports, filePath, className));
             uniques.add(className);
         }
         for (var className : compiler.publicTopLevelTypes()) {
@@ -348,9 +365,10 @@ private void addClassNames(CompilationUnitTree root, String partial, CompletionL
                 list.isIncomplete = true;
                 break;
             }
-            list.items.add(classItem(className));
+            list.items.add(classItem(fileImports, filePath, className));
             uniques.add(className);
         }
+
         LOG.info("...found " + (list.items.size() - previousSize) + " class names");
     }
 
@@ -579,17 +597,34 @@ private CompletionItem packageItem(String name) {
         return i;
     }
 
+    // This version does not add an additionalTextEdit to add the import statement. Useful for if
+    // the completion is for an import statement (which does not need an edit to add itself).
     private CompletionItem classItem(String className) {
+        return classItem(Collections.emptySet(), null, className);
+    }
+
+    private CompletionItem classItem(Set<String> fileImports, Path path, String className) {
         var i = new CompletionItem();
         i.label = simpleName(className).toString();
         i.kind = CompletionItemKind.Class;
         i.detail = className;
         var data = new CompletionData();
         data.className = className;
         i.data = JsonHelper.GSON.toJsonTree(data);
+        i.additionalTextEdits = checkForImports(fileImports, path, className);
         return i;
     }
 
+    private List<TextEdit> checkForImports(Set<String> fileImports, Path path, String className) {
+        final String star = Extractors.packageName(className) + ".*";
+        if (fileImports.contains(className) || fileImports.contains(star)) {
+            return null;
+        }
+
+        AddImport addImport = new AddImport(path, className);
+        return List.of(addImport.rewrite(compiler).get(path));
+    }
+
     private CompletionItem snippetItem(String label, String snippet) {
         var i = new CompletionItem();
         i.label = label;