[Issue 275] Allow @Delegate on no-argument methods

3 files changed, 185 insertions, 68 deletions

diff --git a/src/core/lombok/Delegate.java b/src/core/lombok/Delegate.java
index 6d03d649..9ab9acae 100644
--- a/src/core/lombok/Delegate.java
+++ b/src/core/lombok/Delegate.java
@@ -40,7 +40,7 @@ import java.lang.annotation.Target;
  * that exist in {@link Object}, the {@code canEqual(Object)} method, and any methods that appear in types
  * that are listed in the {@code excludes} property.
  */
-@Target(ElementType.FIELD)
+@Target({ElementType.FIELD, ElementType.METHOD})
 @Retention(RetentionPolicy.SOURCE)
 public @interface Delegate {
 	/**
diff --git a/src/core/lombok/javac/handlers/HandleDelegate.java b/src/core/lombok/javac/handlers/HandleDelegate.java
index f6a81474..3674ae5a 100644
--- a/src/core/lombok/javac/handlers/HandleDelegate.java
+++ b/src/core/lombok/javac/handlers/HandleDelegate.java
@@ -94,23 +94,42 @@ public class HandleDelegate extends JavacAnnotationHandler<Delegate> {
 	
 	@Override public void handle(AnnotationValues<Delegate> annotation, JCAnnotation ast, JavacNode annotationNode) {
 		deleteAnnotationIfNeccessary(annotationNode, Delegate.class);
-		if (annotationNode.up().getKind() != Kind.FIELD) {
-			// As the annotation is legal on fields only, javac itself will take care of printing an error message for this.
+		
+		Type delegateType;
+		Name delegateName = annotationNode.toName(annotationNode.up().getName());
+		DelegateReceiver delegateReceiver;
+		JavacResolution reso = new JavacResolution(annotationNode.getContext());
+		if (annotationNode.up().getKind() == Kind.FIELD) {
+			delegateReceiver = DelegateReceiver.FIELD;
+			delegateType = annotationNode.up().get().type;
+			if (delegateType == null) reso.resolveClassMember(annotationNode.up());
+			delegateType = annotationNode.up().get().type;
+		} else if (annotationNode.up().getKind() == Kind.METHOD) {
+			if (!(annotationNode.up().get() instanceof JCMethodDecl)) {
+				annotationNode.addError("@Delegate is legal only on no-argument methods.");
+				return;
+			}
+			JCMethodDecl methodDecl = (JCMethodDecl) annotationNode.up().get();
+			if (!methodDecl.params.isEmpty()) {
+				annotationNode.addError("@Delegate is legal only on no-argument methods.");
+				return;
+			}
+			delegateReceiver = DelegateReceiver.METHOD;
+			delegateType = methodDecl.restype.type;
+			if (delegateType == null) reso.resolveClassMember(annotationNode.up());
+			delegateType = methodDecl.restype.type;
+		} else {
+			// As the annotation is legal on fields and methods only, javac itself will take care of printing an error message for this.
 			return;
 		}
 		
 		List<Object> delegateTypes = annotation.getActualExpressions("types");
 		List<Object> excludeTypes = annotation.getActualExpressions("excludes");
-		JavacResolution reso = new JavacResolution(annotationNode.getContext());
 		List<Type> toDelegate = new ArrayList<Type>();
 		List<Type> toExclude = new ArrayList<Type>();
 		
 		if (delegateTypes.isEmpty()) {
-			Type type = ((JCVariableDecl)annotationNode.up().get()).type;
-			if (type == null) reso.resolveClassMember(annotationNode.up());
-			//TODO I'm fairly sure the above line (and that entire method) does effectively bupkis!
-			type = ((JCVariableDecl)annotationNode.up().get()).type;
-			if (type != null) toDelegate.add(type);
+			if (delegateType != null) toDelegate.add(delegateType); 
 		} else {
 			for (Object dt : delegateTypes) {
 				if (dt instanceof JCFieldAccess && ((JCFieldAccess)dt).name.toString().equals("class")) {
@@ -167,15 +186,13 @@ public class HandleDelegate extends JavacAnnotationHandler<Delegate> {
 			}
 		}
 		
-		Name delegateFieldName = annotationNode.toName(annotationNode.up().getName());
-		
-		for (MethodSig sig : signaturesToDelegate) generateAndAdd(sig, annotationNode, delegateFieldName);
+		for (MethodSig sig : signaturesToDelegate) generateAndAdd(sig, annotationNode, delegateName, delegateReceiver);
 	}
 	
-	private void generateAndAdd(MethodSig sig, JavacNode annotation, Name delegateFieldName) {
+	private void generateAndAdd(MethodSig sig, JavacNode annotation, Name delegateName, DelegateReceiver delegateReceiver) {
 		List<JCMethodDecl> toAdd = new ArrayList<JCMethodDecl>();
 		try {
-			toAdd.add(createDelegateMethod(sig, annotation, delegateFieldName));
+			toAdd.add(createDelegateMethod(sig, annotation, delegateName, delegateReceiver));
 		} catch (TypeNotConvertibleException e) {
 			annotation.addError("Can't create delegate method for " + sig.name + ": " + e.getMessage());
 			return;
@@ -240,7 +257,7 @@ public class HandleDelegate extends JavacAnnotationHandler<Delegate> {
 		}
 	}
 	
-	private JCMethodDecl createDelegateMethod(MethodSig sig, JavacNode annotation, Name delegateFieldName) throws TypeNotConvertibleException, CantMakeDelegates {
+	private JCMethodDecl createDelegateMethod(MethodSig sig, JavacNode annotation, Name delegateName, DelegateReceiver delegateReceiver) throws TypeNotConvertibleException, CantMakeDelegates {
 		/* public <T, U, ...> ReturnType methodName(ParamType1 name1, ParamType2 name2, ...) throws T1, T2, ... {
 		 *      (return) delegate.<T, U>methodName(name1, name2);
 		 *  }
@@ -288,9 +305,7 @@ public class HandleDelegate extends JavacAnnotationHandler<Delegate> {
 			args.append(maker.Ident(name));
 		}
 		
-		JCExpression delegateFieldRef = maker.Select(maker.Ident(annotation.toName("this")), delegateFieldName);
-		
-		JCExpression delegateCall = maker.Apply(toList(typeArgs), maker.Select(delegateFieldRef, sig.name), toList(args));
+		JCExpression delegateCall = maker.Apply(toList(typeArgs), maker.Select(delegateReceiver.get(annotation, delegateName), sig.name), toList(args));
 		JCStatement body = useReturn ? maker.Return(delegateCall) : maker.Exec(delegateCall);
 		JCBlock bodyBlock = maker.Block(0, com.sun.tools.javac.util.List.of(body));
 		
@@ -367,4 +382,22 @@ public class HandleDelegate extends JavacAnnotationHandler<Delegate> {
 		binding = types.erasure(binding);
 		return binding.toString();
 	}
+	
+	private enum DelegateReceiver {
+		METHOD {
+			public JCExpression get(final JavacNode node, final Name name) {
+				com.sun.tools.javac.util.List<JCExpression> nilExprs = com.sun.tools.javac.util.List.nil();
+				final TreeMaker maker = node.getTreeMaker();
+				return maker.Apply(nilExprs, maker.Select(maker.Ident(node.toName("this")), name), nilExprs);
+			}
+		},
+		FIELD {
+			public JCExpression get(final JavacNode node, final Name name) {
+				final TreeMaker maker = node.getTreeMaker();
+				return maker.Select(maker.Ident(node.toName("this")), name);
+			}
+		};
+		
+		public abstract JCExpression get(final JavacNode node, final Name name);
+	}
 }
diff --git a/src/eclipseAgent/lombok/eclipse/agent/PatchDelegate.java b/src/eclipseAgent/lombok/eclipse/agent/PatchDelegate.java
index affcd4f2..d6310de1 100644
--- a/src/eclipseAgent/lombok/eclipse/agent/PatchDelegate.java
+++ b/src/eclipseAgent/lombok/eclipse/agent/PatchDelegate.java
@@ -103,25 +103,25 @@ public class PatchDelegate {
 		return new String(decl.name);
 	}
 	
-	private static boolean hasDelegateMarkedFields(TypeDeclaration decl) {
+	private static boolean hasDelegateMarkedFieldsOrMethods(TypeDeclaration decl) {
 		if (decl.fields != null) for (FieldDeclaration field : decl.fields) {
 			if (field.annotations == null) continue;
 			for (Annotation ann : field.annotations) {
-				if (ann.type == null) continue;
-				TypeBinding tb = ann.type.resolveType(decl.initializerScope);
-				if (tb == null) continue;
-				if (!charArrayEquals("lombok", tb.qualifiedPackageName())) continue;
-				if (!charArrayEquals("Delegate", tb.qualifiedSourceName())) continue;
-				return true;
+				if (isDelegate(ann, decl)) return true;
+			}
+		}
+		if (decl.methods != null) for (AbstractMethodDeclaration method : decl.methods) {
+			if (method.annotations == null) continue;
+			for (Annotation ann : method.annotations) {
+				if (isDelegate(ann, decl)) return true;
 			}
 		}
-		
 		return false;
 	}
 	
 	public static boolean handleDelegateForType(ClassScope scope) {
 		if (TransformEclipseAST.disableLombok) return false;
-		if (!hasDelegateMarkedFields(scope.referenceContext)) return false;
+		if (!hasDelegateMarkedFieldsOrMethods(scope.referenceContext)) return false;
 		
 		List<ClassScopeEntry> stack = visited.get();
 		StringBuilder corrupted = null;
@@ -145,16 +145,23 @@ public class PatchDelegate {
 			stack.add(entry);
 			
 			try {
-				List<BindingTuple> methodsToDelegate = new ArrayList<BindingTuple>();
 				TypeDeclaration decl = scope.referenceContext;
 				if (decl != null) {
 					CompilationUnitDeclaration cud = scope.compilationUnitScope().referenceContext;
 					EclipseAST eclipseAst = TransformEclipseAST.getAST(cud, true);
-					fillMethodBindings(cud, scope, methodsToDelegate);
+					List<BindingTuple> methodsToDelegate = new ArrayList<BindingTuple>();
+					fillMethodBindingsForFields(cud, scope, methodsToDelegate);
+					if (entry.corruptedPath != null) {
+						eclipseAst.get(scope.referenceContext).addError("No @Delegate methods created because there's a loop: " + entry.corruptedPath);
+					} else {
+						generateDelegateMethods(eclipseAst.get(decl), methodsToDelegate, DelegateReceiver.FIELD);
+					}
+					methodsToDelegate.clear();
+					fillMethodBindingsForMethods(cud, scope, methodsToDelegate);
 					if (entry.corruptedPath != null) {
 						eclipseAst.get(scope.referenceContext).addError("No @Delegate methods created because there's a loop: " + entry.corruptedPath);
 					} else {
-						generateDelegateMethods(eclipseAst.get(decl), methodsToDelegate);
+						generateDelegateMethods(eclipseAst.get(decl), methodsToDelegate, DelegateReceiver.METHOD);
 					}
 				}
 			} finally {
@@ -181,43 +188,18 @@ public class PatchDelegate {
 	private static Map<ASTNode, Object> alreadyApplied = new WeakHashMap<ASTNode, Object>();
 	private static final Object MARKER = new Object();
 	
-	private static void fillMethodBindings(CompilationUnitDeclaration cud, ClassScope scope, List<BindingTuple> methodsToDelegate) {
+	private static void fillMethodBindingsForFields(CompilationUnitDeclaration cud, ClassScope scope, List<BindingTuple> methodsToDelegate) {
 		TypeDeclaration decl = scope.referenceContext;
 		if (decl == null) return;
 		
 		if (decl.fields != null) for (FieldDeclaration field : decl.fields) {
 			if (field.annotations == null) continue;
 			for (Annotation ann : field.annotations) {
-				if (ann.type == null) continue;
-				TypeBinding tb = ann.type.resolveType(decl.initializerScope);
-				if (!charArrayEquals("lombok", tb.qualifiedPackageName())) continue;
-				if (!charArrayEquals("Delegate", tb.qualifiedSourceName())) continue;
+				if (!isDelegate(ann, decl)) continue;
 				if (alreadyApplied.put(ann, MARKER) == MARKER) continue;
 				
-				List<ClassLiteralAccess> rawTypes = new ArrayList<ClassLiteralAccess>();
-				List<ClassLiteralAccess> excludedRawTypes = new ArrayList<ClassLiteralAccess>();
-				for (MemberValuePair pair : ann.memberValuePairs()) {
-					if (charArrayEquals("types", pair.name)) {
-						if (pair.value instanceof ArrayInitializer) {
-							for (Expression expr : ((ArrayInitializer)pair.value).expressions) {
-								if (expr instanceof ClassLiteralAccess) rawTypes.add((ClassLiteralAccess) expr);
-							}
-						}
-						if (pair.value instanceof ClassLiteralAccess) {
-							rawTypes.add((ClassLiteralAccess) pair.value);
-						}
-					}
-					if (charArrayEquals("excludes", pair.name)) {
-						if (pair.value instanceof ArrayInitializer) {
-							for (Expression expr : ((ArrayInitializer)pair.value).expressions) {
-								if (expr instanceof ClassLiteralAccess) excludedRawTypes.add((ClassLiteralAccess) expr);
-							}
-						}
-						if (pair.value instanceof ClassLiteralAccess) {
-							excludedRawTypes.add((ClassLiteralAccess) pair.value);
-						}
-					}
-				}
+				List<ClassLiteralAccess> rawTypes = rawTypes(ann, "types");
+				List<ClassLiteralAccess> excludedRawTypes = rawTypes(ann, "excludes");
 				
 				List<BindingTuple> methodsToExclude = new ArrayList<BindingTuple>();
 				for (ClassLiteralAccess cla : excludedRawTypes) {
@@ -251,6 +233,88 @@ public class PatchDelegate {
 		}
 	}
 	
+	private static void fillMethodBindingsForMethods(CompilationUnitDeclaration cud, ClassScope scope, List<BindingTuple> methodsToDelegate) {
+		TypeDeclaration decl = scope.referenceContext;
+		if (decl == null) return;
+		
+		if (decl.methods != null) for (AbstractMethodDeclaration methodDecl : decl.methods) {
+			if (methodDecl.annotations == null) continue;
+			for (Annotation ann : methodDecl.annotations) {
+				if (!isDelegate(ann, decl)) continue;
+				if (alreadyApplied.put(ann, MARKER) == MARKER) continue;
+				if (!(methodDecl instanceof MethodDeclaration)) {
+					EclipseAST eclipseAst = TransformEclipseAST.getAST(cud, true);
+					eclipseAst.get(ann).addError("@Delegate is legal only on no-argument methods.");
+					continue;
+				}
+				if (methodDecl.arguments != null) {
+					EclipseAST eclipseAst = TransformEclipseAST.getAST(cud, true);
+					eclipseAst.get(ann).addError("@Delegate is legal only on no-argument methods.");
+					continue;
+				}
+				MethodDeclaration method = (MethodDeclaration) methodDecl;
+				
+				List<ClassLiteralAccess> rawTypes = rawTypes(ann, "types");
+				List<ClassLiteralAccess> excludedRawTypes = rawTypes(ann, "excludes");
+				
+				List<BindingTuple> methodsToExclude = new ArrayList<BindingTuple>();
+				for (ClassLiteralAccess cla : excludedRawTypes) {
+					addAllMethodBindings(methodsToExclude, cla.type.resolveType(decl.initializerScope), new HashSet<String>(), method.selector, ann);
+				}
+				
+				Set<String> banList = new HashSet<String>();
+				for (BindingTuple excluded : methodsToExclude) banList.add(printSig(excluded.parameterized));
+				
+				List<BindingTuple> methodsToDelegateForThisAnn = new ArrayList<BindingTuple>();
+				
+				if (rawTypes.isEmpty()) {
+					addAllMethodBindings(methodsToDelegateForThisAnn, method.returnType.resolveType(decl.initializerScope), banList, method.selector, ann);
+				} else {
+					for (ClassLiteralAccess cla : rawTypes) {
+						addAllMethodBindings(methodsToDelegateForThisAnn, cla.type.resolveType(decl.initializerScope), banList, method.selector, ann);
+					}
+				}
+				
+				// Not doing this right now because of problems - see commented-out-method for info.
+				// removeExistingMethods(methodsToDelegate, decl, scope);
+				
+				String dupe = containsDuplicates(methodsToDelegateForThisAnn);
+				if (dupe != null) {
+					EclipseAST eclipseAst = TransformEclipseAST.getAST(cud, true);
+					eclipseAst.get(ann).addError("The method '" + dupe + "' is being delegated by more than one specified type.");
+				} else {
+					methodsToDelegate.addAll(methodsToDelegateForThisAnn);
+				}
+			}
+		}
+	}
+	
+	private static boolean isDelegate(Annotation ann, TypeDeclaration decl) {
+		if (ann.type == null) return false;
+		TypeBinding tb = ann.type.resolveType(decl.initializerScope);
+		if (tb == null) return false;
+		if (!charArrayEquals("lombok", tb.qualifiedPackageName())) return false;
+		if (!charArrayEquals("Delegate", tb.qualifiedSourceName())) return false;
+		return true;
+	}
+	
+	private static List<ClassLiteralAccess> rawTypes(Annotation ann, String name) {
+		List<ClassLiteralAccess> rawTypes = new ArrayList<ClassLiteralAccess>();
+		for (MemberValuePair pair : ann.memberValuePairs()) {
+			if (charArrayEquals(name, pair.name)) {
+				if (pair.value instanceof ArrayInitializer) {
+					for (Expression expr : ((ArrayInitializer)pair.value).expressions) {
+						if (expr instanceof ClassLiteralAccess) rawTypes.add((ClassLiteralAccess) expr);
+					}
+				}
+				if (pair.value instanceof ClassLiteralAccess) {
+					rawTypes.add((ClassLiteralAccess) pair.value);
+				}
+			}
+		}
+		return rawTypes;
+	}
+	
 	/*
 	 * We may someday finish this method. Steps to be completed:
 	 * 
@@ -321,11 +385,11 @@ public class PatchDelegate {
 //		}
 //	}
 	
-	private static void generateDelegateMethods(EclipseNode typeNode, List<BindingTuple> methods) {
+	private static void generateDelegateMethods(EclipseNode typeNode, List<BindingTuple> methods, DelegateReceiver delegateReceiver) {
 		CompilationUnitDeclaration top = (CompilationUnitDeclaration) typeNode.top().get();
 		for (BindingTuple pair : methods) {
 			EclipseNode annNode = typeNode.getAst().get(pair.responsible);
-			MethodDeclaration method = createDelegateMethod(pair.fieldName, typeNode, pair, top.compilationResult, annNode);
+			MethodDeclaration method = createDelegateMethod(pair.fieldName, typeNode, pair, top.compilationResult, annNode, delegateReceiver);
 			if (method != null) { 
 				SetGeneratedByVisitor visitor = new SetGeneratedByVisitor(annNode.get());
 				method.traverse(visitor, ((TypeDeclaration)typeNode.get()).scope);
@@ -473,7 +537,7 @@ public class PatchDelegate {
 		}
 	}
 	
-	private static MethodDeclaration createDelegateMethod(char[] name, EclipseNode typeNode, BindingTuple pair, CompilationResult compilationResult, EclipseNode annNode) {
+	private static MethodDeclaration createDelegateMethod(char[] name, EclipseNode typeNode, BindingTuple pair, CompilationResult compilationResult, EclipseNode annNode, DelegateReceiver delegateReceiver) {
 		/* public <T, U, ...> ReturnType methodName(ParamType1 name1, ParamType2 name2, ...) throws T1, T2, ... {
 		 *      (return) delegate.<T, U>methodName(name1, name2);
 		 *  }
@@ -514,11 +578,7 @@ public class PatchDelegate {
 		call.sourceStart = pS; call.sourceEnd = pE;
 		call.nameSourcePosition = pos(source);
 		setGeneratedBy(call, source);
-		FieldReference fieldRef = new FieldReference(name, pos(source));
-		fieldRef.receiver = new ThisReference(pS, pE);
-		setGeneratedBy(fieldRef, source);
-		setGeneratedBy(fieldRef.receiver, source);
-		call.receiver = fieldRef;
+		call.receiver = delegateReceiver.get(source, name);
 		call.selector = binding.selector;
 		
 		if (binding.typeVariables != null && binding.typeVariables.length > 0) {
@@ -742,7 +802,31 @@ public class PatchDelegate {
 		if (s.length() != c.length) return false;
 		for (int i = 0; i < s.length(); i++) if (s.charAt(i) != c[i]) return false;
 		return true;
+	}
+	
+	private enum DelegateReceiver {
+		METHOD {
+			public Expression get(final ASTNode source, char[] name) {
+				MessageSend call = new MessageSend();
+				call.sourceStart = source.sourceStart; call.sourceEnd = source.sourceEnd;
+				call.nameSourcePosition = pos(source);
+				setGeneratedBy(call, source);
+				call.selector = name;
+				call.receiver = new ThisReference(source.sourceStart, source.sourceEnd);
+				setGeneratedBy(call.receiver, source);
+				return call;
+			}
+		},
+		FIELD {
+			public Expression get(final ASTNode source, char[] name) {
+				FieldReference fieldRef = new FieldReference(name, pos(source));
+				setGeneratedBy(fieldRef, source);
+				fieldRef.receiver = new ThisReference(source.sourceStart, source.sourceEnd);
+				setGeneratedBy(fieldRef.receiver, source);
+				return fieldRef;
+			}
+		};
 		
-		
+		public abstract Expression get(final ASTNode source, char[] name);
 	}
 }