Support "Add all missing imports"

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/commands/OrganizeImportsCommand.java

@@ -33,8 +33,6 @@
 import org.eclipse.jdt.core.IPackageFragmentRoot;
 import org.eclipse.jdt.core.JavaCore;
 import org.eclipse.jdt.core.JavaModelException;
-import org.eclipse.jdt.core.dom.CompilationUnit;
-import org.eclipse.jdt.core.manipulation.OrganizeImportsOperation;
 import org.eclipse.jdt.ls.core.internal.ChangeUtil;
 import org.eclipse.jdt.ls.core.internal.IConstants;
 import org.eclipse.jdt.ls.core.internal.JDTUtils;
@@ -45,7 +43,6 @@
 import org.eclipse.jdt.ls.core.internal.corrections.proposals.CUCorrectionProposal;
 import org.eclipse.jdt.ls.core.internal.corrections.proposals.IProposalRelevance;
 import org.eclipse.jdt.ls.core.internal.handlers.OrganizeImportsHandler;
-import org.eclipse.jface.text.IDocument;
 import org.eclipse.lsp4j.CodeActionKind;
 import org.eclipse.lsp4j.WorkspaceEdit;
 import org.eclipse.ltk.core.refactoring.TextChange;
@@ -176,20 +173,10 @@ public void organizeImportsInPackageFragment(IPackageFragment fragment, Workspac
 	public void organizeImportsInCompilationUnit(ICompilationUnit unit, WorkspaceEdit rootEdit) {
 		try {
 			InnovationContext context = new InnovationContext(unit, 0, unit.getBuffer().getLength() - 1);
-			CUCorrectionProposal proposal = new CUCorrectionProposal("OrganizeImports", CodeActionKind.SourceOrganizeImports, unit, null, IProposalRelevance.ORGANIZE_IMPORTS) {
-				@Override
-				protected void addEdits(IDocument document, TextEdit editRoot) throws CoreException {
-					CompilationUnit astRoot = context.getASTRoot();
-					OrganizeImportsOperation op = new OrganizeImportsOperation(unit, astRoot, true, false, true, null);
-					TextEdit edit = op.createTextEdit(null);
-					TextEdit staticEdit = OrganizeImportsHandler.wrapStaticImports(edit, astRoot, unit);
-					if (staticEdit.getChildrenSize() > 0) {
-						editRoot.addChild(staticEdit);
-					}
-				}
-			};
-
-			addWorkspaceEdit(unit, proposal, rootEdit);
+			CUCorrectionProposal proposal = OrganizeImportsHandler.getOrganizeImportsProposal("OrganizeImports", CodeActionKind.SourceOrganizeImports, unit, IProposalRelevance.ORGANIZE_IMPORTS, context.getASTRoot(), false, false);
+			if (proposal != null) {
+				addWorkspaceEdit(unit, proposal, rootEdit);
+			}
 		} catch (CoreException e) {
 			JavaLanguageServerPlugin.logException("Problem organize imports ", e);
 		}

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/corrections/CorrectionMessages.java

@@ -234,6 +234,7 @@ private CorrectionMessages() {
 	public static String UnresolvedElementsSubProcessor_arraychangetomethod_description;
 	public static String UnresolvedElementsSubProcessor_arraychangetolength_description;
 	public static String UnresolvedElementsSubProcessor_addnewkeyword_description;
+	public static String UnresolvedElementsSubProcessor_add_allMissing_imports_description;
 	public static String JavadocTagsSubProcessor_addjavadoc_method_description;
 	public static String JavadocTagsSubProcessor_addjavadoc_type_description;
 	public static String JavadocTagsSubProcessor_addjavadoc_field_description;

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/corrections/CorrectionMessages.properties

@@ -255,6 +255,7 @@ UnresolvedElementsSubProcessor_methodtargetcast_description=Add cast to method r
 UnresolvedElementsSubProcessor_arraychangetomethod_description=Change to ''{0}(..)''
 UnresolvedElementsSubProcessor_arraychangetolength_description=Change to 'length'
 UnresolvedElementsSubProcessor_addnewkeyword_description=Insert 'new' keyword
+UnresolvedElementsSubProcessor_add_allMissing_imports_description=Add all missing imports
 
 JavadocTagsSubProcessor_addjavadoc_method_description=Add Javadoc for ''{0}''
 JavadocTagsSubProcessor_addjavadoc_type_description=Add Javadoc for ''{0}''

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/corrections/proposals/UnresolvedElementsSubProcessor.java

@@ -94,6 +94,7 @@
 import org.eclipse.jdt.internal.corext.dom.ScopeAnalyzer;
 import org.eclipse.jdt.internal.corext.util.JavaModelUtil;
 import org.eclipse.jdt.internal.ui.text.correction.IProblemLocationCore;
+import org.eclipse.jdt.ls.core.internal.JavaLanguageServerPlugin;
 import org.eclipse.jdt.ls.core.internal.Messages;
 import org.eclipse.jdt.ls.core.internal.contentassist.TypeFilter;
 import org.eclipse.jdt.ls.core.internal.corrections.CorrectionMessages;
@@ -106,6 +107,7 @@
 import org.eclipse.jdt.ls.core.internal.corrections.proposals.ChangeMethodSignatureProposal.InsertDescription;
 import org.eclipse.jdt.ls.core.internal.corrections.proposals.ChangeMethodSignatureProposal.RemoveDescription;
 import org.eclipse.jdt.ls.core.internal.corrections.proposals.ChangeMethodSignatureProposal.SwapDescription;
+import org.eclipse.jdt.ls.core.internal.handlers.OrganizeImportsHandler;
 import org.eclipse.jdt.ls.core.internal.hover.JavaElementLabels;
 import org.eclipse.jdt.ls.core.internal.preferences.PreferenceManager;
 import org.eclipse.lsp4j.CodeActionKind;
@@ -241,7 +243,7 @@ public static void getVariableProposals(IInvocationContext context, IProblemLoca
 			}
 
 			int relevance= Character.isUpperCase(ASTNodes.getSimpleNameIdentifier(node).charAt(0)) ? IProposalRelevance.VARIABLE_TYPE_PROPOSAL_1 : IProposalRelevance.VARIABLE_TYPE_PROPOSAL_2;
-			addSimilarTypeProposals(typeKind, cu, node, relevance + 1, proposals);
+			addSimilarTypeProposals(context, typeKind, cu, node, relevance + 1, proposals);
 
 			typeKind &= ~TypeKinds.ANNOTATIONS;
 			addNewTypeProposals(cu, node, typeKind, relevance, proposals);
@@ -610,7 +612,7 @@ public static void getTypeProposals(IInvocationContext context, IProblemLocation
 		}
 
 		// change to similar type proposals
-		addSimilarTypeProposals(kind, cu, node, IProposalRelevance.SIMILAR_TYPE, proposals);
+		addSimilarTypeProposals(context, kind, cu, node, IProposalRelevance.SIMILAR_TYPE, proposals);
 
 		while (node.getParent() instanceof QualifiedName) {
 			node= (Name) node.getParent();
@@ -658,7 +660,7 @@ static CompilationUnitChange createAddImportChange(ICompilationUnit cu, Name nam
 		return cuChange;
 	}
 
-	private static void addSimilarTypeProposals(int kind, ICompilationUnit cu, Name node, int relevance,
+	private static void addSimilarTypeProposals(IInvocationContext context, int kind, ICompilationUnit cu, Name node, int relevance,
 			Collection<ChangeCorrectionProposal> proposals) throws CoreException {
 		SimilarElement[] elements = SimilarElementsRequestor.findSimilarElement(cu, node, kind);
 
@@ -706,6 +708,12 @@ private static void addSimilarTypeProposals(int kind, ICompilationUnit cu, Name
 				}
 			}
 		}
+		if (proposals.size() > 0) {
+			CUCorrectionProposal proposal = OrganizeImportsHandler.getOrganizeImportsProposal(CorrectionMessages.UnresolvedElementsSubProcessor_add_allMissing_imports_description, CodeActionKind.QuickFix, cu, relevance, context.getASTRoot(), JavaLanguageServerPlugin.getPreferencesManager().getClientPreferences().isAdvancedOrganizeImportsSupported(), true);
+			if (proposal != null) {
+				proposals.add(proposal);
+			}
+		}
 	}
 
 	private static CUCorrectionProposal createTypeRefChangeProposal(ICompilationUnit cu, String fullName, Name node, int relevance, int maxProposals) {

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/handlers/CodeActionHandler.java

@@ -283,9 +283,11 @@ private Optional<Either<Command, CodeAction>> getCodeActionFromProposal(ChangeCo
 			CodeAction codeAction = new CodeAction(name);
 			codeAction.setKind(proposal.getKind());
 			if (command == null) { // lazy resolve the edit.
-				codeAction.setData(new CodeActionData(proposal));
+				// The relevance is in descending order while CodeActionComparator sorts in ascending order
+				codeAction.setData(new CodeActionData(proposal, -proposal.getRelevance()));
 			} else {
 				codeAction.setCommand(command);
+				codeAction.setData(new CodeActionData(null, -proposal.getRelevance()));
 			}
 			if (proposal.getKind() != JavaCodeActionKind.QUICK_ASSIST) {
 				codeAction.setDiagnostics(context.getDiagnostics());

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/handlers/JDTLanguageServer.java

@@ -48,6 +48,7 @@
 import org.eclipse.jdt.ls.core.internal.LanguageServerWorkingCopyOwner;
 import org.eclipse.jdt.ls.core.internal.ServiceStatus;
 import org.eclipse.jdt.ls.core.internal.codemanipulation.GenerateGetterSetterOperation.AccessorField;
+import org.eclipse.jdt.ls.core.internal.handlers.CodeActionHandler.CodeActionData;
 import org.eclipse.jdt.ls.core.internal.handlers.FindLinksHandler.FindLinksParams;
 import org.eclipse.jdt.ls.core.internal.handlers.GenerateAccessorsHandler.AccessorCodeActionParams;
 import org.eclipse.jdt.ls.core.internal.handlers.GenerateAccessorsHandler.GenerateAccessorsParams;
@@ -701,10 +702,17 @@ public CompletableFuture<List<Either<Command, CodeAction>>> codeAction(CodeActio
 	@Override
 	public CompletableFuture<CodeAction> resolveCodeAction(CodeAction params) {
 		logInfo(">> codeAction/resolve");
+		Object data = params.getData();
 		// if no data property is specified, no further resolution the server can provide, so return the original result back.
-		if (params.getData() == null) {
+		if (data == null) {
 			return CompletableFuture.completedFuture(params);
 		}
+		if (data instanceof CodeActionData) {
+			// if the data is CodeActionData and no proposal in the data, return the original result back.
+			if (((CodeActionData) data).getProposal() == null) {
+				return CompletableFuture.completedFuture(params);
+			}
+		}
 		if (CodeActionHandler.codeActionStore.isEmpty()) {
 			return CompletableFuture.completedFuture(params);
 		}

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/handlers/OrganizeImportsHandler.java

@@ -13,6 +13,7 @@
 
 package org.eclipse.jdt.ls.core.internal.handlers;
 
+import java.net.URI;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -21,6 +22,7 @@
 import java.util.function.Function;
 import java.util.stream.Stream;
 
+import org.eclipse.core.resources.IResource;
 import org.eclipse.core.runtime.CoreException;
 import org.eclipse.core.runtime.IProgressMonitor;
 import org.eclipse.core.runtime.NullProgressMonitor;
@@ -50,8 +52,10 @@
 import org.eclipse.jdt.ls.core.internal.JavaLanguageServerPlugin;
 import org.eclipse.jdt.ls.core.internal.JobHelpers;
 import org.eclipse.jdt.ls.core.internal.corrections.SimilarElementsRequestor;
+import org.eclipse.jdt.ls.core.internal.corrections.proposals.CUCorrectionProposal;
 import org.eclipse.jdt.ls.core.internal.preferences.PreferenceManager;
 import org.eclipse.jdt.ls.core.internal.text.correction.SourceAssistProcessor;
+import org.eclipse.jface.text.IDocument;
 import org.eclipse.lsp4j.CodeActionParams;
 import org.eclipse.lsp4j.Range;
 import org.eclipse.lsp4j.WorkspaceEdit;
@@ -67,10 +71,10 @@ public final class OrganizeImportsHandler {
 
 	// For test purpose
 	public static TextEdit organizeImports(ICompilationUnit unit, Function<ImportSelection[], ImportCandidate[]> chooseImports) {
-		return organizeImports(unit, chooseImports, new NullProgressMonitor());
+		return organizeImports(unit, chooseImports, false, new NullProgressMonitor());
 	}
 
-	public static TextEdit organizeImports(ICompilationUnit unit, Function<ImportSelection[], ImportCandidate[]> chooseImports, IProgressMonitor monitor) {
+	public static TextEdit organizeImports(ICompilationUnit unit, Function<ImportSelection[], ImportCandidate[]> chooseImports, boolean restoreExistingImports, IProgressMonitor monitor) {
 		if (unit == null) {
 			return null;
 		}
@@ -80,7 +84,7 @@ public static TextEdit organizeImports(ICompilationUnit unit, Function<ImportSel
 			return null;
 		}
 
-		OrganizeImportsOperation op = new OrganizeImportsOperation(unit, astRoot, true, false, true, (TypeNameMatch[][] openChoices, ISourceRange[] ranges) -> {
+		OrganizeImportsOperation op = new OrganizeImportsOperation(unit, astRoot, true, false, true, chooseImports != null ? (TypeNameMatch[][] openChoices, ISourceRange[] ranges) -> {
 			List<ImportSelection> selections = new ArrayList<>();
 			for (int i = 0; i < openChoices.length; i++) {
 				ImportCandidate[] candidates = Stream.of(openChoices[i]).map((choice) -> new ImportCandidate(choice)).toArray(ImportCandidate[]::new);
@@ -104,7 +108,7 @@ public static TextEdit organizeImports(ICompilationUnit unit, Function<ImportSel
 				typeMaps.put(x.getFullyQualifiedName() + "@" + x.hashCode(), x);
 			});
 			return Stream.of(chosens).filter(chosen -> chosen != null && typeMaps.containsKey(chosen.id)).map(chosen -> typeMaps.get(chosen.id)).toArray(TypeNameMatch[]::new);
-		});
+		} : null, restoreExistingImports);
 		try {
 			JobHelpers.waitForJobs(DocumentLifeCycleHandler.DOCUMENT_LIFE_CYCLE_JOBS, monitor);
 			TextEdit edit = op.createTextEdit(null);
@@ -213,12 +217,36 @@ public static WorkspaceEdit organizeImports(JavaClientConnection connection, Cod
 			return null;
 		}
 
-		TextEdit edit = organizeImports(unit, (selections) -> {
-			Object commandResult = connection.executeClientCommand(CLIENT_COMMAND_ID_CHOOSEIMPORTS, uri, selections);
+		TextEdit edit = organizeImports(unit, getChooseImportsFunction(uri, false), false, monitor);
+		return SourceAssistProcessor.convertToWorkspaceEdit(unit, edit);
+	}
+
+	public static Function<ImportSelection[], ImportCandidate[]> getChooseImportsFunction(String documentUri, boolean restoreExistingImports) {
+		return (selections) -> {
+			Object commandResult = JavaLanguageServerPlugin.getInstance().getClientConnection().executeClientCommand(CLIENT_COMMAND_ID_CHOOSEIMPORTS, documentUri, selections, restoreExistingImports);
 			String json = commandResult == null ? null : new Gson().toJson(commandResult);
 			return JSONUtility.toModel(json, ImportCandidate[].class);
-		}, monitor);
-		return SourceAssistProcessor.convertToWorkspaceEdit(unit, edit);
+		};
+	}
+
+	public static CUCorrectionProposal getOrganizeImportsProposal(String label, String kind, ICompilationUnit cu, int relevance, CompilationUnit astRoot, boolean supportsChooseImports, boolean restoreExistingImports) {
+		IResource resource = cu.getResource();
+		if (resource == null) {
+			return null;
+		}
+		URI uri = resource.getLocationURI();
+		if (uri == null) {
+			return null;
+		}
+		return new CUCorrectionProposal(label, kind, cu, null, relevance + 10) {
+			@Override
+			protected void addEdits(IDocument document, TextEdit editRoot) throws CoreException {
+				TextEdit edit = OrganizeImportsHandler.organizeImports(cu, supportsChooseImports ? OrganizeImportsHandler.getChooseImportsFunction(uri.toString(), restoreExistingImports) : null, restoreExistingImports, new NullProgressMonitor());
+				if (edit != null) {
+					editRoot.addChild(edit);
+				}
+			}
+		};
 	}
 
 	public static class ImportCandidate {

org.eclipse.jdt.ls.core/src/org/eclipse/jdt/ls/core/internal/text/correction/CodeActionComparator.java

@@ -24,6 +24,7 @@
 public class CodeActionComparator implements Comparator<Either<Command, CodeAction>> {
 
 	public static int ORGANIZE_IMPORTS_PRIORITY = 0;
+	public static int ADD_ALL_MISSING_IMPORTS_PRIORITY = 5;
 	public static int GENERATE_ACCESSORS_PRIORITY = 10;
 	public static int GENERATE_CONSTRUCTORS_PRIORITY = 20;
 	public static int GENERATE_HASHCODE_EQUALS_PRIORITY = 30;