Bug 569955 - Consider extending the "Annotate" assist to source editors

Change-Id: Ib6b956edcb7052783e749b94dc2e3cfb4eea9f13
diff --git a/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AbstractAnnotateAssistTests.java b/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AbstractAnnotateAssistTests.java
index 2e8ac71..3eb9ec3 100644
--- a/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AbstractAnnotateAssistTests.java
+++ b/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AbstractAnnotateAssistTests.java
@@ -22,6 +22,7 @@
 import java.io.InputStream;
 import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import org.eclipse.core.runtime.CoreException;
 import org.eclipse.core.runtime.Path;
@@ -36,8 +37,12 @@
 import org.eclipse.jface.text.contentassist.ICompletionProposal;
 import org.eclipse.jface.text.quickassist.IQuickAssistProcessor;
 
+import org.eclipse.ui.part.FileEditorInput;
+
 import org.eclipse.jdt.core.IClasspathAttribute;
 import org.eclipse.jdt.core.IClasspathEntry;
+import org.eclipse.jdt.core.ICompilationUnit;
+import org.eclipse.jdt.core.IJavaElement;
 import org.eclipse.jdt.core.IJavaProject;
 import org.eclipse.jdt.core.JavaCore;
 
@@ -46,8 +51,10 @@
 
 import org.eclipse.jdt.ui.tests.quickfix.JarUtil.ClassFileFilter;
 
+import org.eclipse.jdt.internal.ui.javaeditor.ClassFileEditor;
 import org.eclipse.jdt.internal.ui.javaeditor.JavaEditor;
 import org.eclipse.jdt.internal.ui.javaeditor.JavaSourceViewer;
+import org.eclipse.jdt.internal.ui.text.correction.AssistContext;
 import org.eclipse.jdt.internal.ui.text.correction.ExternalNullAnnotationQuickAssistProcessor;
 import org.eclipse.jdt.internal.ui.text.correction.JavaCorrectionAssistant;
 
@@ -63,16 +70,23 @@
 			((IFolder)parent).create(true, true, null);
 	}
 
-	public List<ICompletionProposal> collectAnnotateProposals(JavaEditor javaEditor, int offset) {
-		JavaSourceViewer viewer= (JavaSourceViewer)javaEditor.getViewer();
-		viewer.setSelection(new TextSelection(offset, 0));
+	public List<ICompletionProposal> collectAnnotateProposals(JavaEditor javaEditor, int offset) throws CoreException {
+		if (javaEditor instanceof ClassFileEditor) {
+			JavaSourceViewer viewer= (JavaSourceViewer)javaEditor.getViewer();
+			viewer.setSelection(new TextSelection(offset, 0));
 
-		JavaCorrectionAssistant correctionAssist= new JavaCorrectionAssistant(javaEditor);
-		IQuickAssistProcessor assistProcessor= new ExternalNullAnnotationQuickAssistProcessor(correctionAssist);
-		ICompletionProposal[] proposals= assistProcessor.computeQuickAssistProposals(viewer.getQuickAssistInvocationContext());
+			JavaCorrectionAssistant correctionAssist= new JavaCorrectionAssistant(javaEditor);
+			IQuickAssistProcessor assistProcessor= new ExternalNullAnnotationQuickAssistProcessor(correctionAssist);
+			ICompletionProposal[] proposals= assistProcessor.computeQuickAssistProposals(viewer.getQuickAssistInvocationContext());
 
-		List<ICompletionProposal> list= Arrays.asList(proposals);
-		return list;
+			List<ICompletionProposal> list= Arrays.asList(proposals);
+			return list;
+		} else {
+			// for source compilation units:
+			IJavaElement element= JavaCore.create(((FileEditorInput) javaEditor.getEditorInput()).getFile());
+			AssistContext context= new AssistContext((ICompilationUnit) element, offset, 0);
+			return collectAssists(context, false).stream().map(ICompletionProposal.class::cast).collect(Collectors.toList());
+		}
 	}
 
 	// === from jdt.core.tests.model: ===
diff --git a/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AnnotateAssistTest1d8.java b/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AnnotateAssistTest1d8.java
index 7902b30..36e965f 100644
--- a/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AnnotateAssistTest1d8.java
+++ b/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/AnnotateAssistTest1d8.java
@@ -29,10 +29,13 @@
 import org.eclipse.core.runtime.Path;
 
 import org.eclipse.core.resources.IFile;
+import org.eclipse.core.resources.IResource;
 
 import org.eclipse.jface.text.IDocument;
 import org.eclipse.jface.text.contentassist.ICompletionProposal;
 
+import org.eclipse.jdt.core.IClasspathAttribute;
+import org.eclipse.jdt.core.IClasspathEntry;
 import org.eclipse.jdt.core.IType;
 import org.eclipse.jdt.core.JavaCore;
 
@@ -1169,4 +1172,87 @@
 			JavaPlugin.getActivePage().closeAllEditors(false);
 		}
 	}
