Javac implementation of @ExtensionMethod.

Casual tests show that it is working. Taken from lombok-pg.

4 files changed, 326 insertions, 17 deletions

diff --git a/src/core/lombok/javac/handlers/HandleData.java b/src/core/lombok/javac/handlers/HandleData.java
index bec98960..62183a15 100644
--- a/src/core/lombok/javac/handlers/HandleData.java
+++ b/src/core/lombok/javac/handlers/HandleData.java
@@ -21,7 +21,7 @@
  */
 package lombok.javac.handlers;
 
-import static lombok.javac.handlers.JavacHandlerUtil.deleteAnnotationIfNeccessary;
+import static lombok.javac.handlers.JavacHandlerUtil.*;
 import lombok.AccessLevel;
 import lombok.Data;
 import lombok.core.AnnotationValues;
@@ -30,9 +30,7 @@ import lombok.javac.JavacNode;
 
 import org.mangosdk.spi.ProviderFor;
 
-import com.sun.tools.javac.code.Flags;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
-import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 
 /**
  * Handles the {@code lombok.Data} annotation for javac.
@@ -42,12 +40,9 @@ public class HandleData extends JavacAnnotationHandler<Data> {
 	@Override public void handle(AnnotationValues<Data> annotation, JCAnnotation ast, JavacNode annotationNode) {
 		deleteAnnotationIfNeccessary(annotationNode, Data.class);
 		JavacNode typeNode = annotationNode.up();
-		JCClassDecl typeDecl = null;
-		if (typeNode.get() instanceof JCClassDecl) typeDecl = (JCClassDecl)typeNode.get();
-		long flags = typeDecl == null ? 0 : typeDecl.mods.flags;
-		boolean notAClass = (flags & (Flags.INTERFACE | Flags.ENUM | Flags.ANNOTATION)) != 0;
+		boolean notAClass = !isClass(typeNode);
 		
-		if (typeDecl == null || notAClass) {
+		if (notAClass) {
 			annotationNode.addError("@Data is only supported on a class.");
 			return;
 		}
diff --git a/src/core/lombok/javac/handlers/HandleExtensionMethod.java b/src/core/lombok/javac/handlers/HandleExtensionMethod.java
new file mode 100644
index 00000000..2df6be45
--- /dev/null
+++ b/src/core/lombok/javac/handlers/HandleExtensionMethod.java
@@ -0,0 +1,208 @@
+/*
+ * Copyright (C) 2012 The Project Lombok Authors.
+ * 
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ * 
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ * 
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+package lombok.javac.handlers;
+
+import static com.sun.tools.javac.code.Flags.*;
+import static lombok.javac.handlers.JavacHandlerUtil.*;
+import static lombok.javac.handlers.JavacResolver.*;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import javax.lang.model.element.ElementKind;
+
+import lombok.core.AnnotationValues;
+import lombok.experimental.ExtensionMethod;
+import lombok.javac.JavacAnnotationHandler;
+import lombok.javac.JavacNode;
+import lombok.javac.ResolutionBased;
+
+import org.mangosdk.spi.ProviderFor;
+
+import com.sun.source.tree.MethodInvocationTree;
+import com.sun.source.util.TreeScanner;
+import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.code.Symbol.MethodSymbol;
+import com.sun.tools.javac.code.Symbol.TypeSymbol;
+import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.code.Type.ClassType;
+import com.sun.tools.javac.code.Type.ErrorType;
+import com.sun.tools.javac.code.Type.ForAll;
+import com.sun.tools.javac.code.Type.MethodType;
+import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.tree.JCTree.JCAnnotation;
+import com.sun.tools.javac.tree.JCTree.JCClassDecl;
+import com.sun.tools.javac.tree.JCTree.JCExpression;
+import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
+import com.sun.tools.javac.tree.JCTree.JCIdent;
+import com.sun.tools.javac.tree.JCTree.JCMethodInvocation;
+import com.sun.tools.javac.tree.TreeMaker;
+
+/**
+ * Handles the {@link ExtensionMethod} annotation for javac.
+ */
+@ProviderFor(JavacAnnotationHandler.class)
+@ResolutionBased
+public class HandleExtensionMethod extends JavacAnnotationHandler<ExtensionMethod> {
+	@Override
+	public void handle(final AnnotationValues<ExtensionMethod> annotation, final JCAnnotation source, final JavacNode annotationNode) {
+		deleteAnnotationIfNeccessary(annotationNode, ExtensionMethod.class);
+		JavacNode typeNode = annotationNode.up();
+		boolean isClassOrEnum = isClassOrEnum(typeNode);
+		
+		if (!isClassOrEnum) {
+			annotationNode.addError("@ExtensionMethod can only be used on a class or an enum");
+			return;
+		}
+		
+		boolean suppressBaseMethods = annotation.getInstance().suppressBaseMethods();
+		
+		List<Object> extensionProviders = annotation.getActualExpressions("value");
+		if (extensionProviders.isEmpty()) {
+			annotationNode.addError(String.format("@%s has no effect since no extension types were specified.", ExtensionMethod.class.getName()));
+			return;
+		}
+		final List<Extension> extensions = getExtensions(annotationNode, extensionProviders);
+		if (extensions.isEmpty()) return;
+		
+		// call HandleVal explicitly to ensure val gets handled before @ExtensionMethdod gets handled.
+		// TODO maybe we should prioritize lombok handler
+		annotationNode.traverse(new HandleVal());
+		new ExtensionMethodReplaceVisitor(annotationNode, extensions, suppressBaseMethods).replace();
+		
+		annotationNode.rebuild();
+	}
+
+
+	private List<Extension> getExtensions(final JavacNode typeNode, final List<Object> extensionProviders) {
+		List<Extension> extensions = new ArrayList<Extension>();
+		for (Object extensionProvider : extensionProviders) {
+			if (!(extensionProvider instanceof JCFieldAccess)) continue;
+			JCFieldAccess provider = (JCFieldAccess) extensionProvider;
+			if (!("class".equals(provider.name.toString()))) continue;
+			Type providerType = CLASS.resolveMember(typeNode, provider.selected);
+			if (providerType == null) continue;
+			if ((providerType.tsym.flags() & (INTERFACE | ANNOTATION)) != 0) continue;
+			
+			extensions.add(getExtension(typeNode, (ClassType) providerType));	
+		}
+		return extensions;
+	}
+	
+	private Extension getExtension(final JavacNode typeNode, final ClassType extensionMethodProviderType) {
+		List<MethodSymbol> extensionMethods = new ArrayList<MethodSymbol>();
+		TypeSymbol tsym = extensionMethodProviderType.asElement();
+		if (tsym != null) for (Symbol member : tsym.getEnclosedElements()) {
+			if (member.getKind() != ElementKind.METHOD) continue;
+			MethodSymbol method = (MethodSymbol) member;
+			if ((method.flags() & (STATIC | PUBLIC)) == 0) continue;
+			if (method.params().isEmpty()) continue;
+			extensionMethods.add(method);
+		}
+		return new Extension(extensionMethods, tsym);
+	}
+	
+	private static class Extension {
+		final List<MethodSymbol> extensionMethods;
+		final TypeSymbol extensionProvider;
+		
+		public Extension(List<MethodSymbol> extensionMethods, TypeSymbol extensionProvider) {
+			this.extensionMethods = extensionMethods;
+			this.extensionProvider = extensionProvider;
+		}
+	}
+	
+	private static class ExtensionMethodReplaceVisitor extends TreeScanner<Void, Void> {
+		final JavacNode annotationNode;
+		final List<Extension> extensions;
+		final boolean suppressBaseMethods;
+		
+		public ExtensionMethodReplaceVisitor(JavacNode annotationNode, List<Extension> extensions, boolean suppressBaseMethods) {
+			this.annotationNode = annotationNode;
+			this.extensions = extensions;
+			this.suppressBaseMethods = suppressBaseMethods;
+		}
+		
+		public void replace() {
+			annotationNode.up().get().accept(this, null);
+		}
+		
+		@Override
+		public Void visitMethodInvocation(final MethodInvocationTree tree, final Void p) {
+			handleMethodCall((JCMethodInvocation) tree);
+			return super.visitMethodInvocation(tree, p);
+		}
+
+		private void handleMethodCall(final JCMethodInvocation methodCall) {
+			JavacNode methodCallNode = annotationNode.getAst().get(methodCall);
+			
+			JavacNode surroundingType = upToTypeNode(methodCallNode);
+			
+			TypeSymbol surroundingTypeSymbol = ((JCClassDecl)surroundingType.get()).sym;
+			JCExpression receiver = receiverOf(methodCall);
+			String methodName = methodNameOf(methodCall);
+			
+			if ("this".equals(methodName) || "super".equals(methodName)) return;
+			Type resolvedMethodCall = CLASS_AND_METHOD.resolveMember(methodCallNode, methodCall);
+			if (resolvedMethodCall == null) return;
+			if (!suppressBaseMethods && !(resolvedMethodCall instanceof ErrorType)) return;
+			Type receiverType = CLASS_AND_METHOD.resolveMember(methodCallNode, receiver);
+			if (receiverType == null) return;
+			if (receiverType.tsym.toString().endsWith(receiver.toString())) return;
+			
+			Types types = Types.instance(annotationNode.getContext());
+			for (Extension extension : extensions) {
+				TypeSymbol extensionProvider = extension.extensionProvider;
+				if (surroundingTypeSymbol == extensionProvider) continue;
+				for (MethodSymbol extensionMethod : extension.extensionMethods) {
+					if (!methodName.equals(extensionMethod.name.toString())) continue;
+					Type extensionMethodType = extensionMethod.type;
+					if (!MethodType.class.isInstance(extensionMethodType) && !ForAll.class.isInstance(extensionMethodType)) continue;
+					Type firstArgType = types.erasure(extensionMethodType.asMethodType().argtypes.get(0));
+					if (!types.isAssignable(receiverType, firstArgType)) continue;
+					methodCall.args = methodCall.args.prepend(receiver);
+					
+					TreeMaker maker = annotationNode.getTreeMaker();
+					JCIdent extensionClassIdent = maker.Ident(annotationNode.toName(extensionProvider.toString()));
+					methodCall.meth = maker.Select(extensionClassIdent, annotationNode.toName(methodName));
+					return;
+				}
+			}
+		}
+		
+		private String methodNameOf(final JCMethodInvocation methodCall) {
+			if (methodCall.meth instanceof JCIdent) {
+				return ((JCIdent) methodCall.meth).name.toString();
+			} else {
+				return ((JCFieldAccess) methodCall.meth).name.toString();
+			}
+		}
+		
+		private JCExpression receiverOf(final JCMethodInvocation methodCall) {
+			if (methodCall.meth instanceof JCIdent) {
+				return annotationNode.getTreeMaker().Ident(annotationNode.toName("this"));
+			} else {
+				return ((JCFieldAccess) methodCall.meth).selected;
+			}
+		}
+	}
+}
diff --git a/src/core/lombok/javac/handlers/JavacHandlerUtil.java b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
index fc7666ec..6dc97fd4 100644
--- a/src/core/lombok/javac/handlers/JavacHandlerUtil.java
+++ b/src/core/lombok/javac/handlers/JavacHandlerUtil.java
@@ -387,9 +387,7 @@ public class JavacHandlerUtil {
 	 * @param node Any node that represents the Type (JCClassDecl) to look in, or any child node thereof.
 	 */
 	public static MemberExistsResult fieldExists(String fieldName, JavacNode node) {
-		while (node != null && !(node.get() instanceof JCClassDecl)) {
-			node = node.up();
-		}
+		node = upToTypeNode(node);
 		
 		if (node != null && node.get() instanceof JCClassDecl) {
 			for (JCTree def : ((JCClassDecl)node.get()).defs) {
@@ -418,9 +416,7 @@ public class JavacHandlerUtil {
 	 * @param params The number of parameters the method should have; varargs count as 0-*. Set to -1 to find any method with the appropriate name regardless of parameter count.
 	 */
 	public static MemberExistsResult methodExists(String methodName, JavacNode node, boolean caseSensitive, int params) {
-		while (node != null && !(node.get() instanceof JCClassDecl)) {
-			node = node.up();
-		}
+		node = upToTypeNode(node);
 		
 		if (node != null && node.get() instanceof JCClassDecl) {
 			for (JCTree def : ((JCClassDecl)node.get()).defs) {
@@ -461,9 +457,7 @@ public class JavacHandlerUtil {
 	 * @param node Any node that represents the Type (JCClassDecl) to look in, or any child node thereof.
 	 */
 	public static MemberExistsResult constructorExists(JavacNode node) {
-		while (node != null && !(node.get() instanceof JCClassDecl)) {
-			node = node.up();
-		}
+		node = upToTypeNode(node);
 		
 		if (node != null && node.get() instanceof JCClassDecl) {
 			for (JCTree def : ((JCClassDecl)node.get()).defs) {
@@ -871,4 +865,28 @@ public class JavacHandlerUtil {
 		}
 		return out.toList();
 	}
+	
+	static boolean isClass(JavacNode typeNode) {
+		return isClassAndDoesNotHaveFlags(typeNode, Flags.INTERFACE | Flags.ENUM | Flags.ANNOTATION);
+	}
+	
+	static boolean isClassOrEnum(JavacNode typeNode) {
+		return isClassAndDoesNotHaveFlags(typeNode, Flags.INTERFACE | Flags.ANNOTATION);
+	}
+	
+	private static boolean isClassAndDoesNotHaveFlags(JavacNode typeNode, int flags) {
+		JCClassDecl typeDecl = null;
+		if (typeNode.get() instanceof JCClassDecl) typeDecl = (JCClassDecl)typeNode.get();
+		else return false;
+		
+		long typeDeclflags = typeDecl == null ? 0 : typeDecl.mods.flags;
+		return (typeDeclflags & flags) == 0;
+	}
+	
+	public static JavacNode upToTypeNode(JavacNode node) {
+		if (node == null) throw new NullPointerException("node");
+		while ((node != null) && !(node.get() instanceof JCClassDecl)) node = node.up();
+		
+		return node;
+	}
 }
diff --git a/src/core/lombok/javac/handlers/JavacResolver.java b/src/core/lombok/javac/handlers/JavacResolver.java
new file mode 100644
index 00000000..62f98841
--- /dev/null
+++ b/src/core/lombok/javac/handlers/JavacResolver.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2012 The Project Lombok Authors.
+ * 
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ * 
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ * 
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+package lombok.javac.handlers;
+
+import lombok.javac.JavacNode;
+import lombok.javac.JavacResolution;
+
+import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.tree.JCTree.JCBlock;
+import com.sun.tools.javac.tree.JCTree.JCExpression;
+import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
+import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
+
+public enum JavacResolver {
+	CLASS {
+		@Override
+		public Type resolveMember(JavacNode node, JCExpression expr) {
+			Type type = expr.type;
+			if (type == null) {
+				try {
+					new JavacResolution(node.getContext()).resolveClassMember(node);
+					type = expr.type;
+				} catch (Exception ignore) {
+				}
+			}
+			return type;
+		}
+	},
+	METHOD {
+		public Type resolveMember(JavacNode node, JCExpression expr) {
+			Type type = expr.type;
+			if (type == null) {
+				try {
+					JCExpression resolvedExpression = ((JCExpression) new JavacResolution(node.getContext()).resolveMethodMember(node).get(expr));
+					if (resolvedExpression != null) type = resolvedExpression.type;
+				} catch (Exception ignore) {
+				}
+			}
+			return type;
+		}
+	},
+	CLASS_AND_METHOD {
+		@Override
+		public Type resolveMember(JavacNode node, JCExpression expr) {
+			Type type = METHOD.resolveMember(node, expr);
+			if (type == null) {
+				JavacNode classNode = node;
+				while (classNode != null && noneOf(classNode.get(), JCBlock.class, JCMethodDecl.class, JCVariableDecl.class)) {
+					classNode = classNode.up();
+				}
+				if (classNode != null) {
+					type = CLASS.resolveMember(classNode, expr);
+				}
+			}
+			return type;
+		}
+		
+		private boolean noneOf(Object o, Class<?>... clazzes) {
+			for (Class<?> clazz : clazzes) {
+				if (clazz.isInstance(o)) return false;
+			}
+			return true;
+		}
+	};
+	
+	
+
+	public abstract Type resolveMember(final JavacNode node, final JCExpression expr);
+}