correct return type and statement for @Singular

5 files changed, 36 insertions, 46 deletions

diff --git a/src/core/lombok/javac/handlers/HandleSuperBuilder.java b/src/core/lombok/javac/handlers/HandleSuperBuilder.java
index 0bd9ee2e..1228dfd3 100644
--- a/src/core/lombok/javac/handlers/HandleSuperBuilder.java
+++ b/src/core/lombok/javac/handlers/HandleSuperBuilder.java
@@ -497,15 +497,17 @@ public class HandleSuperBuilder extends JavacAnnotationHandler<SuperBuilder> {
 	
 	public void makeSetterMethodsForBuilder(JavacNode builderType, BuilderFieldData fieldNode, JavacNode source) {
 		boolean deprecate = isFieldDeprecated(fieldNode.originalFieldNode);
+		JavacTreeMaker maker = builderType.getTreeMaker();
+		JCExpression returnType = maker.Ident(builderType.toName("B"));
+		JCReturn returnStatement = maker.Return(maker.Apply(List.<JCExpression>nil(), maker.Ident(builderType.toName(SELF_METHOD)), List.<JCExpression>nil()));
 		if (fieldNode.singularData == null || fieldNode.singularData.getSingularizer() == null) {
-			makeSimpleSetterMethodForBuilder(builderType, deprecate, fieldNode.createdFields.get(0), fieldNode.nameOfSetFlag, source, true, true);
+			makeSimpleSetterMethodForBuilder(builderType, deprecate, fieldNode.createdFields.get(0), fieldNode.nameOfSetFlag, source, true, true, returnType, returnStatement);
 		} else {
-			// TODO: Fix singular methods to return self().
-			fieldNode.singularData.getSingularizer().generateMethods(fieldNode.singularData, deprecate, builderType, source.get(), true, true);
+			fieldNode.singularData.getSingularizer().generateMethods(fieldNode.singularData, deprecate, builderType, source.get(), true, returnType, returnStatement);
 		}
 	}
 	
-	private void makeSimpleSetterMethodForBuilder(JavacNode builderType, boolean deprecate, JavacNode fieldNode, Name nameOfSetFlag, JavacNode source, boolean fluent, boolean chain) {
+	private void makeSimpleSetterMethodForBuilder(JavacNode builderType, boolean deprecate, JavacNode fieldNode, Name nameOfSetFlag, JavacNode source, boolean fluent, boolean chain, JCExpression returnType, JCReturn returnStatement) {
 		Name fieldName = ((JCVariableDecl) fieldNode.get()).name;
 		
 		for (JavacNode child : builderType.down()) {
@@ -519,9 +521,6 @@ public class HandleSuperBuilder extends JavacAnnotationHandler<SuperBuilder> {
 		
 		JavacTreeMaker maker = fieldNode.getTreeMaker();
 		
-		JCExpression returnType = maker.Ident(builderType.toName("B"));
-		JCReturn returnStatement = maker.Return(maker.Apply(List.<JCExpression>nil(), maker.Ident(builderType.toName(SELF_METHOD)), List.<JCExpression>nil()));
-
 		JCMethodDecl newMethod = HandleSetter.createSetter(Flags.PUBLIC, deprecate, fieldNode, maker, setterName, nameOfSetFlag, returnType, returnStatement, source, List.<JCAnnotation>nil(), List.<JCAnnotation>nil());
 		
 		injectMethod(builderType, newMethod);
@@ -557,6 +556,7 @@ public class HandleSuperBuilder extends JavacAnnotationHandler<SuperBuilder> {
 		JCExpression extending = null;
 		if (parentBuilderClass != null) {
 			// If the annotated class extends another class, we want this builder to extend the builder of the superclass.
+			// FIXME: The extends clause should look like "Parent.ParentBuilder", not just "ParentBuilder" (risk of name clashes).
 			extending = maker.TypeApply(maker.Ident(tdParent.toName(parentBuilderClass)), 
 					List.<JCExpression>of(maker.Ident(tdParent.toName("C")), maker.Ident(tdParent.toName("B"))));
 			// TODO: type params from annotated class
diff --git a/src/core/lombok/javac/handlers/JavacSingularsRecipes.java b/src/core/lombok/javac/handlers/JavacSingularsRecipes.java
index 4279ec63..ea5d1a54 100644
--- a/src/core/lombok/javac/handlers/JavacSingularsRecipes.java
+++ b/src/core/lombok/javac/handlers/JavacSingularsRecipes.java
@@ -194,7 +194,30 @@ public class JavacSingularsRecipes {
 		}
 		
 		public abstract java.util.List<JavacNode> generateFields(SingularData data, JavacNode builderType, JCTree source);
-		public abstract void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, boolean chain);
+		
+		/**
+		 * Generates the singular, plural, and clear methods for the given
+		 * {@link SingularData}.<br>
+		 * Uses the given <code>builderType</code> as return type if
+		 * <code>chain == true</code>, <code>void</code> otherwise.. If you need more
+		 * control over the return type and value, use
+		 * {@link #generateMethods(SingularData, boolean, JavacNode, JCTree, boolean, JCExpression, JCStatement)}.
+		 */
+		public void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, boolean chain) {
+			JavacTreeMaker maker = builderType.getTreeMaker();
+			JCExpression returnType = chain ?
+					cloneSelfType(builderType) : 
+						maker.Type(createVoidType(builderType.getSymbolTable(), CTC_VOID));
+			JCStatement returnStatement = chain ? maker.Return(maker.Ident(builderType.toName("this"))) : null;
+			generateMethods(data, deprecate, builderType, source, fluent, returnType, returnStatement);
+		}
+		/**
+		 * Generates the singular, plural, and clear methods for the given
+		 * {@link SingularData}.<br>
+		 * Uses the given <code>returnType</code> and <code>returnStatement</code> for the generated methods. 
+		 */
+		public abstract void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, JCExpression returnType, JCStatement returnStatement);
+		
 		public abstract void appendBuildCode(SingularData data, JavacNode builderType, JCTree source, ListBuffer<JCStatement> statements, Name targetVariableName, String builderVariable);
 		
 		public boolean requiresCleaning() {
diff --git a/src/core/lombok/javac/handlers/singulars/JavacGuavaSingularizer.java b/src/core/lombok/javac/handlers/singulars/JavacGuavaSingularizer.java
index 1798d52c..9a6a46f2 100644
--- a/src/core/lombok/javac/handlers/singulars/JavacGuavaSingularizer.java
+++ b/src/core/lombok/javac/handlers/singulars/JavacGuavaSingularizer.java
@@ -36,7 +36,6 @@ import lombok.javac.handlers.JavacSingularsRecipes.JavacSingularizer;
 import lombok.javac.handlers.JavacSingularsRecipes.SingularData;
 
 import com.sun.tools.javac.code.Flags;
-import com.sun.tools.javac.code.Symtab;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
@@ -70,19 +69,10 @@ abstract class JavacGuavaSingularizer extends JavacSingularizer {
 		return Collections.singletonList(injectFieldAndMarkGenerated(builderType, buildField));
 	}
 	
-	@Override public void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, boolean chain) {
+	@Override public void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, JCExpression returnType, JCStatement returnStatement) {
 		JavacTreeMaker maker = builderType.getTreeMaker();
-		Symtab symbolTable = builderType.getSymbolTable();
-		JCExpression returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		JCStatement returnStatement = chain ? maker.Return(maker.Ident(builderType.toName("this"))) : null;
 		generateSingularMethod(deprecate, maker, returnType, returnStatement, data, builderType, source, fluent);
-
-		returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		returnStatement = chain ? maker.Return(maker.Ident(builderType.toName("this"))) : null;
 		generatePluralMethod(deprecate, maker, returnType, returnStatement, data, builderType, source, fluent);
-		
-		returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		returnStatement = chain ? maker.Return(maker.Ident(builderType.toName("this"))) : null;
 		generateClearMethod(deprecate, maker, returnType, returnStatement, data, builderType, source);
 	}
 	
diff --git a/src/core/lombok/javac/handlers/singulars/JavacJavaUtilListSetSingularizer.java b/src/core/lombok/javac/handlers/singulars/JavacJavaUtilListSetSingularizer.java
index 196ce45d..dfffb06b 100644
--- a/src/core/lombok/javac/handlers/singulars/JavacJavaUtilListSetSingularizer.java
+++ b/src/core/lombok/javac/handlers/singulars/JavacJavaUtilListSetSingularizer.java
@@ -33,7 +33,6 @@ import lombok.javac.handlers.JavacHandlerUtil;
 import lombok.javac.handlers.JavacSingularsRecipes.SingularData;
 
 import com.sun.tools.javac.code.Flags;
-import com.sun.tools.javac.code.Symtab;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
@@ -76,26 +75,15 @@ abstract class JavacJavaUtilListSetSingularizer extends JavacJavaUtilSingularize
 		return Collections.singletonList(injectFieldAndMarkGenerated(builderType, buildField));
 	}
 	
-	@Override public void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, boolean chain) {
+	@Override public void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, JCExpression returnType, JCStatement returnStatement) {
 		if (useGuavaInstead(builderType)) {
-			guavaListSetSingularizer.generateMethods(data, deprecate, builderType, source, fluent, chain);
+			guavaListSetSingularizer.generateMethods(data, deprecate, builderType, source, fluent, returnType, returnStatement);
 			return;
 		}
 		
 		JavacTreeMaker maker = builderType.getTreeMaker();
-		Symtab symbolTable = builderType.getSymbolTable();
-		Name thisName = builderType.toName("this");
-		
-		JCExpression returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		JCStatement returnStatement = chain ? maker.Return(maker.Ident(thisName)) : null;
 		generateSingularMethod(deprecate, maker, returnType, returnStatement, data, builderType, source, fluent);
-		
-		returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		returnStatement = chain ? maker.Return(maker.Ident(thisName)) : null;
 		generatePluralMethod(deprecate, maker, returnType, returnStatement, data, builderType, source, fluent);
-		
-		returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		returnStatement = chain ? maker.Return(maker.Ident(thisName)) : null;
 		generateClearMethod(deprecate, maker, returnType, returnStatement, data, builderType, source);
 	}
 	
diff --git a/src/core/lombok/javac/handlers/singulars/JavacJavaUtilMapSingularizer.java b/src/core/lombok/javac/handlers/singulars/JavacJavaUtilMapSingularizer.java
index fbcf776e..930ec3a4 100644
--- a/src/core/lombok/javac/handlers/singulars/JavacJavaUtilMapSingularizer.java
+++ b/src/core/lombok/javac/handlers/singulars/JavacJavaUtilMapSingularizer.java
@@ -37,7 +37,6 @@ import lombok.javac.handlers.JavacSingularsRecipes.SingularData;
 import org.mangosdk.spi.ProviderFor;
 
 import com.sun.tools.javac.code.Flags;
-import com.sun.tools.javac.code.Symtab;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
 import com.sun.tools.javac.tree.JCTree.JCExpression;
@@ -101,25 +100,15 @@ public class JavacJavaUtilMapSingularizer extends JavacJavaUtilSingularizer {
 		return Arrays.asList(keyFieldNode, valueFieldNode);
 	}
 	
-	@Override public void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, boolean chain) {
+	@Override public void generateMethods(SingularData data, boolean deprecate, JavacNode builderType, JCTree source, boolean fluent, JCExpression returnType, JCStatement returnStatement) {
 		if (useGuavaInstead(builderType)) {
-			guavaMapSingularizer.generateMethods(data, deprecate, builderType, source, fluent, chain);
+			guavaMapSingularizer.generateMethods(data, deprecate, builderType, source, fluent, returnType, returnStatement);
 			return;
 		}
 		
 		JavacTreeMaker maker = builderType.getTreeMaker();
-		Symtab symbolTable = builderType.getSymbolTable();
-		
-		JCExpression returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		JCStatement returnStatement = chain ? maker.Return(maker.Ident(builderType.toName("this"))) : null;
 		generateSingularMethod(deprecate, maker, returnType, returnStatement, data, builderType, source, fluent);
-		
-		returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		returnStatement = chain ? maker.Return(maker.Ident(builderType.toName("this"))) : null;
 		generatePluralMethod(deprecate, maker, returnType, returnStatement, data, builderType, source, fluent);
-		
-		returnType = chain ? cloneSelfType(builderType) : maker.Type(createVoidType(symbolTable, CTC_VOID));
-		returnStatement = chain ? maker.Return(maker.Ident(builderType.toName("this"))) : null;
 		generateClearMethod(deprecate, maker, returnType, returnStatement, data, builderType, source);
 	}