+
+	/**
+	 * Assert two proposals ("@NonNull" and "@Nullable") on a simple return type (type variable).
+	 * Apply the second proposal and check the effect.
+	 *
+	 * Similar to AnnotateAssistTest1d5.testAnnotateReturn2() but annotating source not binary.
+	 *
+	 * @throws Exception Any exception
+	 */
+	@Test
+	public void testAnnotateReturnInSourceFolder() throws Exception {
+		String MY_MAP_PATH= "pack/age/MyMap";
+		String[] pathAndContents= new String[] {
+					MY_MAP_PATH+".java",
+					"package pack.age;\n" +
+					"public interface MyMap<K,V> {\n" +
+					"    public V get(K key);\n" +
+					"}\n"
+				};
+		JarUtil.createSourceDir(pathAndContents, fJProject1.getProject().getLocation()+"/src");
+		fJProject1.getProject().refreshLocal(IResource.DEPTH_INFINITE, null);
+
+		IClasspathEntry[] rawClasspath= fJProject1.getRawClasspath();
+		fJProject1.setRawClasspath(new IClasspathEntry[] {
+				rawClasspath[0],
+				JavaCore.newSourceEntry(fJProject1.getPath().append("src"))
+			},
+			null);
+		IType type= fJProject1.findType(MY_MAP_PATH.replace('/', '.'));
+		JavaEditor javaEditor= (JavaEditor) JavaUI.openInEditor(type);
+
+		try {
+			int offset= pathAndContents[1].indexOf("V get");
+
+			List<ICompletionProposal> list= collectAnnotateProposals(javaEditor, offset);
+			assertNumberOfProposals(list, 0); // no annotation path defined
+
+			fJProject1.setRawClasspath(new IClasspathEntry[] {
+					rawClasspath[0],
+					JavaCore.newSourceEntry(fJProject1.getPath().append("src"), null, null, null,
+							new IClasspathAttribute[] {
+									JavaCore.newClasspathAttribute(IClasspathAttribute.EXTERNAL_ANNOTATION_PATH, ANNOTATION_PATH)
+					})
+				},
+				null);
+
+			list= collectAnnotateProposals(javaEditor, offset);
+			assertCorrectLabels(list);
+			assertNumberOfProposals(list, 2);
+
+			ICompletionProposal proposal= findProposalByName("Annotate as '@NonNull V'", list);
+			String expectedInfo=
+					"<dl><dt>get</dt>" +
+					"<dd>(TK;)TV;</dd>" +
+					"<dd>(TK;)T<b>1</b>V;</dd>" + // <= 1
+					"</dl>";
+			assertEquals("expect detail", expectedInfo, proposal.getAdditionalProposalInfo());
+
+			proposal= findProposalByName("Annotate as '@Nullable V'", list);
+			expectedInfo=
+					"<dl><dt>get</dt>" +
+					"<dd>(TK;)TV;</dd>" +
+					"<dd>(TK;)T<b>0</b>V;</dd>" + // <= 0
+					"</dl>";
+			assertEquals("expect detail", expectedInfo, proposal.getAdditionalProposalInfo());
+
+			IDocument document= javaEditor.getDocumentProvider().getDocument(javaEditor.getEditorInput());
+			proposal.apply(document);
+
+			IFile annotationFile= fJProject1.getProject().getFile(new Path(ANNOTATION_PATH).append(MY_MAP_PATH + ".eea"));
+			assertTrue("Annotation file should have been created", annotationFile.exists());
+
+			String expectedContent=
+					"class pack/age/MyMap\n" +
+					"get\n" +
+					" (TK;)TV;\n" +
+					" (TK;)T0V;\n";
+			checkContentOfFile("annotation file content", annotationFile, expectedContent);
+		} finally {
+			JavaPlugin.getActivePage().closeAllEditors(false);
+		}
+	}
+
 }
