1 files changed, 53 insertions, 13 deletions
diff --git a/src/lombok/javac/apt/Processor.java b/src/lombok/javac/apt/Processor.java
index 864ef52c..2780ce16 100644
--- a/src/lombok/javac/apt/Processor.java
+++ b/src/lombok/javac/apt/Processor.java
@@ -1,5 +1,8 @@
 package lombok.javac.apt;
 
+import java.util.ArrayList;
+import java.util.IdentityHashMap;
+import java.util.List;
 import java.util.Set;
 
 import javax.annotation.processing.AbstractProcessor;
@@ -12,10 +15,17 @@ import javax.lang.model.element.Element;
 import javax.lang.model.element.TypeElement;
 
 import lombok.javac.HandlerLibrary;
-import lombok.javac.JavacNode;
+import lombok.javac.JavacAST;
+import lombok.javac.JavacASTAdapter;
+import lombok.javac.JavacAST.Node;
 
 import com.sun.source.util.Trees;
 import com.sun.tools.javac.processing.JavacProcessingEnvironment;
+import com.sun.tools.javac.tree.JCTree.JCAnnotation;
+import com.sun.tools.javac.tree.JCTree.JCClassDecl;
+import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
+import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
+import com.sun.tools.javac.tree.JCTree.JCVariableDecl;
 
 
 @SupportedAnnotationTypes("*")
@@ -39,23 +49,53 @@ public class Processor extends AbstractProcessor {
 		if ( processingEnv == null ) return false;
 		
 		
-		for ( TypeElement annotationType : annotations ) {
-			if ( !handlers.hasHandlerFor(annotationType) ) continue;
-			for ( Element element : roundEnv.getElementsAnnotatedWith(annotationType) ) {
-				handlers.handleAnnotation(createNode(element), annotationType);
-			}
+		IdentityHashMap<JCCompilationUnit, Void> units = new IdentityHashMap<JCCompilationUnit, Void>();
+		for ( Element element : roundEnv.getRootElements() ) units.put(toUnit(element), null);
+		
+		List<JavacAST> asts = new ArrayList<JavacAST>();
+		
+		for ( JCCompilationUnit unit : units.keySet() ) asts.add(new JavacAST(trees, processingEnv, unit));
+		
+		for ( JavacAST ast : asts ) {
+			ast.traverse(new AnnotationVisitor());
+			handlers.handleAST(ast);
+		}
+		return false;
+	}
+	
+	private class AnnotationVisitor extends JavacASTAdapter {
+		@Override public void visitAnnotationOnType(JCClassDecl type, Node annotationNode, JCAnnotation annotation) {
+			if ( annotationNode.isHandled() ) return;
+			handlers.handleAnnotation((JCCompilationUnit) annotationNode.top().get(), annotationNode, annotation);
+			annotationNode.setHandled();
 		}
 		
-		for ( Element element : roundEnv.getRootElements() ) {
-			if ( element instanceof TypeElement ) {
-				handlers.handleType((TypeElement)element);
-			}
+		@Override public void visitAnnotationOnField(JCVariableDecl field, Node annotationNode, JCAnnotation annotation) {
+			if ( annotationNode.isHandled() ) return;
+			handlers.handleAnnotation((JCCompilationUnit) annotationNode.top().get(), annotationNode, annotation);
+			annotationNode.setHandled();
 		}
 		
-		return false;
+		@Override public void visitAnnotationOnMethod(JCMethodDecl method, Node annotationNode, JCAnnotation annotation) {
+			if ( annotationNode.isHandled() ) return;
+			handlers.handleAnnotation((JCCompilationUnit) annotationNode.top().get(), annotationNode, annotation);
+			annotationNode.setHandled();
+		}
+		
+		@Override public void visitAnnotationOnMethodArgument(JCVariableDecl argument, JCMethodDecl method, Node annotationNode, JCAnnotation annotation) {
+			if ( annotationNode.isHandled() ) return;
+			handlers.handleAnnotation((JCCompilationUnit) annotationNode.top().get(), annotationNode, annotation);
+			annotationNode.setHandled();
+		}
+		
+		@Override public void visitAnnotationOnLocal(JCVariableDecl local, Node annotationNode, JCAnnotation annotation) {
+			if ( annotationNode.isHandled() ) return;
+			handlers.handleAnnotation((JCCompilationUnit) annotationNode.top().get(), annotationNode, annotation);
+			annotationNode.setHandled();
+		}
 	}
 	
-	private JavacNode createNode(Element element) {
-		return new JavacNode(trees, processingEnv, element);
+	private JCCompilationUnit toUnit(Element element) {
+		return (JCCompilationUnit) trees.getPath(element).getCompilationUnit();
 	}
 }