diff --git a/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/QuickFixTest.java b/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/QuickFixTest.java
index 240ad98..a8ef0ca 100644
--- a/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/QuickFixTest.java
+++ b/org.eclipse.jdt.ui.tests/ui/org/eclipse/jdt/ui/tests/quickfix/QuickFixTest.java
@@ -357,7 +357,7 @@
 		return false;
 	}
 
-	protected static final ArrayList<IJavaCompletionProposal> collectAssists(IInvocationContext context, boolean includeLinkedRename) throws CoreException {
+	public static final ArrayList<IJavaCompletionProposal> collectAssists(IInvocationContext context, boolean includeLinkedRename) throws CoreException {
 		Class<?>[] filteredTypes= includeLinkedRename ? null : new Class[] { LinkedNamesAssistProposal.class, RenameRefactoringProposal.class };
 		return collectAssists(context, filteredTypes);
 	}
diff --git a/org.eclipse.jdt.ui/core extension/org/eclipse/jdt/internal/corext/fix/ExternalNullAnnotationChangeProposals.java b/org.eclipse.jdt.ui/core extension/org/eclipse/jdt/internal/corext/fix/ExternalNullAnnotationChangeProposals.java
index 0f2f756..7b8e0b9 100644
--- a/org.eclipse.jdt.ui/core extension/org/eclipse/jdt/internal/corext/fix/ExternalNullAnnotationChangeProposals.java
+++ b/org.eclipse.jdt.ui/core extension/org/eclipse/jdt/internal/corext/fix/ExternalNullAnnotationChangeProposals.java
@@ -28,7 +28,6 @@
 
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
-import java.util.ArrayList;
 import java.util.List;
 
 import org.eclipse.swt.graphics.Image;
@@ -376,7 +375,7 @@
 	}
 
 	/* Quick assist on class file, propose changes on any type detail. */
-	public static void collectExternalAnnotationProposals(ICompilationUnit cu, ASTNode coveringNode, int offset, ArrayList<IJavaCompletionProposal> resultingCollection) {
+	public static void collectExternalAnnotationProposals(ICompilationUnit cu, ASTNode coveringNode, int offset, List<IJavaCompletionProposal> resultingCollection) {
 
 		IJavaProject javaProject= cu.getJavaProject();
 		if (JavaCore.DISABLED.equals(javaProject.getOption(JavaCore.COMPILER_ANNOTATION_NULL_ANALYSIS, true)))
@@ -642,7 +641,7 @@
 	/* Create one proposal from each of the three given renderers. */
 	static void createProposalsForType(ICompilationUnit cu, ASTNode type, int dims,
 			int outerDims, boolean annotateVarargs, int offset,
-			TypeRenderer rendererNonNull, TypeRenderer rendererNullable, TypeRenderer rendererRemove, ProposalCreator creator, ArrayList<IJavaCompletionProposal> resultingCollection) {
+			TypeRenderer rendererNonNull, TypeRenderer rendererNullable, TypeRenderer rendererRemove, ProposalCreator creator, List<IJavaCompletionProposal> resultingCollection) {
 		SignatureAnnotationChangeProposal operation;
 		String label;
 		// propose adding @NonNull:
diff --git a/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/AdvancedQuickAssistProcessor.java b/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/AdvancedQuickAssistProcessor.java
index f061ee5..c484662 100644
--- a/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/AdvancedQuickAssistProcessor.java
+++ b/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/AdvancedQuickAssistProcessor.java
@@ -165,7 +165,8 @@
 					|| getJoinIfListInIfElseIfProposals(context, coveringNode, coveredNodes, null)
 					|| getConvertSwitchToIfProposals(context, coveringNode, null)
 					|| getConvertIfElseToSwitchProposals(context, coveringNode, null)
-					|| GetterSetterCorrectionSubProcessor.addGetterSetterProposal(context, coveringNode, null, null);
+					|| GetterSetterCorrectionSubProcessor.addGetterSetterProposal(context, coveringNode, null, null)
+					|| ExternalNullAnnotationQuickAssistProcessor.canAssist(context);
 		}
 		return false;
 	}
@@ -208,6 +209,8 @@
 				getConvertSwitchToIfProposals(context, coveringNode, resultingCollections);
 				getConvertIfElseToSwitchProposals(context, coveringNode, resultingCollections);
 				GetterSetterCorrectionSubProcessor.addGetterSetterProposal(context, coveringNode, locations, resultingCollections);
+
+				ExternalNullAnnotationQuickAssistProcessor.getAnnotateProposals(context, resultingCollections);
 			}
 
 			return resultingCollections.toArray(new IJavaCompletionProposal[resultingCollections.size()]);
diff --git a/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/ExternalNullAnnotationQuickAssistProcessor.java b/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/ExternalNullAnnotationQuickAssistProcessor.java
index c7727ff..7536a15 100644
--- a/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/ExternalNullAnnotationQuickAssistProcessor.java
+++ b/org.eclipse.jdt.ui/ui/org/eclipse/jdt/internal/ui/text/correction/ExternalNullAnnotationQuickAssistProcessor.java
@@ -14,6 +14,7 @@
 package org.eclipse.jdt.internal.ui.text.correction;
 
 import java.util.ArrayList;
+import java.util.List;
 
 import org.eclipse.jface.text.contentassist.ICompletionProposal;
 import org.eclipse.jface.text.quickassist.IQuickAssistInvocationContext;
@@ -25,13 +26,16 @@
 
 import org.eclipse.jdt.core.IClassFile;
 import org.eclipse.jdt.core.ICompilationUnit;
+import org.eclipse.jdt.core.IJavaProject;
 import org.eclipse.jdt.core.JavaCore;
 import org.eclipse.jdt.core.JavaModelException;
 import org.eclipse.jdt.core.WorkingCopyOwner;
 
 import org.eclipse.jdt.internal.corext.fix.ExternalNullAnnotationChangeProposals;
 
+import org.eclipse.jdt.ui.text.java.IInvocationContext;
 import org.eclipse.jdt.ui.text.java.IJavaCompletionProposal;
+import org.eclipse.jdt.ui.text.java.correction.ICommandAccess;
 
 import org.eclipse.jdt.internal.ui.JavaPlugin;
 import org.eclipse.jdt.internal.ui.javaeditor.IClassFileEditorInput;
@@ -112,4 +116,27 @@
 			}
 		}
 	}
+
+	/**
+	 * API for regular assist processor working on a source file / compilation unit.
+	 * @param context the invocation context
+	 * @param proposals list of proposals to which new proposals may be added
+	 */
+	public static void getAnnotateProposals(IInvocationContext context, List<ICommandAccess> proposals) {
+		List<IJavaCompletionProposal> eeaList= new ArrayList<>();
+		ExternalNullAnnotationChangeProposals.collectExternalAnnotationProposals(context.getCompilationUnit(),
+				context.getCoveringNode(), context.getSelectionOffset(), eeaList);
+		for (IJavaCompletionProposal proposal : eeaList)
+			proposals.add((ICommandAccess) proposal);
+	}
+
+	public static boolean canAssist(IInvocationContext context) {
+		ICompilationUnit cu= context.getCompilationUnit();
+		if (cu != null) {
+			IJavaProject javaProject= cu.getJavaProject();
+			return javaProject.getOption(JavaCore.COMPILER_ANNOTATION_NULL_ANALYSIS, true).equals(JavaCore.ENABLED)
+					&& ExternalNullAnnotationChangeProposals.hasAnnotationPathInWorkspace(javaProject, cu);
+		}
+		return false;
+	}
 }